1 //! B+-tree node pool. 2 3 #[cfg(test)] 4 use super::Comparator; 5 use super::{Forest, Node, NodeData}; 6 use crate::entity::PrimaryMap; 7 #[cfg(test)] 8 use core::fmt; 9 use core::ops::{Index, IndexMut}; 10 11 /// A pool of nodes, including a free list. 12 pub(super) struct NodePool<F: Forest> { 13 nodes: PrimaryMap<Node, NodeData<F>>, 14 freelist: Option<Node>, 15 } 16 17 impl<F: Forest> NodePool<F> { 18 /// Allocate a new empty pool of nodes. 19 pub fn new() -> Self { 20 Self { 21 nodes: PrimaryMap::new(), 22 freelist: None, 23 } 24 } 25 26 /// Free all nodes. 27 pub fn clear(&mut self) { 28 self.nodes.clear(); 29 self.freelist = None; 30 } 31 32 /// Allocate a new node containing `data`. 33 pub fn alloc_node(&mut self, data: NodeData<F>) -> Node { 34 debug_assert!(!data.is_free(), "can't allocate free node"); 35 match self.freelist { 36 Some(node) => { 37 // Remove this node from the free list. 38 match self.nodes[node] { 39 NodeData::Free { next } => self.freelist = next, 40 _ => panic!("Invalid {} on free list", node), 41 } 42 self.nodes[node] = data; 43 node 44 } 45 None => { 46 // The free list is empty. Allocate a new node. 47 self.nodes.push(data) 48 } 49 } 50 } 51 52 /// Free a node. 53 pub fn free_node(&mut self, node: Node) { 54 // Quick check for a double free. 55 debug_assert!(!self.nodes[node].is_free(), "{node} is already free"); 56 self.nodes[node] = NodeData::Free { 57 next: self.freelist, 58 }; 59 self.freelist = Some(node); 60 } 61 62 /// Free the entire tree rooted at `node`. 63 pub fn free_tree(&mut self, node: Node) { 64 if let NodeData::Inner { size, tree, .. } = self[node] { 65 // Note that we have to capture `tree` by value to avoid borrow checker trouble. 66 for i in 0..usize::from(size + 1) { 67 // Recursively free sub-trees. This recursion can never be deeper than `MAX_PATH`, 68 // and since most trees have less than a handful of nodes, it is worthwhile to 69 // avoid the heap allocation for an iterative tree traversal. 70 self.free_tree(tree[i]); 71 } 72 } 73 self.free_node(node); 74 } 75 } 76 77 #[cfg(test)] 78 impl<F: Forest> NodePool<F> { 79 /// Verify the consistency of the tree rooted at `node`. 80 pub fn verify_tree<C: Comparator<F::Key>>(&self, node: Node, comp: &C) 81 where 82 NodeData<F>: fmt::Display, 83 F::Key: fmt::Display, 84 { 85 use crate::entity::EntitySet; 86 use alloc::vec::Vec; 87 use core::borrow::Borrow; 88 use core::cmp::Ordering; 89 90 // The root node can't be an inner node with just a single sub-tree. It should have been 91 // pruned. 92 if let NodeData::Inner { size, .. } = self[node] { 93 assert!(size > 0, "Root must have more than one sub-tree"); 94 } 95 96 let mut done = match self[node] { 97 NodeData::Inner { size, .. } | NodeData::Leaf { size, .. } => { 98 EntitySet::with_capacity(size.into()) 99 } 100 _ => EntitySet::new(), 101 }; 102 103 let mut todo = Vec::new(); 104 105 // Todo-list entries are: 106 // 1. Optional LHS key which must be <= all node entries. 107 // 2. The node reference. 108 // 3. Optional RHS key which must be > all node entries. 109 todo.push((None, node, None)); 110 111 while let Some((lkey, node, rkey)) = todo.pop() { 112 assert!(done.insert(node), "Node appears more than once in tree"); 113 let mut lower = lkey; 114 115 match self[node] { 116 NodeData::Inner { size, keys, tree } => { 117 let size = size as usize; 118 let capacity = tree.len(); 119 let keys = &keys[0..size]; 120 121 // Verify occupancy. 122 // Right-most nodes can be small, but others must be at least half full. 123 assert!( 124 rkey.is_none() || (size + 1) * 2 >= capacity, 125 "Only {}/{} entries in {}:{}, upper={}", 126 size + 1, 127 capacity, 128 node, 129 self[node], 130 rkey.unwrap() 131 ); 132 133 // Queue up the sub-trees, checking for duplicates. 134 for i in 0..size + 1 { 135 // Get an upper bound for node[i]. 136 let upper = keys.get(i).cloned().or(rkey); 137 138 // Check that keys are strictly monotonic. 139 if let (Some(a), Some(b)) = (lower, upper) { 140 assert_eq!( 141 comp.cmp(a, b), 142 Ordering::Less, 143 "Key order {} < {} failed in {}: {}", 144 a, 145 b, 146 node, 147 self[node] 148 ); 149 } 150 151 // Queue up the sub-tree. 152 todo.push((lower, tree[i], upper)); 153 154 // Set a lower bound for the next tree. 155 lower = upper; 156 } 157 } 158 NodeData::Leaf { size, keys, .. } => { 159 let size = size as usize; 160 let capacity = keys.borrow().len(); 161 let keys = &keys.borrow()[0..size]; 162 163 // Verify occupancy. 164 // Right-most nodes can be small, but others must be at least half full. 165 assert!(size > 0, "Leaf {node} is empty"); 166 assert!( 167 rkey.is_none() || size * 2 >= capacity, 168 "Only {}/{} entries in {}:{}, upper={}", 169 size, 170 capacity, 171 node, 172 self[node], 173 rkey.unwrap() 174 ); 175 176 for i in 0..size + 1 { 177 let upper = keys.get(i).cloned().or(rkey); 178 179 // Check that keys are strictly monotonic. 180 if let (Some(a), Some(b)) = (lower, upper) { 181 let wanted = if i == 0 { 182 Ordering::Equal 183 } else { 184 Ordering::Less 185 }; 186 assert_eq!( 187 comp.cmp(a, b), 188 wanted, 189 "Key order for {} - {} failed in {}: {}", 190 a, 191 b, 192 node, 193 self[node] 194 ); 195 } 196 197 // Set a lower bound for the next key. 198 lower = upper; 199 } 200 } 201 NodeData::Free { .. } => panic!("Free {} reached", node), 202 } 203 } 204 } 205 } 206 207 impl<F: Forest> Index<Node> for NodePool<F> { 208 type Output = NodeData<F>; 209 210 fn index(&self, index: Node) -> &Self::Output { 211 self.nodes.index(index) 212 } 213 } 214 215 impl<F: Forest> IndexMut<Node> for NodePool<F> { 216 fn index_mut(&mut self, index: Node) -> &mut Self::Output { 217 self.nodes.index_mut(index) 218 } 219 } 220