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