1 //! Forest of maps.
2 
3 use super::{Comparator, Forest, INNER_SIZE, Node, NodeData, NodePool, Path};
4 use crate::packed_option::PackedOption;
5 #[cfg(test)]
6 use alloc::string::String;
7 #[cfg(test)]
8 use core::fmt;
9 use core::marker::PhantomData;
10 
11 /// Tag type defining forest types for a map.
12 struct MapTypes<K, V>(PhantomData<(K, V)>);
13 
14 impl<K, V> Forest for MapTypes<K, V>
15 where
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.
34 pub struct MapForest<K, V>
35 where
36     K: Copy,
37     V: Copy,
38 {
39     nodes: NodePool<MapTypes<K, V>>,
40 }
41 
42 impl<K, V> MapForest<K, V>
43 where
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)]
71 pub struct Map<K, V>
72 where
73     K: Copy,
74     V: Copy,
75 {
76     root: PackedOption<Node>,
77     unused: PhantomData<(K, V)>,
78 }
79 
80 impl<K, V> Map<K, V>
81 where
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 
206 impl<K, V> Default for Map<K, V>
207 where
208     K: Copy,
209     V: Copy,
210 {
211     fn default() -> Self {
212         Self::new()
213     }
214 }
215 
216 #[cfg(test)]
217 impl<K, V> Map<K, V>
218 where
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.
250 pub struct MapCursor<'a, K, V, C>
251 where
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 
262 impl<'a, K, V, C> MapCursor<'a, K, V, C>
263 where
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`.
380 pub struct MapIter<'a, K, V>
381 where
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 
390 impl<'a, K, V> Iterator for MapIter<'a, K, V>
391 where
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)]
409 impl<'a, K, V, C> MapCursor<'a, K, V, C>
410 where
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)]
428 mod 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 }
922