cranelift_bforest/
map.rs

1//! Forest of maps.
2
3use super::{Comparator, Forest, Node, NodeData, NodePool, Path, INNER_SIZE};
4use crate::packed_option::PackedOption;
5#[cfg(test)]
6use alloc::string::String;
7#[cfg(test)]
8use core::fmt;
9use core::marker::PhantomData;
10
11/// Tag type defining forest types for a map.
12struct MapTypes<K, V>(PhantomData<(K, V)>);
13
14impl<K, V> Forest for MapTypes<K, V>
15where
16    K: Copy,
17    V: Copy,
18{
19    type Key = K;
20    type Value = V;
21    type LeafKeys = [K; INNER_SIZE - 1];
22    type LeafValues = [V; INNER_SIZE - 1];
23
24    fn splat_key(key: Self::Key) -> Self::LeafKeys {
25        [key; INNER_SIZE - 1]
26    }
27
28    fn splat_value(value: Self::Value) -> Self::LeafValues {
29        [value; INNER_SIZE - 1]
30    }
31}
32
33/// Memory pool for a forest of `Map` instances.
34pub struct MapForest<K, V>
35where
36    K: Copy,
37    V: Copy,
38{
39    nodes: NodePool<MapTypes<K, V>>,
40}
41
42impl<K, V> MapForest<K, V>
43where
44    K: Copy,
45    V: Copy,
46{
47    /// Create a new empty forest.
48    pub fn new() -> Self {
49        Self {
50            nodes: NodePool::new(),
51        }
52    }
53
54    /// Clear all maps in the forest.
55    ///
56    /// All `Map` instances belong to this forest are invalidated and should no longer be used.
57    pub fn clear(&mut self) {
58        self.nodes.clear();
59    }
60}
61
62/// B-tree mapping from `K` to `V`.
63///
64/// This is not a general-purpose replacement for `BTreeMap`. See the [module
65/// documentation](index.html) for more information about design tradeoffs.
66///
67/// Maps can be cloned, but that operation should only be used as part of cloning the whole forest
68/// they belong to. *Cloning a map does not allocate new memory for the clone*. It creates an alias
69/// of the same memory.
70#[derive(Clone)]
71pub struct Map<K, V>
72where
73    K: Copy,
74    V: Copy,
75{
76    root: PackedOption<Node>,
77    unused: PhantomData<(K, V)>,
78}
79
80impl<K, V> Map<K, V>
81where
82    K: Copy,
83    V: Copy,
84{
85    /// Make an empty map.
86    pub fn new() -> Self {
87        Self {
88            root: None.into(),
89            unused: PhantomData,
90        }
91    }
92
93    /// Is this an empty map?
94    pub fn is_empty(&self) -> bool {
95        self.root.is_none()
96    }
97
98    /// Get the value stored for `key`.
99    pub fn get<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> Option<V> {
100        self.root
101            .expand()
102            .and_then(|root| Path::default().find(key, root, &forest.nodes, comp))
103    }
104
105    /// Look up the value stored for `key`.
106    ///
107    /// If it exists, return the stored key-value pair.
108    ///
109    /// Otherwise, return the last key-value pair with a key that is less than or equal to `key`.
110    ///
111    /// If no stored keys are less than or equal to `key`, return `None`.
112    pub fn get_or_less<C: Comparator<K>>(
113        &self,
114        key: K,
115        forest: &MapForest<K, V>,
116        comp: &C,
117    ) -> Option<(K, V)> {
118        self.root.expand().and_then(|root| {
119            let mut path = Path::default();
120            match path.find(key, root, &forest.nodes, comp) {
121                Some(v) => Some((key, v)),
122                None => path.prev(root, &forest.nodes),
123            }
124        })
125    }
126
127    /// Insert `key, value` into the map and return the old value stored for `key`, if any.
128    pub fn insert<C: Comparator<K>>(
129        &mut self,
130        key: K,
131        value: V,
132        forest: &mut MapForest<K, V>,
133        comp: &C,
134    ) -> Option<V> {
135        self.cursor(forest, comp).insert(key, value)
136    }
137
138    /// Remove `key` from the map and return the removed value for `key`, if any.
139    pub fn remove<C: Comparator<K>>(
140        &mut self,
141        key: K,
142        forest: &mut MapForest<K, V>,
143        comp: &C,
144    ) -> Option<V> {
145        let mut c = self.cursor(forest, comp);
146        if c.goto(key).is_some() {
147            c.remove()
148        } else {
149            None
150        }
151    }
152
153    /// Remove all entries.
154    pub fn clear(&mut self, forest: &mut MapForest<K, V>) {
155        if let Some(root) = self.root.take() {
156            forest.nodes.free_tree(root);
157        }
158    }
159
160    /// Retains only the elements specified by the predicate.
161    ///
162    /// Remove all key-value pairs where the predicate returns false.
163    ///
164    /// The predicate is allowed to update the values stored in the map.
165    pub fn retain<F>(&mut self, forest: &mut MapForest<K, V>, mut predicate: F)
166    where
167        F: FnMut(K, &mut V) -> bool,
168    {
169        let mut path = Path::default();
170        if let Some(root) = self.root.expand() {
171            path.first(root, &forest.nodes);
172        }
173        while let Some((node, entry)) = path.leaf_pos() {
174            let keep = {
175                let (ks, vs) = forest.nodes[node].unwrap_leaf_mut();
176                predicate(ks[entry], &mut vs[entry])
177            };
178            if keep {
179                path.next(&forest.nodes);
180            } else {
181                self.root = path.remove(&mut forest.nodes).into();
182            }
183        }
184    }
185
186    /// Create a cursor for navigating this map. The cursor is initially positioned off the end of
187    /// the map.
188    pub fn cursor<'a, C: Comparator<K>>(
189        &'a mut self,
190        forest: &'a mut MapForest<K, V>,
191        comp: &'a C,
192    ) -> MapCursor<'a, K, V, C> {
193        MapCursor::new(self, forest, comp)
194    }
195
196    /// Create an iterator traversing this map. The iterator type is `(K, V)`.
197    pub fn iter<'a>(&'a self, forest: &'a MapForest<K, V>) -> MapIter<'a, K, V> {
198        MapIter {
199            root: self.root,
200            pool: &forest.nodes,
201            path: Path::default(),
202        }
203    }
204}
205
206impl<K, V> Default for Map<K, V>
207where
208    K: Copy,
209    V: Copy,
210{
211    fn default() -> Self {
212        Self::new()
213    }
214}
215
216#[cfg(test)]
217impl<K, V> Map<K, V>
218where
219    K: Copy + fmt::Display,
220    V: Copy,
221{
222    /// Verify consistency.
223    fn verify<C: Comparator<K>>(&self, forest: &MapForest<K, V>, comp: &C)
224    where
225        NodeData<MapTypes<K, V>>: fmt::Display,
226    {
227        if let Some(root) = self.root.expand() {
228            forest.nodes.verify_tree(root, comp);
229        }
230    }
231
232    /// Get a text version of the path to `key`.
233    fn tpath<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> String {
234        use alloc::string::ToString;
235        match self.root.expand() {
236            None => "map(empty)".to_string(),
237            Some(root) => {
238                let mut path = Path::default();
239                path.find(key, root, &forest.nodes, comp);
240                path.to_string()
241            }
242        }
243    }
244}
245
246/// A position in a `Map` used to navigate and modify the ordered map.
247///
248/// A cursor always points at a key-value pair in the map, or "off the end" which is a position
249/// after the last entry in the map.
250pub struct MapCursor<'a, K, V, C>
251where
252    K: 'a + Copy,
253    V: 'a + Copy,
254    C: 'a + Comparator<K>,
255{
256    root: &'a mut PackedOption<Node>,
257    pool: &'a mut NodePool<MapTypes<K, V>>,
258    comp: &'a C,
259    path: Path<MapTypes<K, V>>,
260}
261
262impl<'a, K, V, C> MapCursor<'a, K, V, C>
263where
264    K: Copy,
265    V: Copy,
266    C: Comparator<K>,
267{
268    /// Create a cursor with a default (off-the-end) location.
269    fn new(container: &'a mut Map<K, V>, forest: &'a mut MapForest<K, V>, comp: &'a C) -> Self {
270        Self {
271            root: &mut container.root,
272            pool: &mut forest.nodes,
273            comp,
274            path: Path::default(),
275        }
276    }
277
278    /// Is this cursor pointing to an empty map?
279    pub fn is_empty(&self) -> bool {
280        self.root.is_none()
281    }
282
283    /// Move cursor to the next key-value pair and return it.
284    ///
285    /// If the cursor reaches the end, return `None` and leave the cursor at the off-the-end
286    /// position.
287    pub fn next(&mut self) -> Option<(K, V)> {
288        self.path.next(self.pool)
289    }
290
291    /// Move cursor to the previous key-value pair and return it.
292    ///
293    /// If the cursor is already pointing at the first entry, leave it there and return `None`.
294    pub fn prev(&mut self) -> Option<(K, V)> {
295        self.root
296            .expand()
297            .and_then(|root| self.path.prev(root, self.pool))
298    }
299
300    /// Get the current key, or `None` if the cursor is at the end.
301    pub fn key(&self) -> Option<K> {
302        self.path
303            .leaf_pos()
304            .and_then(|(node, entry)| self.pool[node].unwrap_leaf().0.get(entry).cloned())
305    }
306
307    /// Get the current value, or `None` if the cursor is at the end.
308    pub fn value(&self) -> Option<V> {
309        self.path
310            .leaf_pos()
311            .and_then(|(node, entry)| self.pool[node].unwrap_leaf().1.get(entry).cloned())
312    }
313
314    /// Get a mutable reference to the current value, or `None` if the cursor is at the end.
315    pub fn value_mut(&mut self) -> Option<&mut V> {
316        self.path
317            .leaf_pos()
318            .and_then(move |(node, entry)| self.pool[node].unwrap_leaf_mut().1.get_mut(entry))
319    }
320
321    /// Move this cursor to `key`.
322    ///
323    /// If `key` is in the map, place the cursor at `key` and return the corresponding value.
324    ///
325    /// If `key` is not in the set, place the cursor at the next larger element (or the end) and
326    /// return `None`.
327    pub fn goto(&mut self, elem: K) -> Option<V> {
328        self.root.expand().and_then(|root| {
329            let v = self.path.find(elem, root, self.pool, self.comp);
330            if v.is_none() {
331                self.path.normalize(self.pool);
332            }
333            v
334        })
335    }
336
337    /// Move this cursor to the first element.
338    pub fn goto_first(&mut self) -> Option<V> {
339        self.root.map(|root| self.path.first(root, self.pool).1)
340    }
341
342    /// Insert `(key, value))` into the map and leave the cursor at the inserted pair.
343    ///
344    /// If the map did not contain `key`, return `None`.
345    ///
346    /// If `key` is already present, replace the existing with `value` and return the old value.
347    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
348        match self.root.expand() {
349            None => {
350                let root = self.pool.alloc_node(NodeData::leaf(key, value));
351                *self.root = root.into();
352                self.path.set_root_node(root);
353                None
354            }
355            Some(root) => {
356                // TODO: Optimize the case where `self.path` is already at the correct insert pos.
357                let old = self.path.find(key, root, self.pool, self.comp);
358                if old.is_some() {
359                    *self.path.value_mut(self.pool) = value;
360                } else {
361                    *self.root = self.path.insert(key, value, self.pool).into();
362                }
363                old
364            }
365        }
366    }
367
368    /// Remove the current entry (if any) and return the mapped value.
369    /// This advances the cursor to the next entry after the removed one.
370    pub fn remove(&mut self) -> Option<V> {
371        let value = self.value();
372        if value.is_some() {
373            *self.root = self.path.remove(self.pool).into();
374        }
375        value
376    }
377}
378
379/// An iterator visiting the key-value pairs of a `Map`.
380pub struct MapIter<'a, K, V>
381where
382    K: 'a + Copy,
383    V: 'a + Copy,
384{
385    root: PackedOption<Node>,
386    pool: &'a NodePool<MapTypes<K, V>>,
387    path: Path<MapTypes<K, V>>,
388}
389
390impl<'a, K, V> Iterator for MapIter<'a, K, V>
391where
392    K: 'a + Copy,
393    V: 'a + Copy,
394{
395    type Item = (K, V);
396
397    fn next(&mut self) -> Option<Self::Item> {
398        // We use `self.root` to indicate if we need to go to the first element. Reset to `None`
399        // once we've returned the first element. This also works for an empty tree since the
400        // `path.next()` call returns `None` when the path is empty. This also fuses the iterator.
401        match self.root.take() {
402            Some(root) => Some(self.path.first(root, self.pool)),
403            None => self.path.next(self.pool),
404        }
405    }
406}
407
408#[cfg(test)]
409impl<'a, K, V, C> MapCursor<'a, K, V, C>
410where
411    K: Copy + fmt::Display,
412    V: Copy + fmt::Display,
413    C: Comparator<K>,
414{
415    fn verify(&self) {
416        self.path.verify(self.pool);
417        self.root.map(|root| self.pool.verify_tree(root, self.comp));
418    }
419
420    /// Get a text version of the path to the current position.
421    fn tpath(&self) -> String {
422        use alloc::string::ToString;
423        self.path.to_string()
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use alloc::vec::Vec;
431    use core::mem;
432
433    #[test]
434    fn node_size() {
435        // check that nodes are cache line sized when keys and values are 32 bits.
436        type F = MapTypes<u32, u32>;
437        assert_eq!(mem::size_of::<NodeData<F>>(), 64);
438    }
439
440    #[test]
441    fn empty() {
442        let mut f = MapForest::<u32, f32>::new();
443        f.clear();
444
445        let mut m = Map::<u32, f32>::new();
446        assert!(m.is_empty());
447        m.clear(&mut f);
448
449        assert_eq!(m.get(7, &f, &()), None);
450        assert_eq!(m.iter(&f).next(), None);
451        assert_eq!(m.get_or_less(7, &f, &()), None);
452        m.retain(&mut f, |_, _| unreachable!());
453
454        let mut c = m.cursor(&mut f, &());
455        assert!(c.is_empty());
456        assert_eq!(c.key(), None);
457        assert_eq!(c.value(), None);
458        assert_eq!(c.next(), None);
459        assert_eq!(c.prev(), None);
460        c.verify();
461        assert_eq!(c.tpath(), "<empty path>");
462        assert_eq!(c.goto_first(), None);
463        assert_eq!(c.tpath(), "<empty path>");
464    }
465
466    #[test]
467    fn inserting() {
468        let f = &mut MapForest::<u32, f32>::new();
469        let mut m = Map::<u32, f32>::new();
470
471        // The first seven values stay in a single leaf node.
472        assert_eq!(m.insert(50, 5.0, f, &()), None);
473        assert_eq!(m.insert(50, 5.5, f, &()), Some(5.0));
474        assert_eq!(m.insert(20, 2.0, f, &()), None);
475        assert_eq!(m.insert(80, 8.0, f, &()), None);
476        assert_eq!(m.insert(40, 4.0, f, &()), None);
477        assert_eq!(m.insert(60, 6.0, f, &()), None);
478        assert_eq!(m.insert(90, 9.0, f, &()), None);
479        assert_eq!(m.insert(200, 20.0, f, &()), None);
480
481        m.verify(f, &());
482
483        assert_eq!(
484            m.iter(f).collect::<Vec<_>>(),
485            [
486                (20, 2.0),
487                (40, 4.0),
488                (50, 5.5),
489                (60, 6.0),
490                (80, 8.0),
491                (90, 9.0),
492                (200, 20.0),
493            ]
494        );
495
496        assert_eq!(m.get(0, f, &()), None);
497        assert_eq!(m.get(20, f, &()), Some(2.0));
498        assert_eq!(m.get(30, f, &()), None);
499        assert_eq!(m.get(40, f, &()), Some(4.0));
500        assert_eq!(m.get(50, f, &()), Some(5.5));
501        assert_eq!(m.get(60, f, &()), Some(6.0));
502        assert_eq!(m.get(70, f, &()), None);
503        assert_eq!(m.get(80, f, &()), Some(8.0));
504        assert_eq!(m.get(100, f, &()), None);
505
506        assert_eq!(m.get_or_less(0, f, &()), None);
507        assert_eq!(m.get_or_less(20, f, &()), Some((20, 2.0)));
508        assert_eq!(m.get_or_less(30, f, &()), Some((20, 2.0)));
509        assert_eq!(m.get_or_less(40, f, &()), Some((40, 4.0)));
510        assert_eq!(m.get_or_less(200, f, &()), Some((200, 20.0)));
511        assert_eq!(m.get_or_less(201, f, &()), Some((200, 20.0)));
512
513        {
514            let mut c = m.cursor(f, &());
515            assert_eq!(c.prev(), Some((200, 20.0)));
516            assert_eq!(c.prev(), Some((90, 9.0)));
517            assert_eq!(c.prev(), Some((80, 8.0)));
518            assert_eq!(c.prev(), Some((60, 6.0)));
519            assert_eq!(c.prev(), Some((50, 5.5)));
520            assert_eq!(c.prev(), Some((40, 4.0)));
521            assert_eq!(c.prev(), Some((20, 2.0)));
522            assert_eq!(c.prev(), None);
523        }
524
525        // Test some removals where the node stays healthy.
526        assert_eq!(m.tpath(50, f, &()), "node0[2]");
527        assert_eq!(m.tpath(80, f, &()), "node0[4]");
528        assert_eq!(m.tpath(200, f, &()), "node0[6]");
529
530        assert_eq!(m.remove(80, f, &()), Some(8.0));
531        assert_eq!(m.tpath(50, f, &()), "node0[2]");
532        assert_eq!(m.tpath(80, f, &()), "node0[4]");
533        assert_eq!(m.tpath(200, f, &()), "node0[5]");
534        assert_eq!(m.remove(80, f, &()), None);
535        m.verify(f, &());
536
537        assert_eq!(m.remove(20, f, &()), Some(2.0));
538        assert_eq!(m.tpath(50, f, &()), "node0[1]");
539        assert_eq!(m.tpath(80, f, &()), "node0[3]");
540        assert_eq!(m.tpath(200, f, &()), "node0[4]");
541        assert_eq!(m.remove(20, f, &()), None);
542        m.verify(f, &());
543
544        // [ 40 50 60 90 200 ]
545
546        {
547            let mut c = m.cursor(f, &());
548            assert_eq!(c.goto_first(), Some(4.0));
549            assert_eq!(c.key(), Some(40));
550            assert_eq!(c.value(), Some(4.0));
551            assert_eq!(c.next(), Some((50, 5.5)));
552            assert_eq!(c.next(), Some((60, 6.0)));
553            assert_eq!(c.next(), Some((90, 9.0)));
554            assert_eq!(c.next(), Some((200, 20.0)));
555            c.verify();
556            assert_eq!(c.next(), None);
557            c.verify();
558        }
559
560        // Removals from the root leaf node beyond underflow.
561        assert_eq!(m.remove(200, f, &()), Some(20.0));
562        assert_eq!(m.remove(40, f, &()), Some(4.0));
563        assert_eq!(m.remove(60, f, &()), Some(6.0));
564        m.verify(f, &());
565        assert_eq!(m.remove(50, f, &()), Some(5.5));
566        m.verify(f, &());
567        assert_eq!(m.remove(90, f, &()), Some(9.0));
568        m.verify(f, &());
569        assert!(m.is_empty());
570    }
571
572    #[test]
573    fn split_level0_leaf() {
574        // Various ways of splitting a full leaf node at level 0.
575        let f = &mut MapForest::<u32, f32>::new();
576
577        fn full_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
578            let mut m = Map::new();
579            for n in 1..8 {
580                m.insert(n * 10, n as f32 * 1.1, f, &());
581            }
582            m
583        }
584
585        // Insert at front of leaf.
586        let mut m = full_leaf(f);
587        m.insert(5, 4.2, f, &());
588        m.verify(f, &());
589        assert_eq!(m.get(5, f, &()), Some(4.2));
590
591        // Retain even entries, with altered values.
592        m.retain(f, |k, v| {
593            *v = (k / 10) as f32;
594            (k % 20) == 0
595        });
596        assert_eq!(
597            m.iter(f).collect::<Vec<_>>(),
598            [(20, 2.0), (40, 4.0), (60, 6.0)]
599        );
600
601        // Insert at back of leaf.
602        let mut m = full_leaf(f);
603        m.insert(80, 4.2, f, &());
604        m.verify(f, &());
605        assert_eq!(m.get(80, f, &()), Some(4.2));
606
607        // Insert before middle (40).
608        let mut m = full_leaf(f);
609        m.insert(35, 4.2, f, &());
610        m.verify(f, &());
611        assert_eq!(m.get(35, f, &()), Some(4.2));
612
613        // Insert after middle (40).
614        let mut m = full_leaf(f);
615        m.insert(45, 4.2, f, &());
616        m.verify(f, &());
617        assert_eq!(m.get(45, f, &()), Some(4.2));
618
619        m.clear(f);
620        assert!(m.is_empty());
621    }
622
623    #[test]
624    fn split_level1_leaf() {
625        // Various ways of splitting a full leaf node at level 1.
626        let f = &mut MapForest::<u32, f32>::new();
627
628        // Return a map whose root node is a full inner node, and the leaf nodes are all full
629        // containing:
630        //
631        // 110, 120, ..., 170
632        // 210, 220, ..., 270
633        // ...
634        // 810, 820, ..., 870
635        fn full(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
636            let mut m = Map::new();
637
638            // Start by inserting elements in order.
639            // This should leave 8 leaf nodes with 4 elements in each.
640            for row in 1..9 {
641                for col in 1..5 {
642                    m.insert(row * 100 + col * 10, row as f32 + col as f32 * 0.1, f, &());
643                }
644            }
645
646            // Then top up the leaf nodes without splitting them.
647            for row in 1..9 {
648                for col in 5..8 {
649                    m.insert(row * 100 + col * 10, row as f32 + col as f32 * 0.1, f, &());
650                }
651            }
652
653            m
654        }
655
656        let mut m = full(f);
657        // Verify geometry. Get get node2 as the root and leaves node0, 1, 3, ...
658        m.verify(f, &());
659        assert_eq!(m.tpath(110, f, &()), "node2[0]--node0[0]");
660        assert_eq!(m.tpath(140, f, &()), "node2[0]--node0[3]");
661        assert_eq!(m.tpath(210, f, &()), "node2[1]--node1[0]");
662        assert_eq!(m.tpath(270, f, &()), "node2[1]--node1[6]");
663        assert_eq!(m.tpath(310, f, &()), "node2[2]--node3[0]");
664        assert_eq!(m.tpath(810, f, &()), "node2[7]--node8[0]");
665        assert_eq!(m.tpath(870, f, &()), "node2[7]--node8[6]");
666
667        {
668            let mut c = m.cursor(f, &());
669            assert_eq!(c.goto_first(), Some(1.1));
670            assert_eq!(c.key(), Some(110));
671        }
672
673        // Front of first leaf.
674        m.insert(0, 4.2, f, &());
675        m.verify(f, &());
676        assert_eq!(m.get(0, f, &()), Some(4.2));
677
678        // First leaf split 4-4 after appending to LHS.
679        f.clear();
680        m = full(f);
681        m.insert(135, 4.2, f, &());
682        m.verify(f, &());
683        assert_eq!(m.get(135, f, &()), Some(4.2));
684
685        // First leaf split 4-4 after prepending to RHS.
686        f.clear();
687        m = full(f);
688        m.insert(145, 4.2, f, &());
689        m.verify(f, &());
690        assert_eq!(m.get(145, f, &()), Some(4.2));
691
692        // First leaf split 4-4 after appending to RHS.
693        f.clear();
694        m = full(f);
695        m.insert(175, 4.2, f, &());
696        m.verify(f, &());
697        assert_eq!(m.get(175, f, &()), Some(4.2));
698
699        // Left-middle leaf split, ins LHS.
700        f.clear();
701        m = full(f);
702        m.insert(435, 4.2, f, &());
703        m.verify(f, &());
704        assert_eq!(m.get(435, f, &()), Some(4.2));
705
706        // Left-middle leaf split, ins RHS.
707        f.clear();
708        m = full(f);
709        m.insert(445, 4.2, f, &());
710        m.verify(f, &());
711        assert_eq!(m.get(445, f, &()), Some(4.2));
712
713        // Right-middle leaf split, ins LHS.
714        f.clear();
715        m = full(f);
716        m.insert(535, 4.2, f, &());
717        m.verify(f, &());
718        assert_eq!(m.get(535, f, &()), Some(4.2));
719
720        // Right-middle leaf split, ins RHS.
721        f.clear();
722        m = full(f);
723        m.insert(545, 4.2, f, &());
724        m.verify(f, &());
725        assert_eq!(m.get(545, f, &()), Some(4.2));
726
727        // Last leaf split, ins LHS.
728        f.clear();
729        m = full(f);
730        m.insert(835, 4.2, f, &());
731        m.verify(f, &());
732        assert_eq!(m.get(835, f, &()), Some(4.2));
733
734        // Last leaf split, ins RHS.
735        f.clear();
736        m = full(f);
737        m.insert(845, 4.2, f, &());
738        m.verify(f, &());
739        assert_eq!(m.get(845, f, &()), Some(4.2));
740
741        // Front of last leaf.
742        f.clear();
743        m = full(f);
744        m.insert(805, 4.2, f, &());
745        m.verify(f, &());
746        assert_eq!(m.get(805, f, &()), Some(4.2));
747
748        m.clear(f);
749        m.verify(f, &());
750    }
751
752    // Make a tree with two barely healthy leaf nodes:
753    // [ 10 20 30 40 ] [ 50 60 70 80 ]
754    fn two_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
755        f.clear();
756        let mut m = Map::new();
757        for n in 1..9 {
758            m.insert(n * 10, n as f32, f, &());
759        }
760        m
761    }
762
763    #[test]
764    fn remove_level1() {
765        let f = &mut MapForest::<u32, f32>::new();
766        let mut m = two_leaf(f);
767
768        // Verify geometry.
769        m.verify(f, &());
770        assert_eq!(m.tpath(10, f, &()), "node2[0]--node0[0]");
771        assert_eq!(m.tpath(40, f, &()), "node2[0]--node0[3]");
772        assert_eq!(m.tpath(49, f, &()), "node2[0]--node0[4]");
773        assert_eq!(m.tpath(50, f, &()), "node2[1]--node1[0]");
774        assert_eq!(m.tpath(80, f, &()), "node2[1]--node1[3]");
775
776        // Remove the front entry from a node that stays healthy.
777        assert_eq!(m.insert(55, 5.5, f, &()), None);
778        assert_eq!(m.remove(50, f, &()), Some(5.0));
779        m.verify(f, &());
780        assert_eq!(m.tpath(49, f, &()), "node2[0]--node0[4]");
781        assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[4]");
782        assert_eq!(m.tpath(55, f, &()), "node2[1]--node1[0]");
783
784        // Remove the front entry from the first leaf node: No critical key to update.
785        assert_eq!(m.insert(15, 1.5, f, &()), None);
786        assert_eq!(m.remove(10, f, &()), Some(1.0));
787        m.verify(f, &());
788
789        // [ 15 20 30 40 ] [ 55 60 70 80 ]
790
791        // Remove the front entry from a right-most node that underflows.
792        // No rebalancing for the right-most node. Still need critical key update.
793        assert_eq!(m.remove(55, f, &()), Some(5.5));
794        m.verify(f, &());
795        assert_eq!(m.tpath(55, f, &()), "node2[0]--node0[4]");
796        assert_eq!(m.tpath(60, f, &()), "node2[1]--node1[0]");
797
798        // [ 15 20 30 40 ] [ 60 70 80 ]
799
800        // Replenish the right leaf.
801        assert_eq!(m.insert(90, 9.0, f, &()), None);
802        assert_eq!(m.insert(100, 10.0, f, &()), None);
803        m.verify(f, &());
804        assert_eq!(m.tpath(55, f, &()), "node2[0]--node0[4]");
805        assert_eq!(m.tpath(60, f, &()), "node2[1]--node1[0]");
806
807        // [ 15 20 30 40 ] [ 60 70 80 90 100 ]
808
809        // Removing one entry from the left leaf should trigger a rebalancing from the right
810        // sibling.
811        assert_eq!(m.remove(20, f, &()), Some(2.0));
812        m.verify(f, &());
813
814        // [ 15 30 40 60 ] [ 70 80 90 100 ]
815        // Check that the critical key was updated correctly.
816        assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[3]");
817        assert_eq!(m.tpath(60, f, &()), "node2[0]--node0[3]");
818        assert_eq!(m.tpath(70, f, &()), "node2[1]--node1[0]");
819
820        // Remove front entry from the left-most leaf node, underflowing.
821        // This should cause two leaf nodes to be merged and the root node to go away.
822        assert_eq!(m.remove(15, f, &()), Some(1.5));
823        m.verify(f, &());
824    }
825
826    #[test]
827    fn remove_level1_rightmost() {
828        let f = &mut MapForest::<u32, f32>::new();
829        let mut m = two_leaf(f);
830
831        // [ 10 20 30 40 ] [ 50 60 70 80 ]
832
833        // Remove entries from the right leaf. This doesn't trigger a rebalancing.
834        assert_eq!(m.remove(60, f, &()), Some(6.0));
835        assert_eq!(m.remove(80, f, &()), Some(8.0));
836        assert_eq!(m.remove(50, f, &()), Some(5.0));
837        m.verify(f, &());
838
839        // [ 10 20 30 40 ] [ 70 ]
840        assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[4]");
841        assert_eq!(m.tpath(70, f, &()), "node2[1]--node1[0]");
842
843        // Removing the last entry from the right leaf should cause a collapse.
844        assert_eq!(m.remove(70, f, &()), Some(7.0));
845        m.verify(f, &());
846    }
847
848    // Make a 3-level tree with barely healthy nodes.
849    // 1 root, 8 inner nodes, 7*4+5=33 leaf nodes, 4 entries each.
850    fn level3_sparse(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
851        f.clear();
852        let mut m = Map::new();
853        for n in 1..133 {
854            m.insert(n * 10, n as f32, f, &());
855        }
856        m
857    }
858
859    #[test]
860    fn level3_removes() {
861        let f = &mut MapForest::<u32, f32>::new();
862        let mut m = level3_sparse(f);
863        m.verify(f, &());
864
865        // Check geometry.
866        // Root: node11
867        // [ node2 170 node10 330 node16 490 node21 650 node26 810 node31 970 node36 1130 node41 ]
868        // L1: node11
869        assert_eq!(m.tpath(0, f, &()), "node11[0]--node2[0]--node0[0]");
870        assert_eq!(m.tpath(10000, f, &()), "node11[7]--node41[4]--node40[4]");
871
872        // 650 is a critical key in the middle of the root.
873        assert_eq!(m.tpath(640, f, &()), "node11[3]--node21[3]--node19[3]");
874        assert_eq!(m.tpath(650, f, &()), "node11[4]--node26[0]--node20[0]");
875
876        // Deleting 640 triggers a rebalance from node19 to node 20, cascading to n21 -> n26.
877        assert_eq!(m.remove(640, f, &()), Some(64.0));
878        m.verify(f, &());
879        assert_eq!(m.tpath(650, f, &()), "node11[3]--node26[3]--node20[3]");
880
881        // 1130 is in the first leaf of the last L1 node. Deleting it triggers a rebalance node35
882        // -> node37, but no rebalance above where there is no right sibling.
883        assert_eq!(m.tpath(1130, f, &()), "node11[6]--node41[0]--node35[0]");
884        assert_eq!(m.tpath(1140, f, &()), "node11[6]--node41[0]--node35[1]");
885        assert_eq!(m.remove(1130, f, &()), Some(113.0));
886        m.verify(f, &());
887        assert_eq!(m.tpath(1140, f, &()), "node11[6]--node41[0]--node37[0]");
888    }
889
890    #[test]
891    fn insert_many() {
892        let f = &mut MapForest::<u32, f32>::new();
893        let mut m = Map::<u32, f32>::new();
894
895        let mm = 4096;
896        let mut x = 0;
897
898        for n in 0..mm {
899            assert_eq!(m.insert(x, n as f32, f, &()), None);
900            m.verify(f, &());
901
902            x = (x + n + 1) % mm;
903        }
904
905        x = 0;
906        for n in 0..mm {
907            assert_eq!(m.get(x, f, &()), Some(n as f32));
908            x = (x + n + 1) % mm;
909        }
910
911        x = 0;
912        for n in 0..mm {
913            assert_eq!(m.remove(x, f, &()), Some(n as f32));
914            m.verify(f, &());
915
916            x = (x + n + 1) % mm;
917        }
918
919        assert!(m.is_empty());
920    }
921}