Skip to main content

cranelift_bforest/
path.rs

1//! A path from the root of a B+-tree to a leaf node.
2
3use super::node::Removed;
4use super::{Comparator, Forest, MAX_PATH, Node, NodeData, NodePool, slice_insert, slice_shift};
5use core::borrow::Borrow;
6use core::marker::PhantomData;
7use wasmtime_core::error::OutOfMemory;
8
9#[cfg(test)]
10use core::fmt;
11
12pub(super) struct Path<F: Forest> {
13    /// Number of path entries including the root and leaf nodes.
14    size: usize,
15
16    /// Path of node references from the root to a leaf node.
17    node: [Node; MAX_PATH],
18
19    /// Entry number in each node.
20    entry: [u8; MAX_PATH],
21
22    unused: PhantomData<F>,
23}
24
25impl<F: Forest> Default for Path<F> {
26    fn default() -> Self {
27        Self {
28            size: 0,
29            node: [Node(0); MAX_PATH],
30            entry: [0; MAX_PATH],
31            unused: PhantomData,
32        }
33    }
34}
35
36impl<F: Forest> Path<F> {
37    /// Reset path by searching for `key` starting from `root`.
38    ///
39    /// If `key` is in the tree, returns the corresponding value and leaved the path pointing at
40    /// the entry. Otherwise returns `None` and:
41    ///
42    /// - A key smaller than all stored keys returns a path to the first entry of the first leaf.
43    /// - A key larger than all stored keys returns a path to one beyond the last element of the
44    ///   last leaf.
45    /// - A key between the stored keys of adjacent leaf nodes returns a path to one beyond the
46    ///   last entry of the first of the leaf nodes.
47    ///
48    pub fn find(
49        &mut self,
50        key: F::Key,
51        root: Node,
52        pool: &NodePool<F>,
53        comp: &dyn Comparator<F::Key>,
54    ) -> Option<F::Value> {
55        let mut node = root;
56        for level in 0.. {
57            self.size = level + 1;
58            self.node[level] = node;
59            match pool[node] {
60                NodeData::Inner { size, keys, tree } => {
61                    // Invariant: `tree[i]` contains keys smaller than
62                    // `keys[i]`, greater or equal to `keys[i-1]`.
63                    let i = match comp.search(key, &keys[0..size.into()]) {
64                        // We hit an existing key, so follow the >= branch.
65                        Ok(i) => i + 1,
66                        // Key is less than `keys[i]`, so follow the < branch.
67                        Err(i) => i,
68                    };
69                    self.entry[level] = i as u8;
70                    node = tree[i];
71                }
72                NodeData::Leaf { size, keys, vals } => {
73                    // For a leaf we want either the found key or an insert position.
74                    return match comp.search(key, &keys.borrow()[0..size.into()]) {
75                        Ok(i) => {
76                            self.entry[level] = i as u8;
77                            Some(vals.borrow()[i])
78                        }
79                        Err(i) => {
80                            self.entry[level] = i as u8;
81                            None
82                        }
83                    };
84                }
85                NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
86            }
87        }
88        unreachable!();
89    }
90
91    /// Move path to the first entry of the tree starting at `root` and return it.
92    pub fn first(&mut self, root: Node, pool: &NodePool<F>) -> (F::Key, F::Value) {
93        let mut node = root;
94        for level in 0.. {
95            self.size = level + 1;
96            self.node[level] = node;
97            self.entry[level] = 0;
98            match pool[node] {
99                NodeData::Inner { tree, .. } => node = tree[0],
100                NodeData::Leaf { keys, vals, .. } => return (keys.borrow()[0], vals.borrow()[0]),
101                NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
102            }
103        }
104        unreachable!();
105    }
106
107    /// Move this path to the next key-value pair and return it.
108    pub fn next(&mut self, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
109        match self.leaf_pos() {
110            None => return None,
111            Some((node, entry)) => {
112                let (keys, vals) = pool[node].unwrap_leaf();
113                if entry + 1 < keys.len() {
114                    self.entry[self.size - 1] += 1;
115                    return Some((keys[entry + 1], vals[entry + 1]));
116                }
117            }
118        }
119
120        // The current leaf node is exhausted. Move to the next one.
121        let leaf_level = self.size - 1;
122        self.next_node(leaf_level, pool).map(|node| {
123            let (keys, vals) = pool[node].unwrap_leaf();
124            (keys[0], vals[0])
125        })
126    }
127
128    /// Move this path to the previous key-value pair and return it.
129    ///
130    /// If the path is at the off-the-end position, go to the last key-value pair.
131    ///
132    /// If the path is already at the first key-value pair, leave it there and return `None`.
133    pub fn prev(&mut self, root: Node, pool: &NodePool<F>) -> Option<(F::Key, F::Value)> {
134        // We use `size == 0` as a generic off-the-end position.
135        if self.size == 0 {
136            self.goto_subtree_last(0, root, pool);
137            let (node, entry) = self.leaf_pos().unwrap();
138            let (keys, vals) = pool[node].unwrap_leaf();
139            return Some((keys[entry], vals[entry]));
140        }
141
142        match self.leaf_pos() {
143            None => return None,
144            Some((node, entry)) => {
145                if entry > 0 {
146                    self.entry[self.size - 1] -= 1;
147                    let (keys, vals) = pool[node].unwrap_leaf();
148                    return Some((keys[entry - 1], vals[entry - 1]));
149                }
150            }
151        }
152
153        // The current leaf node is exhausted. Move to the previous one.
154        self.prev_leaf(pool).map(|node| {
155            let (keys, vals) = pool[node].unwrap_leaf();
156            let e = self.leaf_entry();
157            (keys[e], vals[e])
158        })
159    }
160
161    /// Move path to the first entry of the next node at level, if one exists.
162    ///
163    /// Returns the new node if it exists.
164    ///
165    /// Reset the path to `size = 0` and return `None` if there is no next node.
166    fn next_node(&mut self, level: usize, pool: &NodePool<F>) -> Option<Node> {
167        match self.right_sibling_branch_level(level, pool) {
168            None => {
169                self.size = 0;
170                None
171            }
172            Some(bl) => {
173                let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
174                self.entry[bl] += 1;
175                let mut node = bnodes[usize::from(self.entry[bl])];
176
177                for l in bl + 1..level {
178                    self.node[l] = node;
179                    self.entry[l] = 0;
180                    node = pool[node].unwrap_inner().1[0];
181                }
182
183                self.node[level] = node;
184                self.entry[level] = 0;
185                Some(node)
186            }
187        }
188    }
189
190    /// Move the path to the last entry of the previous leaf node, if one exists.
191    ///
192    /// Returns the new leaf node if it exists.
193    ///
194    /// Leave the path unchanged and returns `None` if we are already at the first leaf node.
195    fn prev_leaf(&mut self, pool: &NodePool<F>) -> Option<Node> {
196        self.left_sibling_branch_level(self.size - 1).map(|bl| {
197            let entry = self.entry[bl] - 1;
198            self.entry[bl] = entry;
199            let (_, bnodes) = pool[self.node[bl]].unwrap_inner();
200            self.goto_subtree_last(bl + 1, bnodes[usize::from(entry)], pool)
201        })
202    }
203
204    /// Move this path to the last position for the sub-tree at `level, root`.
205    fn goto_subtree_last(&mut self, level: usize, root: Node, pool: &NodePool<F>) -> Node {
206        let mut node = root;
207        for l in level.. {
208            self.node[l] = node;
209            match pool[node] {
210                NodeData::Inner { size, ref tree, .. } => {
211                    self.entry[l] = size;
212                    node = tree[usize::from(size)];
213                }
214                NodeData::Leaf { size, .. } => {
215                    self.entry[l] = size - 1;
216                    self.size = l + 1;
217                    break;
218                }
219                NodeData::Free { .. } => panic!("Free {} reached from {}", node, root),
220            }
221        }
222        node
223    }
224
225    /// Set the root node and point the path at the first entry of the node.
226    pub fn set_root_node(&mut self, root: Node) {
227        self.size = 1;
228        self.node[0] = root;
229        self.entry[0] = 0;
230    }
231
232    /// Get the current leaf node and entry, if any.
233    pub fn leaf_pos(&self) -> Option<(Node, usize)> {
234        let i = self.size.wrapping_sub(1);
235        self.node.get(i).map(|&n| (n, self.entry[i].into()))
236    }
237
238    /// Get the current leaf node.
239    fn leaf_node(&self) -> Node {
240        self.node[self.size - 1]
241    }
242
243    /// Get the current entry in the leaf node.
244    fn leaf_entry(&self) -> usize {
245        self.entry[self.size - 1].into()
246    }
247
248    /// Is this path pointing to the first entry in the tree?
249    /// This corresponds to the smallest key.
250    fn at_first_entry(&self) -> bool {
251        self.entry[0..self.size].iter().all(|&i| i == 0)
252    }
253
254    /// Get a mutable reference to the current value.
255    /// This assumes that there is a current value.
256    pub fn value_mut<'a>(&self, pool: &'a mut NodePool<F>) -> &'a mut F::Value {
257        &mut pool[self.leaf_node()].unwrap_leaf_mut().1[self.leaf_entry()]
258    }
259
260    /// Insert the key-value pair at the current position.
261    /// The current position must be the correct insertion location for the key.
262    /// This function does not check for duplicate keys. Use `find` or similar for that.
263    /// Returns the new root node.
264    pub fn insert(
265        &mut self,
266        key: F::Key,
267        value: F::Value,
268        pool: &mut NodePool<F>,
269    ) -> Result<Node, OutOfMemory> {
270        if !self.try_leaf_insert(key, value, pool) {
271            self.split_and_insert(key, value, pool)?;
272        }
273        Ok(self.node[0])
274    }
275
276    /// Try to insert `key, value` at the current position, but fail and return false if the leaf
277    /// node is full.
278    fn try_leaf_insert(&self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> bool {
279        let index = self.leaf_entry();
280
281        // The case `index == 0` should only ever happen when there are no earlier leaf nodes,
282        // otherwise we should have appended to the previous leaf node instead. This invariant
283        // means that we don't need to update keys stored in inner nodes here.
284        debug_assert!(index > 0 || self.at_first_entry());
285
286        pool[self.leaf_node()].try_leaf_insert(index, key, value)
287    }
288
289    /// Split the current leaf node and then insert `key, value`.
290    /// This should only be used if `try_leaf_insert()` fails.
291    fn split_and_insert(
292        &mut self,
293        mut key: F::Key,
294        value: F::Value,
295        pool: &mut NodePool<F>,
296    ) -> Result<(), OutOfMemory> {
297        let orig_root = self.node[0];
298
299        // Loop invariant: We need to split the node at `level` and then retry a failed insertion.
300        // The items to insert are either `(key, ins_node)` or `(key, value)`.
301        let mut ins_node = None;
302        let mut split;
303        for level in (0..self.size).rev() {
304            // Split the current node.
305            let mut node = self.node[level];
306            let mut entry = self.entry[level].into();
307            split = pool[node].split(entry);
308            let rhs_node = pool.alloc_node(split.rhs_data)?;
309
310            // Should the path be moved to the new RHS node?
311            // Prefer the smaller node if we're right in the middle.
312            // Prefer to append to LHS all other things being equal.
313            //
314            // When inserting into an inner node (`ins_node.is_some()`), we must point to a valid
315            // entry in the current node since the new entry is inserted *after* the insert
316            // location.
317            if entry > split.lhs_entries
318                || (entry == split.lhs_entries
319                    && (split.lhs_entries > split.rhs_entries || ins_node.is_some()))
320            {
321                node = rhs_node;
322                entry -= split.lhs_entries;
323                self.node[level] = node;
324                self.entry[level] = entry as u8;
325            }
326
327            // Now that we have a not-full node, it must be possible to insert.
328            match ins_node {
329                None => {
330                    let inserted = pool[node].try_leaf_insert(entry, key, value);
331                    debug_assert!(inserted);
332                    // If we inserted at the front of the new rhs_node leaf, we need to propagate
333                    // the inserted key as the critical key instead of the previous front key.
334                    if entry == 0 && node == rhs_node {
335                        split.crit_key = key;
336                    }
337                }
338                Some(n) => {
339                    let inserted = pool[node].try_inner_insert(entry, key, n);
340                    debug_assert!(inserted);
341                    // The lower level was moved to the new RHS node, so make sure that is
342                    // reflected here.
343                    if n == self.node[level + 1] {
344                        self.entry[level] += 1;
345                    }
346                }
347            }
348
349            // We are now done with the current level, but `rhs_node` must be inserted in the inner
350            // node above us. If we're already at level 0, the root node needs to be split.
351            key = split.crit_key;
352            ins_node = Some(rhs_node);
353            if level > 0 {
354                let pnode = &mut pool[self.node[level - 1]];
355                let pentry = self.entry[level - 1].into();
356                if pnode.try_inner_insert(pentry, key, rhs_node) {
357                    // If this level level was moved to the new RHS node, update parent entry.
358                    if node == rhs_node {
359                        self.entry[level - 1] += 1;
360                    }
361                    return Ok(());
362                }
363            }
364        }
365
366        // If we get here we have split the original root node and need to add an extra level.
367        let rhs_node = ins_node.expect("empty path");
368        let root = pool.alloc_node(NodeData::inner(orig_root, key, rhs_node))?;
369        let entry = if self.node[0] == rhs_node { 1 } else { 0 };
370        self.size += 1;
371        slice_insert(&mut self.node[0..self.size], 0, root);
372        slice_insert(&mut self.entry[0..self.size], 0, entry);
373        Ok(())
374    }
375
376    /// Remove the key-value pair at the current position and advance the path to the next
377    /// key-value pair, leaving the path in a normalized state.
378    ///
379    /// Return the new root node.
380    pub fn remove(&mut self, pool: &mut NodePool<F>) -> Option<Node> {
381        let e = self.leaf_entry();
382        match pool[self.leaf_node()].leaf_remove(e) {
383            Removed::Healthy => {
384                if e == 0 {
385                    self.update_crit_key(pool)
386                }
387                Some(self.node[0])
388            }
389            status => self.balance_nodes(status, pool),
390        }
391    }
392
393    /// Get the critical key for the current node at `level`.
394    ///
395    /// The critical key is less than or equal to all keys in the sub-tree at `level` and greater
396    /// than all keys to the left of the current node at `level`.
397    ///
398    /// The left-most node at any level does not have a critical key.
399    fn current_crit_key(&self, level: usize, pool: &NodePool<F>) -> Option<F::Key> {
400        // Find the level containing the critical key for the current node.
401        self.left_sibling_branch_level(level).map(|bl| {
402            let (keys, _) = pool[self.node[bl]].unwrap_inner();
403            keys[usize::from(self.entry[bl]) - 1]
404        })
405    }
406
407    /// Update the critical key after removing the front entry of the leaf node.
408    fn update_crit_key(&mut self, pool: &mut NodePool<F>) {
409        // Find the inner level containing the critical key for the current leaf node.
410        let crit_level = match self.left_sibling_branch_level(self.size - 1) {
411            None => return,
412            Some(l) => l,
413        };
414        let crit_kidx = self.entry[crit_level] - 1;
415
416        // Extract the new critical key from the leaf node.
417        let crit_key = pool[self.leaf_node()].leaf_crit_key();
418        let crit_node = self.node[crit_level];
419
420        match pool[crit_node] {
421            NodeData::Inner {
422                size, ref mut keys, ..
423            } => {
424                debug_assert!(crit_kidx < size);
425                keys[usize::from(crit_kidx)] = crit_key;
426            }
427            _ => panic!("Expected inner node"),
428        }
429    }
430
431    /// Given that the current leaf node is in an unhealthy (underflowed or even empty) status,
432    /// balance it with sibling nodes.
433    ///
434    /// Return the new root node.
435    fn balance_nodes(&mut self, status: Removed, pool: &mut NodePool<F>) -> Option<Node> {
436        // The current leaf node is not in a healthy state, and its critical key may have changed
437        // too.
438        //
439        // Start by dealing with a changed critical key for the leaf level.
440        if status != Removed::Empty && self.leaf_entry() == 0 {
441            self.update_crit_key(pool);
442        }
443
444        let leaf_level = self.size - 1;
445        if self.heal_level(status, leaf_level, pool) {
446            // Tree has become empty.
447            self.size = 0;
448            return None;
449        }
450
451        // Discard the root node if it has shrunk to a single sub-tree.
452        let mut ns = 0;
453        while let NodeData::Inner {
454            size: 0, ref tree, ..
455        } = pool[self.node[ns]]
456        {
457            ns += 1;
458            self.node[ns] = tree[0];
459        }
460
461        if ns > 0 {
462            for l in 0..ns {
463                pool.free_node(self.node[l]);
464            }
465
466            // Shift the whole array instead of just 0..size because `self.size` may be cleared
467            // here if the path is pointing off-the-end.
468            slice_shift(&mut self.node, ns);
469            slice_shift(&mut self.entry, ns);
470
471            if self.size > 0 {
472                self.size -= ns;
473            }
474        }
475
476        // Return the root node, even when `size=0` indicating that we're at the off-the-end
477        // position.
478        Some(self.node[0])
479    }
480
481    /// After removing an entry from the node at `level`, check its health and rebalance as needed.
482    ///
483    /// Leave the path up to and including `level` in a normalized state where all entries are in
484    /// bounds.
485    ///
486    /// Returns true if the tree becomes empty.
487    fn heal_level(&mut self, status: Removed, level: usize, pool: &mut NodePool<F>) -> bool {
488        match status {
489            Removed::Healthy => {}
490            Removed::Rightmost => {
491                // The rightmost entry was removed from the current node, so move the path so it
492                // points at the first entry of the next node at this level.
493                debug_assert_eq!(
494                    usize::from(self.entry[level]),
495                    pool[self.node[level]].entries()
496                );
497                self.next_node(level, pool);
498            }
499            Removed::Underflow => self.underflowed_node(level, pool),
500            Removed::Empty => return self.empty_node(level, pool),
501        }
502        false
503    }
504
505    /// The current node at `level` has underflowed, meaning that it is below half capacity but
506    /// not completely empty.
507    ///
508    /// Handle this by balancing entries with the right sibling node.
509    ///
510    /// Leave the path up to and including `level` in a valid state that points to the same entry.
511    fn underflowed_node(&mut self, level: usize, pool: &mut NodePool<F>) {
512        // Look for a right sibling node at this level. If none exists, we allow the underflowed
513        // node to persist as the right-most node at its level.
514        if let Some((crit_key, rhs_node)) = self.right_sibling(level, pool) {
515            // New critical key for the updated right sibling node.
516            let new_ck: Option<F::Key>;
517            let empty;
518            // Make a COPY of the sibling node to avoid fighting the borrow checker.
519            let mut rhs = pool[rhs_node];
520            match pool[self.node[level]].balance(crit_key, &mut rhs) {
521                None => {
522                    // Everything got moved to the RHS node.
523                    new_ck = self.current_crit_key(level, pool);
524                    empty = true;
525                }
526                Some(key) => {
527                    // Entries moved from RHS node.
528                    new_ck = Some(key);
529                    empty = false;
530                }
531            }
532            // Put back the updated RHS node data.
533            pool[rhs_node] = rhs;
534            // Update the critical key for the RHS node unless it has become a left-most
535            // node.
536            if let Some(ck) = new_ck {
537                self.update_right_crit_key(level, ck, pool);
538            }
539            if empty {
540                let empty_tree = self.empty_node(level, pool);
541                debug_assert!(!empty_tree);
542            }
543
544            // Any Removed::Rightmost state must have been cleared above by merging nodes. If the
545            // current entry[level] was one off the end of the node, it will now point at a proper
546            // entry.
547            debug_assert!(usize::from(self.entry[level]) < pool[self.node[level]].entries());
548        } else if usize::from(self.entry[level]) >= pool[self.node[level]].entries() {
549            // There's no right sibling at this level, so the node can't be rebalanced.
550            // Check if we are in an off-the-end position.
551            self.size = 0;
552        }
553    }
554
555    /// The current node at `level` has become empty.
556    ///
557    /// Remove the node from its parent node and leave the path in a normalized state. This means
558    /// that the path at this level will go through the right sibling of this node.
559    ///
560    /// If the current node has no right sibling, set `self.size = 0`.
561    ///
562    /// Returns true if the tree becomes empty.
563    fn empty_node(&mut self, level: usize, pool: &mut NodePool<F>) -> bool {
564        pool.free_node(self.node[level]);
565        if level == 0 {
566            // We just deleted the root node, so the tree is now empty.
567            return true;
568        }
569
570        // Get the right sibling node before recursively removing nodes.
571        let rhs_node = self.right_sibling(level, pool).map(|(_, n)| n);
572
573        // Remove the current sub-tree from the parent node.
574        let pl = level - 1;
575        let pe = self.entry[pl].into();
576        let status = pool[self.node[pl]].inner_remove(pe);
577        self.heal_level(status, pl, pool);
578
579        // Finally update the path at this level.
580        match rhs_node {
581            // We'll leave `self.entry[level]` unchanged. It can be non-zero after moving node
582            // entries to the right sibling node.
583            Some(rhs) => self.node[level] = rhs,
584            // We have no right sibling, so we must have deleted the right-most
585            // entry. The path should be moved to the "off-the-end" position.
586            None => self.size = 0,
587        }
588        false
589    }
590
591    /// Find the level where the right sibling to the current node at `level` branches off.
592    ///
593    /// This will be an inner node with two adjacent sub-trees: In one the current node at level is
594    /// a right-most node, in the other, the right sibling is a left-most node.
595    ///
596    /// Returns `None` if the current node is a right-most node so no right sibling exists.
597    fn right_sibling_branch_level(&self, level: usize, pool: &NodePool<F>) -> Option<usize> {
598        (0..level).rposition(|l| match pool[self.node[l]] {
599            NodeData::Inner { size, .. } => self.entry[l] < size,
600            _ => panic!("Expected inner node"),
601        })
602    }
603
604    /// Find the level where the left sibling to the current node at `level` branches off.
605    fn left_sibling_branch_level(&self, level: usize) -> Option<usize> {
606        self.entry[0..level].iter().rposition(|&e| e != 0)
607    }
608
609    /// Get the right sibling node to the current node at `level`.
610    /// Also return the critical key between the current node and the right sibling.
611    fn right_sibling(&self, level: usize, pool: &NodePool<F>) -> Option<(F::Key, Node)> {
612        // Find the critical level: The deepest level where two sibling subtrees contain the
613        // current node and its right sibling.
614        self.right_sibling_branch_level(level, pool).map(|bl| {
615            // Extract the critical key and the `bl+1` node.
616            let be = usize::from(self.entry[bl]);
617            let crit_key;
618            let mut node;
619            {
620                let (keys, tree) = pool[self.node[bl]].unwrap_inner();
621                crit_key = keys[be];
622                node = tree[be + 1];
623            }
624
625            // Follow left-most links back down to `level`.
626            for _ in bl + 1..level {
627                node = pool[node].unwrap_inner().1[0];
628            }
629
630            (crit_key, node)
631        })
632    }
633
634    /// Update the critical key for the right sibling node at `level`.
635    fn update_right_crit_key(&self, level: usize, crit_key: F::Key, pool: &mut NodePool<F>) {
636        let bl = self
637            .right_sibling_branch_level(level, pool)
638            .expect("No right sibling exists");
639        match pool[self.node[bl]] {
640            NodeData::Inner { ref mut keys, .. } => {
641                keys[usize::from(self.entry[bl])] = crit_key;
642            }
643            _ => panic!("Expected inner node"),
644        }
645    }
646
647    /// Normalize the path position such that it is either pointing at a real entry or `size=0`
648    /// indicating "off-the-end".
649    pub fn normalize(&mut self, pool: &NodePool<F>) {
650        if let Some((leaf, entry)) = self.leaf_pos() {
651            if entry >= pool[leaf].entries() {
652                let leaf_level = self.size - 1;
653                self.next_node(leaf_level, pool);
654            }
655        }
656    }
657}
658
659#[cfg(test)]
660impl<F: Forest> Path<F> {
661    /// Check the internal consistency of this path.
662    pub fn verify(&self, pool: &NodePool<F>) {
663        for level in 0..self.size {
664            match pool[self.node[level]] {
665                NodeData::Inner { size, tree, .. } => {
666                    assert!(level < self.size - 1, "Expected leaf node at level {level}");
667                    assert!(
668                        self.entry[level] <= size,
669                        "OOB inner entry {}/{} at level {}",
670                        self.entry[level],
671                        size,
672                        level
673                    );
674                    assert_eq!(
675                        self.node[level + 1],
676                        tree[usize::from(self.entry[level])],
677                        "Node mismatch at level {level}"
678                    );
679                }
680                NodeData::Leaf { size, .. } => {
681                    assert_eq!(level, self.size - 1, "Expected inner node");
682                    assert!(
683                        self.entry[level] <= size,
684                        "OOB leaf entry {}/{}",
685                        self.entry[level],
686                        size,
687                    );
688                }
689                NodeData::Free { .. } => {
690                    panic!("Free {} in path", self.node[level]);
691                }
692            }
693        }
694    }
695}
696
697#[cfg(test)]
698impl<F: Forest> fmt::Display for Path<F> {
699    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
700        if self.size == 0 {
701            write!(f, "<empty path>")
702        } else {
703            write!(f, "{}[{}]", self.node[0], self.entry[0])?;
704            for i in 1..self.size {
705                write!(f, "--{}[{}]", self.node[i], self.entry[i])?;
706            }
707            Ok(())
708        }
709    }
710}
711
712#[cfg(test)]
713mod tests {
714    use wasmtime_core::alloc::PanicOnOom;
715
716    use super::*;
717    use core::cmp::Ordering;
718
719    struct TC();
720
721    impl Comparator<i32> for TC {
722        fn cmp(&self, a: i32, b: i32) -> Ordering {
723            a.cmp(&b)
724        }
725    }
726
727    struct TF();
728
729    impl Forest for TF {
730        type Key = i32;
731        type Value = char;
732        type LeafKeys = [i32; 7];
733        type LeafValues = [char; 7];
734
735        fn splat_key(key: Self::Key) -> Self::LeafKeys {
736            [key; 7]
737        }
738
739        fn splat_value(value: Self::Value) -> Self::LeafValues {
740            [value; 7]
741        }
742    }
743
744    #[test]
745    fn search_single_leaf() {
746        // Testing Path::new() for trees with a single leaf node.
747        let mut pool = NodePool::<TF>::new();
748        let root = pool.alloc_node(NodeData::leaf(10, 'a')).panic_on_oom();
749        let mut p = Path::default();
750        let comp = TC();
751
752        // Search for key less than stored key.
753        assert_eq!(p.find(5, root, &pool, &comp), None);
754        assert_eq!(p.size, 1);
755        assert_eq!(p.node[0], root);
756        assert_eq!(p.entry[0], 0);
757
758        // Search for stored key.
759        assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
760        assert_eq!(p.size, 1);
761        assert_eq!(p.node[0], root);
762        assert_eq!(p.entry[0], 0);
763
764        // Search for key greater than stored key.
765        assert_eq!(p.find(15, root, &pool, &comp), None);
766        assert_eq!(p.size, 1);
767        assert_eq!(p.node[0], root);
768        assert_eq!(p.entry[0], 1);
769
770        // Modify leaf node to contain two values.
771        match pool[root] {
772            NodeData::Leaf {
773                ref mut size,
774                ref mut keys,
775                ref mut vals,
776            } => {
777                *size = 2;
778                keys[1] = 20;
779                vals[1] = 'b';
780            }
781            _ => unreachable!(),
782        }
783
784        // Search for key between stored keys.
785        assert_eq!(p.find(15, root, &pool, &comp), None);
786        assert_eq!(p.size, 1);
787        assert_eq!(p.node[0], root);
788        assert_eq!(p.entry[0], 1);
789
790        // Search for key greater than stored keys.
791        assert_eq!(p.find(25, root, &pool, &comp), None);
792        assert_eq!(p.size, 1);
793        assert_eq!(p.node[0], root);
794        assert_eq!(p.entry[0], 2);
795    }
796
797    #[test]
798    fn search_single_inner() {
799        // Testing Path::new() for trees with a single inner node and two leaves.
800        let mut pool = NodePool::<TF>::new();
801        let leaf1 = pool.alloc_node(NodeData::leaf(10, 'a')).panic_on_oom();
802        let leaf2 = pool.alloc_node(NodeData::leaf(20, 'b')).panic_on_oom();
803        let root = pool
804            .alloc_node(NodeData::inner(leaf1, 20, leaf2))
805            .panic_on_oom();
806        let mut p = Path::default();
807        let comp = TC();
808
809        // Search for key less than stored keys.
810        assert_eq!(p.find(5, root, &pool, &comp), None);
811        assert_eq!(p.size, 2);
812        assert_eq!(p.node[0], root);
813        assert_eq!(p.entry[0], 0);
814        assert_eq!(p.node[1], leaf1);
815        assert_eq!(p.entry[1], 0);
816
817        assert_eq!(p.find(10, root, &pool, &comp), Some('a'));
818        assert_eq!(p.size, 2);
819        assert_eq!(p.node[0], root);
820        assert_eq!(p.entry[0], 0);
821        assert_eq!(p.node[1], leaf1);
822        assert_eq!(p.entry[1], 0);
823
824        // Midway between the two leaf nodes.
825        assert_eq!(p.find(15, root, &pool, &comp), None);
826        assert_eq!(p.size, 2);
827        assert_eq!(p.node[0], root);
828        assert_eq!(p.entry[0], 0);
829        assert_eq!(p.node[1], leaf1);
830        assert_eq!(p.entry[1], 1);
831
832        assert_eq!(p.find(20, root, &pool, &comp), Some('b'));
833        assert_eq!(p.size, 2);
834        assert_eq!(p.node[0], root);
835        assert_eq!(p.entry[0], 1);
836        assert_eq!(p.node[1], leaf2);
837        assert_eq!(p.entry[1], 0);
838
839        assert_eq!(p.find(25, root, &pool, &comp), None);
840        assert_eq!(p.size, 2);
841        assert_eq!(p.node[0], root);
842        assert_eq!(p.entry[0], 1);
843        assert_eq!(p.node[1], leaf2);
844        assert_eq!(p.entry[1], 1);
845    }
846}