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::{
10     marker::PhantomData,
11     mem,
12     ops::{Bound, RangeBounds},
13 };
14 use wasmtime_core::{alloc::PanicOnOom as _, error::OutOfMemory};
15 
16 /// Tag type defining forest types for a map.
17 struct MapTypes<K, V>(PhantomData<(K, V)>);
18 
19 impl<K, V> Forest for MapTypes<K, V>
20 where
21     K: Copy,
22     V: Copy,
23 {
24     type Key = K;
25     type Value = V;
26     type LeafKeys = [K; INNER_SIZE - 1];
27     type LeafValues = [V; INNER_SIZE - 1];
28 
splat_key(key: Self::Key) -> Self::LeafKeys29     fn splat_key(key: Self::Key) -> Self::LeafKeys {
30         [key; INNER_SIZE - 1]
31     }
32 
splat_value(value: Self::Value) -> Self::LeafValues33     fn splat_value(value: Self::Value) -> Self::LeafValues {
34         [value; INNER_SIZE - 1]
35     }
36 }
37 
38 /// Memory pool for a forest of `Map` instances.
39 pub struct MapForest<K, V>
40 where
41     K: Copy,
42     V: Copy,
43 {
44     nodes: NodePool<MapTypes<K, V>>,
45 }
46 
47 impl<K, V> MapForest<K, V>
48 where
49     K: Copy,
50     V: Copy,
51 {
52     /// Create a new empty forest.
new() -> Self53     pub fn new() -> Self {
54         Self {
55             nodes: NodePool::new(),
56         }
57     }
58 
59     /// Clear all maps in the forest.
60     ///
61     /// All `Map` instances belong to this forest are invalidated and should no longer be used.
clear(&mut self)62     pub fn clear(&mut self) {
63         self.nodes.clear();
64     }
65 }
66 
67 /// B-tree mapping from `K` to `V`.
68 ///
69 /// This is not a general-purpose replacement for `BTreeMap`. See the [module
70 /// documentation](index.html) for more information about design tradeoffs.
71 ///
72 /// Maps can be cloned, but that operation should only be used as part of cloning the whole forest
73 /// they belong to. *Cloning a map does not allocate new memory for the clone*. It creates an alias
74 /// of the same memory.
75 #[derive(Clone)]
76 pub struct Map<K, V>
77 where
78     K: Copy,
79     V: Copy,
80 {
81     root: PackedOption<Node>,
82     unused: PhantomData<(K, V)>,
83 }
84 
85 impl<K, V> Map<K, V>
86 where
87     K: Copy,
88     V: Copy,
89 {
90     /// Make an empty map.
new() -> Self91     pub fn new() -> Self {
92         Self {
93             root: None.into(),
94             unused: PhantomData,
95         }
96     }
97 
98     /// Is this an empty map?
is_empty(&self) -> bool99     pub fn is_empty(&self) -> bool {
100         self.root.is_none()
101     }
102 
103     /// Get the value stored for `key`.
get<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> Option<V>104     pub fn get<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> Option<V> {
105         self.root
106             .expand()
107             .and_then(|root| Path::default().find(key, root, &forest.nodes, comp))
108     }
109 
110     /// Look up the value stored for `key`.
111     ///
112     /// If it exists, return the stored key-value pair.
113     ///
114     /// Otherwise, return the last key-value pair with a key that is less than or equal to `key`.
115     ///
116     /// If no stored keys are less than or equal to `key`, return `None`.
get_or_less<C: Comparator<K>>( &self, key: K, forest: &MapForest<K, V>, comp: &C, ) -> Option<(K, V)>117     pub fn get_or_less<C: Comparator<K>>(
118         &self,
119         key: K,
120         forest: &MapForest<K, V>,
121         comp: &C,
122     ) -> Option<(K, V)> {
123         self.root.expand().and_then(|root| {
124             let mut path = Path::default();
125             match path.find(key, root, &forest.nodes, comp) {
126                 Some(v) => Some((key, v)),
127                 None => path.prev(root, &forest.nodes),
128             }
129         })
130     }
131 
132     /// Insert `key, value` into the map and return the old value stored for `key`, if any.
insert<C: Comparator<K>>( &mut self, key: K, value: V, forest: &mut MapForest<K, V>, comp: &C, ) -> Option<V>133     pub fn insert<C: Comparator<K>>(
134         &mut self,
135         key: K,
136         value: V,
137         forest: &mut MapForest<K, V>,
138         comp: &C,
139     ) -> Option<V> {
140         self.cursor_mut(forest, comp).insert(key, value)
141     }
142 
143     /// Like `insert` but returns an error on allocation failure.
try_insert<C: Comparator<K>>( &mut self, key: K, value: V, forest: &mut MapForest<K, V>, comp: &C, ) -> Result<Option<V>, OutOfMemory>144     pub fn try_insert<C: Comparator<K>>(
145         &mut self,
146         key: K,
147         value: V,
148         forest: &mut MapForest<K, V>,
149         comp: &C,
150     ) -> Result<Option<V>, OutOfMemory> {
151         self.cursor_mut(forest, comp).try_insert(key, value)
152     }
153 
154     /// Remove `key` from the map and return the removed value for `key`, if any.
remove<C: Comparator<K>>( &mut self, key: K, forest: &mut MapForest<K, V>, comp: &C, ) -> Option<V>155     pub fn remove<C: Comparator<K>>(
156         &mut self,
157         key: K,
158         forest: &mut MapForest<K, V>,
159         comp: &C,
160     ) -> Option<V> {
161         let mut c = self.cursor_mut(forest, comp);
162         if c.goto(key).is_some() {
163             c.remove()
164         } else {
165             None
166         }
167     }
168 
169     /// Remove all entries.
clear(&mut self, forest: &mut MapForest<K, V>)170     pub fn clear(&mut self, forest: &mut MapForest<K, V>) {
171         if let Some(root) = self.root.take() {
172             forest.nodes.free_tree(root);
173         }
174     }
175 
176     /// Retains only the elements specified by the predicate.
177     ///
178     /// Remove all key-value pairs where the predicate returns false.
179     ///
180     /// The predicate is allowed to update the values stored in the map.
retain<F>(&mut self, forest: &mut MapForest<K, V>, mut predicate: F) where F: FnMut(K, &mut V) -> bool,181     pub fn retain<F>(&mut self, forest: &mut MapForest<K, V>, mut predicate: F)
182     where
183         F: FnMut(K, &mut V) -> bool,
184     {
185         let mut path = Path::default();
186         if let Some(root) = self.root.expand() {
187             path.first(root, &forest.nodes);
188         }
189         while let Some((node, entry)) = path.leaf_pos() {
190             let keep = {
191                 let (ks, vs) = forest.nodes[node].unwrap_leaf_mut();
192                 predicate(ks[entry], &mut vs[entry])
193             };
194             if keep {
195                 path.next(&forest.nodes);
196             } else {
197                 self.root = path.remove(&mut forest.nodes).into();
198             }
199         }
200     }
201 
202     /// Create an immutable cursor for navigating this map.
203     ///
204     /// The cursor is initially positioned off the end of the map.
cursor<'a, C: Comparator<K>>( &'a self, forest: &'a MapForest<K, V>, comp: &'a C, ) -> MapCursor<'a, K, V, C>205     pub fn cursor<'a, C: Comparator<K>>(
206         &'a self,
207         forest: &'a MapForest<K, V>,
208         comp: &'a C,
209     ) -> MapCursor<'a, K, V, C> {
210         MapCursor::new(self, forest, comp)
211     }
212 
213     /// Create a mutable cursor for navigating this map.
214     ///
215     /// The cursor is initially positioned off the end of the map.
cursor_mut<'a, C: Comparator<K>>( &'a mut self, forest: &'a mut MapForest<K, V>, comp: &'a C, ) -> MapCursorMut<'a, K, V, C>216     pub fn cursor_mut<'a, C: Comparator<K>>(
217         &'a mut self,
218         forest: &'a mut MapForest<K, V>,
219         comp: &'a C,
220     ) -> MapCursorMut<'a, K, V, C> {
221         MapCursorMut::new(self, forest, comp)
222     }
223 
224     /// Create an iterator traversing this map. The iterator type is `(K, V)`.
iter<'a>(&'a self, forest: &'a MapForest<K, V>) -> MapIter<'a, K, V>225     pub fn iter<'a>(&'a self, forest: &'a MapForest<K, V>) -> MapIter<'a, K, V> {
226         MapIter {
227             root: self.root,
228             pool: &forest.nodes,
229             path: Path::default(),
230         }
231     }
232 
233     /// Consume this map and its forest, yielding `(K, V)` pairs from the map.
into_iter(self, forest: MapForest<K, V>) -> MapIntoIter<K, V>234     pub fn into_iter(self, forest: MapForest<K, V>) -> MapIntoIter<K, V> {
235         MapIntoIter {
236             root: self.root,
237             pool: forest.nodes,
238             path: Path::default(),
239         }
240     }
241 
242     /// Iterate over the entries within the given range.
range<'a, R, C>( &'a self, range: R, forest: &'a MapForest<K, V>, comp: &'a C, ) -> MapRange<'a, K, V, C> where R: RangeBounds<K>, C: Comparator<K>,243     pub fn range<'a, R, C>(
244         &'a self,
245         range: R,
246         forest: &'a MapForest<K, V>,
247         comp: &'a C,
248     ) -> MapRange<'a, K, V, C>
249     where
250         R: RangeBounds<K>,
251         C: Comparator<K>,
252     {
253         MapRange {
254             cursor: self.cursor(forest, comp),
255             start: copied_bound(range.start_bound()),
256             end: copied_bound(range.end_bound()),
257             direction: Direction::None,
258         }
259     }
260 }
261 
copied_bound<K>(start_bound: Bound<&K>) -> Bound<K> where K: Copy,262 fn copied_bound<K>(start_bound: Bound<&K>) -> Bound<K>
263 where
264     K: Copy,
265 {
266     match start_bound {
267         Bound::Unbounded => Bound::Unbounded,
268         Bound::Included(x) => Bound::Included(*x),
269         Bound::Excluded(x) => Bound::Excluded(*x),
270     }
271 }
272 
273 impl<K, V> Default for Map<K, V>
274 where
275     K: Copy,
276     V: Copy,
277 {
default() -> Self278     fn default() -> Self {
279         Self::new()
280     }
281 }
282 
283 #[cfg(test)]
284 impl<K, V> Map<K, V>
285 where
286     K: Copy + fmt::Display,
287     V: Copy,
288 {
289     /// Verify consistency.
verify<C: Comparator<K>>(&self, forest: &MapForest<K, V>, comp: &C) where NodeData<MapTypes<K, V>>: fmt::Display,290     fn verify<C: Comparator<K>>(&self, forest: &MapForest<K, V>, comp: &C)
291     where
292         NodeData<MapTypes<K, V>>: fmt::Display,
293     {
294         if let Some(root) = self.root.expand() {
295             forest.nodes.verify_tree(root, comp);
296         }
297     }
298 
299     /// Get a text version of the path to `key`.
tpath<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> String300     fn tpath<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> String {
301         use alloc::string::ToString;
302         match self.root.expand() {
303             None => "map(empty)".to_string(),
304             Some(root) => {
305                 let mut path = Path::default();
306                 path.find(key, root, &forest.nodes, comp);
307                 path.to_string()
308             }
309         }
310     }
311 }
312 
313 struct MapCursorRaw<K, V>
314 where
315     K: Copy,
316     V: Copy,
317 {
318     path: Path<MapTypes<K, V>>,
319 }
320 
321 impl<K, V> MapCursorRaw<K, V>
322 where
323     K: Copy,
324     V: Copy,
325 {
new() -> Self326     fn new() -> Self {
327         Self {
328             path: Path::default(),
329         }
330     }
331 
is_empty(&self, root: &PackedOption<Node>) -> bool332     fn is_empty(&self, root: &PackedOption<Node>) -> bool {
333         root.is_none()
334     }
335 
next(&mut self, pool: &NodePool<MapTypes<K, V>>) -> Option<(K, V)>336     fn next(&mut self, pool: &NodePool<MapTypes<K, V>>) -> Option<(K, V)> {
337         self.path.next(pool)
338     }
339 
prev( &mut self, root: &PackedOption<Node>, pool: &NodePool<MapTypes<K, V>>, ) -> Option<(K, V)>340     fn prev(
341         &mut self,
342         root: &PackedOption<Node>,
343         pool: &NodePool<MapTypes<K, V>>,
344     ) -> Option<(K, V)> {
345         root.expand().and_then(|root| self.path.prev(root, pool))
346     }
347 
key(&self, pool: &NodePool<MapTypes<K, V>>) -> Option<K>348     fn key(&self, pool: &NodePool<MapTypes<K, V>>) -> Option<K> {
349         self.path
350             .leaf_pos()
351             .and_then(|(node, entry)| pool[node].unwrap_leaf().0.get(entry).cloned())
352     }
353 
value(&self, pool: &NodePool<MapTypes<K, V>>) -> Option<V>354     fn value(&self, pool: &NodePool<MapTypes<K, V>>) -> Option<V> {
355         self.path
356             .leaf_pos()
357             .and_then(|(node, entry)| pool[node].unwrap_leaf().1.get(entry).cloned())
358     }
359 
value_mut<'a>(&self, pool: &'a mut NodePool<MapTypes<K, V>>) -> Option<&'a mut V>360     fn value_mut<'a>(&self, pool: &'a mut NodePool<MapTypes<K, V>>) -> Option<&'a mut V> {
361         self.path
362             .leaf_pos()
363             .and_then(move |(node, entry)| pool[node].unwrap_leaf_mut().1.get_mut(entry))
364     }
365 
goto( &mut self, root: &PackedOption<Node>, pool: &NodePool<MapTypes<K, V>>, elem: K, comp: &impl Comparator<K>, ) -> Option<V>366     fn goto(
367         &mut self,
368         root: &PackedOption<Node>,
369         pool: &NodePool<MapTypes<K, V>>,
370         elem: K,
371         comp: &impl Comparator<K>,
372     ) -> Option<V> {
373         root.expand().and_then(|root| {
374             let v = self.path.find(elem, root, pool, comp);
375             if v.is_none() {
376                 self.path.normalize(pool);
377             }
378             v
379         })
380     }
381 
goto_first( &mut self, root: &PackedOption<Node>, pool: &NodePool<MapTypes<K, V>>, ) -> Option<V>382     fn goto_first(
383         &mut self,
384         root: &PackedOption<Node>,
385         pool: &NodePool<MapTypes<K, V>>,
386     ) -> Option<V> {
387         root.map(|root| self.path.first(root, pool).1)
388     }
389 
goto_end(&mut self)390     fn goto_end(&mut self) {
391         self.path = Path::default();
392     }
393 
insert( &mut self, root: &mut PackedOption<Node>, pool: &mut NodePool<MapTypes<K, V>>, key: K, value: V, comp: &impl Comparator<K>, ) -> Option<V>394     fn insert(
395         &mut self,
396         root: &mut PackedOption<Node>,
397         pool: &mut NodePool<MapTypes<K, V>>,
398         key: K,
399         value: V,
400         comp: &impl Comparator<K>,
401     ) -> Option<V> {
402         self.try_insert(root, pool, key, value, comp).panic_on_oom()
403     }
404 
try_insert( &mut self, root: &mut PackedOption<Node>, pool: &mut NodePool<MapTypes<K, V>>, key: K, value: V, comp: &impl Comparator<K>, ) -> Result<Option<V>, OutOfMemory>405     fn try_insert(
406         &mut self,
407         root: &mut PackedOption<Node>,
408         pool: &mut NodePool<MapTypes<K, V>>,
409         key: K,
410         value: V,
411         comp: &impl Comparator<K>,
412     ) -> Result<Option<V>, OutOfMemory> {
413         match root.expand() {
414             None => {
415                 let node = pool.alloc_node(NodeData::leaf(key, value))?;
416                 *root = node.into();
417                 self.path.set_root_node(node);
418                 Ok(None)
419             }
420             Some(r) => {
421                 // TODO: Optimize the case where `self.path` is already at the correct insert pos.
422                 let old = self.path.find(key, r, pool, comp);
423                 if old.is_some() {
424                     *self.path.value_mut(pool) = value;
425                 } else {
426                     *root = self.path.insert(key, value, pool)?.into();
427                 }
428                 Ok(old)
429             }
430         }
431     }
432 
remove( &mut self, root: &mut PackedOption<Node>, pool: &mut NodePool<MapTypes<K, V>>, ) -> Option<V>433     fn remove(
434         &mut self,
435         root: &mut PackedOption<Node>,
436         pool: &mut NodePool<MapTypes<K, V>>,
437     ) -> Option<V> {
438         let value = self.value(pool);
439         if value.is_some() {
440             *root = self.path.remove(pool).into();
441         }
442         value
443     }
444 }
445 
446 /// A position in a `Map` used to navigate and modify the ordered map.
447 ///
448 /// A cursor always points at a key-value pair in the map, or "off the end" which is a position
449 /// after the last entry in the map.
450 pub struct MapCursor<'a, K, V, C>
451 where
452     K: 'a + Copy,
453     V: 'a + Copy,
454     C: 'a + Comparator<K>,
455 {
456     root: &'a PackedOption<Node>,
457     pool: &'a NodePool<MapTypes<K, V>>,
458     comp: &'a C,
459     raw: MapCursorRaw<K, V>,
460 }
461 
462 impl<'a, K, V, C> MapCursor<'a, K, V, C>
463 where
464     K: Copy,
465     V: Copy,
466     C: Comparator<K>,
467 {
468     /// Create a cursor with a default (off-the-end) location.
new(container: &'a Map<K, V>, forest: &'a MapForest<K, V>, comp: &'a C) -> Self469     fn new(container: &'a Map<K, V>, forest: &'a MapForest<K, V>, comp: &'a C) -> Self {
470         Self {
471             root: &container.root,
472             pool: &forest.nodes,
473             comp,
474             raw: MapCursorRaw::new(),
475         }
476     }
477 
478     /// Is this cursor pointing to an empty map?
is_empty(&self) -> bool479     pub fn is_empty(&self) -> bool {
480         self.raw.is_empty(self.root)
481     }
482 
483     /// Move cursor to the next key-value pair and return it.
484     ///
485     /// If the cursor reaches the end, return `None` and leave the cursor at the off-the-end
486     /// position.
next(&mut self) -> Option<(K, V)>487     pub fn next(&mut self) -> Option<(K, V)> {
488         self.raw.next(self.pool)
489     }
490 
491     /// Move cursor to the previous key-value pair and return it.
492     ///
493     /// If the cursor is already pointing at the first entry, leave it there and return `None`.
prev(&mut self) -> Option<(K, V)>494     pub fn prev(&mut self) -> Option<(K, V)> {
495         self.raw.prev(self.root, self.pool)
496     }
497 
498     /// Get the current key, or `None` if the cursor is at the end.
key(&self) -> Option<K>499     pub fn key(&self) -> Option<K> {
500         self.raw.key(self.pool)
501     }
502 
503     /// Get the current value, or `None` if the cursor is at the end.
value(&self) -> Option<V>504     pub fn value(&self) -> Option<V> {
505         self.raw.value(self.pool)
506     }
507 
508     /// Move this cursor to `key`.
509     ///
510     /// If `key` is in the map, place the cursor at `key` and return the corresponding value.
511     ///
512     /// If `key` is not in the set, place the cursor at the next larger element (or the end) and
513     /// return `None`.
goto(&mut self, elem: K) -> Option<V>514     pub fn goto(&mut self, elem: K) -> Option<V> {
515         self.raw.goto(self.root, self.pool, elem, self.comp)
516     }
517 
518     /// Move this cursor to the first element.
goto_first(&mut self) -> Option<V>519     pub fn goto_first(&mut self) -> Option<V> {
520         self.raw.goto_first(self.root, self.pool)
521     }
522 
523     /// Move the cursor to the off-the-end location.
goto_end(&mut self)524     pub fn goto_end(&mut self) {
525         self.raw.goto_end();
526     }
527 }
528 
529 /// A position in a `Map` used to navigate and modify the ordered map.
530 ///
531 /// A cursor always points at a key-value pair in the map, or "off the end" which is a position
532 /// after the last entry in the map.
533 pub struct MapCursorMut<'a, K, V, C>
534 where
535     K: 'a + Copy,
536     V: 'a + Copy,
537     C: 'a + Comparator<K>,
538 {
539     root: &'a mut PackedOption<Node>,
540     pool: &'a mut NodePool<MapTypes<K, V>>,
541     comp: &'a C,
542     raw: MapCursorRaw<K, V>,
543 }
544 
545 impl<'a, K, V, C> MapCursorMut<'a, K, V, C>
546 where
547     K: Copy,
548     V: Copy,
549     C: Comparator<K>,
550 {
551     /// Create a cursor with a default (off-the-end) location.
new(container: &'a mut Map<K, V>, forest: &'a mut MapForest<K, V>, comp: &'a C) -> Self552     fn new(container: &'a mut Map<K, V>, forest: &'a mut MapForest<K, V>, comp: &'a C) -> Self {
553         Self {
554             root: &mut container.root,
555             pool: &mut forest.nodes,
556             comp,
557             raw: MapCursorRaw::new(),
558         }
559     }
560 
561     /// Is this cursor pointing to an empty map?
is_empty(&self) -> bool562     pub fn is_empty(&self) -> bool {
563         self.raw.is_empty(self.root)
564     }
565 
566     /// Move cursor to the next key-value pair and return it.
567     ///
568     /// If the cursor reaches the end, return `None` and leave the cursor at the off-the-end
569     /// position.
next(&mut self) -> Option<(K, V)>570     pub fn next(&mut self) -> Option<(K, V)> {
571         self.raw.next(self.pool)
572     }
573 
574     /// Move cursor to the previous key-value pair and return it.
575     ///
576     /// If the cursor is already pointing at the first entry, leave it there and return `None`.
prev(&mut self) -> Option<(K, V)>577     pub fn prev(&mut self) -> Option<(K, V)> {
578         self.raw.prev(self.root, self.pool)
579     }
580 
581     /// Get the current key, or `None` if the cursor is at the end.
key(&self) -> Option<K>582     pub fn key(&self) -> Option<K> {
583         self.raw.key(self.pool)
584     }
585 
586     /// Get the current value, or `None` if the cursor is at the end.
value(&self) -> Option<V>587     pub fn value(&self) -> Option<V> {
588         self.raw.value(self.pool)
589     }
590 
591     /// Get a mutable reference to the current value, or `None` if the cursor is at the end.
value_mut(&mut self) -> Option<&mut V>592     pub fn value_mut(&mut self) -> Option<&mut V> {
593         self.raw.value_mut(self.pool)
594     }
595 
596     /// Move this cursor to `key`.
597     ///
598     /// If `key` is in the map, place the cursor at `key` and return the corresponding value.
599     ///
600     /// If `key` is not in the set, place the cursor at the next larger element (or the end) and
601     /// return `None`.
goto(&mut self, elem: K) -> Option<V>602     pub fn goto(&mut self, elem: K) -> Option<V> {
603         self.raw.goto(self.root, self.pool, elem, self.comp)
604     }
605 
606     /// Move this cursor to the first element.
goto_first(&mut self) -> Option<V>607     pub fn goto_first(&mut self) -> Option<V> {
608         self.raw.goto_first(self.root, self.pool)
609     }
610 
611     /// Insert `(key, value))` into the map and leave the cursor at the inserted pair.
612     ///
613     /// If the map did not contain `key`, return `None`.
614     ///
615     /// If `key` is already present, replace the existing with `value` and return the old value.
insert(&mut self, key: K, value: V) -> Option<V>616     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
617         self.raw.insert(self.root, self.pool, key, value, self.comp)
618     }
619 
620     /// Like `insert` but returns an error on allocation failure.
try_insert(&mut self, key: K, value: V) -> Result<Option<V>, OutOfMemory>621     pub fn try_insert(&mut self, key: K, value: V) -> Result<Option<V>, OutOfMemory> {
622         self.raw
623             .try_insert(self.root, self.pool, key, value, self.comp)
624     }
625 
626     /// Remove the current entry (if any) and return the mapped value.
627     /// This advances the cursor to the next entry after the removed one.
remove(&mut self) -> Option<V>628     pub fn remove(&mut self) -> Option<V> {
629         self.raw.remove(self.root, self.pool)
630     }
631 }
632 
633 /// An iterator visiting the key-value pairs of a `Map`.
634 pub struct MapIter<'a, K, V>
635 where
636     K: 'a + Copy,
637     V: 'a + Copy,
638 {
639     root: PackedOption<Node>,
640     pool: &'a NodePool<MapTypes<K, V>>,
641     path: Path<MapTypes<K, V>>,
642 }
643 
644 impl<'a, K, V> Iterator for MapIter<'a, K, V>
645 where
646     K: 'a + Copy,
647     V: 'a + Copy,
648 {
649     type Item = (K, V);
650 
next(&mut self) -> Option<Self::Item>651     fn next(&mut self) -> Option<Self::Item> {
652         // We use `self.root` to indicate if we need to go to the first element. Reset to `None`
653         // once we've returned the first element. This also works for an empty tree since the
654         // `path.next()` call returns `None` when the path is empty. This also fuses the iterator.
655         match self.root.take() {
656             Some(root) => Some(self.path.first(root, self.pool)),
657             None => self.path.next(self.pool),
658         }
659     }
660 }
661 
662 /// A consuming iterator of a `Map`, yielding its key-value pairs.
663 pub struct MapIntoIter<K, V>
664 where
665     K: Copy,
666     V: Copy,
667 {
668     root: PackedOption<Node>,
669     pool: NodePool<MapTypes<K, V>>,
670     path: Path<MapTypes<K, V>>,
671 }
672 
673 impl<K, V> Iterator for MapIntoIter<K, V>
674 where
675     K: Copy,
676     V: Copy,
677 {
678     type Item = (K, V);
679 
next(&mut self) -> Option<Self::Item>680     fn next(&mut self) -> Option<Self::Item> {
681         // We use `self.root` to indicate if we need to go to the first element. Reset to `None`
682         // once we've returned the first element. This also works for an empty tree since the
683         // `path.next()` call returns `None` when the path is empty. This also fuses the iterator.
684         match self.root.take() {
685             Some(root) => Some(self.path.first(root, &self.pool)),
686             None => self.path.next(&self.pool),
687         }
688     }
689 }
690 
691 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
692 enum Direction {
693     None,
694     Next,
695     NextBack,
696 }
697 
698 /// An immutable iterator over a range of values, returned by `Map::range`.
699 pub struct MapRange<'a, K, V, C>
700 where
701     K: Copy,
702     V: Copy,
703     C: Comparator<K>,
704 {
705     cursor: MapCursor<'a, K, V, C>,
706     start: Bound<K>,
707     end: Bound<K>,
708 
709     /// The direction of the last call into the iterator, a.k.a. whether
710     /// `self.cursor` is
711     ///
712     /// 1. `Direction::Next`: just before `self.start`,
713     ///
714     /// 2. `Direction::NextBack`: just after `self.end`, or
715     ///
716     /// 3. `Direction::None`: not initialized to any location yet.
717     direction: Direction,
718 }
719 
720 impl<'a, K, V, C> Iterator for MapRange<'a, K, V, C>
721 where
722     K: Copy,
723     V: Copy,
724     C: Comparator<K>,
725 {
726     type Item = (K, V);
727 
next(&mut self) -> Option<Self::Item>728     fn next(&mut self) -> Option<Self::Item> {
729         let old_direction = mem::replace(&mut self.direction, Direction::Next);
730 
731         let key_val = (|| {
732             if old_direction == Direction::Next {
733                 self.cursor.next()
734             } else {
735                 match self.start {
736                     Bound::Included(k) => {
737                         self.cursor.goto(k);
738                         Some((self.cursor.key()?, self.cursor.value()?))
739                     }
740                     Bound::Excluded(k) => {
741                         self.cursor.goto(k);
742                         if self
743                             .cursor
744                             .key()
745                             .is_some_and(|key| self.cursor.comp.cmp(k, key).is_eq())
746                         {
747                             self.cursor.next()
748                         } else {
749                             Some((self.cursor.key()?, self.cursor.value()?))
750                         }
751                     }
752                     Bound::Unbounded => {
753                         let val = self.cursor.goto_first()?;
754                         Some((self.cursor.key()?, val))
755                     }
756                 }
757             }
758         })();
759 
760         match key_val {
761             Some((key, val)) if self.is_in_bounds(key) => {
762                 self.start = Bound::Excluded(key);
763                 Some((key, val))
764             }
765             _ => {
766                 self.cursor.goto_end();
767                 None
768             }
769         }
770     }
771 }
772 
773 impl<'a, K, V, C> DoubleEndedIterator for MapRange<'a, K, V, C>
774 where
775     K: Copy,
776     V: Copy,
777     C: Comparator<K>,
778 {
next_back(&mut self) -> Option<Self::Item>779     fn next_back(&mut self) -> Option<Self::Item> {
780         let old_direction = mem::replace(&mut self.direction, Direction::NextBack);
781 
782         let key_val = (|| {
783             if old_direction == Direction::NextBack {
784                 self.cursor.prev()
785             } else {
786                 match self.end {
787                     Bound::Included(k) => match self.cursor.goto(k) {
788                         Some(_) => Some((self.cursor.key()?, self.cursor.value()?)),
789                         None => self.cursor.prev(),
790                     },
791                     Bound::Excluded(k) => {
792                         self.cursor.goto(k);
793                         if self
794                             .cursor
795                             .key()
796                             .is_some_and(|key| self.cursor.comp.cmp(k, key).is_eq())
797                         {
798                             self.cursor.prev()
799                         } else {
800                             Some((self.cursor.key()?, self.cursor.value()?))
801                         }
802                     }
803                     Bound::Unbounded => {
804                         self.cursor.goto_end();
805                         self.cursor.prev()
806                     }
807                 }
808             }
809         })();
810 
811         match key_val {
812             Some((key, val)) if self.is_in_bounds(key) => {
813                 self.end = Bound::Excluded(key);
814                 Some((key, val))
815             }
816             _ => {
817                 self.cursor.goto_end();
818                 None
819             }
820         }
821     }
822 }
823 
824 impl<'a, K, V, C> MapRange<'a, K, V, C>
825 where
826     K: Copy,
827     V: Copy,
828     C: Comparator<K>,
829 {
is_in_bounds(&self, key: K) -> bool830     fn is_in_bounds(&self, key: K) -> bool {
831         self.is_in_start_bound(key) && self.is_in_end_bound(key)
832     }
833 
is_in_start_bound(&self, key: K) -> bool834     fn is_in_start_bound(&self, key: K) -> bool {
835         match self.start {
836             Bound::Included(k) => self.cursor.comp.cmp(key, k).is_ge(),
837             Bound::Excluded(k) => self.cursor.comp.cmp(key, k).is_gt(),
838             Bound::Unbounded => true,
839         }
840     }
841 
is_in_end_bound(&self, key: K) -> bool842     fn is_in_end_bound(&self, key: K) -> bool {
843         match self.end {
844             Bound::Included(k) => self.cursor.comp.cmp(k, key).is_ge(),
845             Bound::Excluded(k) => self.cursor.comp.cmp(k, key).is_gt(),
846             Bound::Unbounded => true,
847         }
848     }
849 }
850 
851 #[cfg(test)]
852 impl<'a, K, V, C> MapCursorMut<'a, K, V, C>
853 where
854     K: Copy + fmt::Display,
855     V: Copy + fmt::Display,
856     C: Comparator<K>,
857 {
verify(&self)858     fn verify(&self) {
859         self.raw.path.verify(self.pool);
860         self.root.map(|root| self.pool.verify_tree(root, self.comp));
861     }
862 
863     /// Get a text version of the path to the current position.
tpath(&self) -> String864     fn tpath(&self) -> String {
865         use alloc::string::ToString;
866         self.raw.path.to_string()
867     }
868 }
869 
870 #[cfg(test)]
871 mod tests {
872     use super::*;
873     use alloc::{vec, vec::Vec};
874     use core::mem;
875 
876     #[test]
node_size()877     fn node_size() {
878         // check that nodes are cache line sized when keys and values are 32 bits.
879         type F = MapTypes<u32, u32>;
880         assert_eq!(mem::size_of::<NodeData<F>>(), 64);
881     }
882 
883     #[test]
empty()884     fn empty() {
885         let mut f = MapForest::<u32, f32>::new();
886         f.clear();
887 
888         let mut m = Map::<u32, f32>::new();
889         assert!(m.is_empty());
890         m.clear(&mut f);
891 
892         assert_eq!(m.get(7, &f, &()), None);
893         assert_eq!(m.iter(&f).next(), None);
894         assert_eq!(m.get_or_less(7, &f, &()), None);
895         m.retain(&mut f, |_, _| unreachable!());
896 
897         let mut c = m.cursor_mut(&mut f, &());
898         assert!(c.is_empty());
899         assert_eq!(c.key(), None);
900         assert_eq!(c.value(), None);
901         assert_eq!(c.next(), None);
902         assert_eq!(c.prev(), None);
903         c.verify();
904         assert_eq!(c.tpath(), "<empty path>");
905         assert_eq!(c.goto_first(), None);
906         assert_eq!(c.tpath(), "<empty path>");
907     }
908 
909     #[test]
inserting()910     fn inserting() {
911         let f = &mut MapForest::<u32, f32>::new();
912         let mut m = Map::<u32, f32>::new();
913 
914         // The first seven values stay in a single leaf node.
915         assert_eq!(m.insert(50, 5.0, f, &()), None);
916         assert_eq!(m.insert(50, 5.5, f, &()), Some(5.0));
917         assert_eq!(m.insert(20, 2.0, f, &()), None);
918         assert_eq!(m.insert(80, 8.0, f, &()), None);
919         assert_eq!(m.insert(40, 4.0, f, &()), None);
920         assert_eq!(m.insert(60, 6.0, f, &()), None);
921         assert_eq!(m.insert(90, 9.0, f, &()), None);
922         assert_eq!(m.insert(200, 20.0, f, &()), None);
923 
924         m.verify(f, &());
925 
926         assert_eq!(
927             m.iter(f).collect::<Vec<_>>(),
928             [
929                 (20, 2.0),
930                 (40, 4.0),
931                 (50, 5.5),
932                 (60, 6.0),
933                 (80, 8.0),
934                 (90, 9.0),
935                 (200, 20.0),
936             ]
937         );
938 
939         assert_eq!(m.get(0, f, &()), None);
940         assert_eq!(m.get(20, f, &()), Some(2.0));
941         assert_eq!(m.get(30, f, &()), None);
942         assert_eq!(m.get(40, f, &()), Some(4.0));
943         assert_eq!(m.get(50, f, &()), Some(5.5));
944         assert_eq!(m.get(60, f, &()), Some(6.0));
945         assert_eq!(m.get(70, f, &()), None);
946         assert_eq!(m.get(80, f, &()), Some(8.0));
947         assert_eq!(m.get(100, f, &()), None);
948 
949         assert_eq!(m.get_or_less(0, f, &()), None);
950         assert_eq!(m.get_or_less(20, f, &()), Some((20, 2.0)));
951         assert_eq!(m.get_or_less(30, f, &()), Some((20, 2.0)));
952         assert_eq!(m.get_or_less(40, f, &()), Some((40, 4.0)));
953         assert_eq!(m.get_or_less(200, f, &()), Some((200, 20.0)));
954         assert_eq!(m.get_or_less(201, f, &()), Some((200, 20.0)));
955 
956         {
957             let mut c = m.cursor_mut(f, &());
958             assert_eq!(c.prev(), Some((200, 20.0)));
959             assert_eq!(c.prev(), Some((90, 9.0)));
960             assert_eq!(c.prev(), Some((80, 8.0)));
961             assert_eq!(c.prev(), Some((60, 6.0)));
962             assert_eq!(c.prev(), Some((50, 5.5)));
963             assert_eq!(c.prev(), Some((40, 4.0)));
964             assert_eq!(c.prev(), Some((20, 2.0)));
965             assert_eq!(c.prev(), None);
966         }
967 
968         // Test some removals where the node stays healthy.
969         assert_eq!(m.tpath(50, f, &()), "node0[2]");
970         assert_eq!(m.tpath(80, f, &()), "node0[4]");
971         assert_eq!(m.tpath(200, f, &()), "node0[6]");
972 
973         assert_eq!(m.remove(80, f, &()), Some(8.0));
974         assert_eq!(m.tpath(50, f, &()), "node0[2]");
975         assert_eq!(m.tpath(80, f, &()), "node0[4]");
976         assert_eq!(m.tpath(200, f, &()), "node0[5]");
977         assert_eq!(m.remove(80, f, &()), None);
978         m.verify(f, &());
979 
980         assert_eq!(m.remove(20, f, &()), Some(2.0));
981         assert_eq!(m.tpath(50, f, &()), "node0[1]");
982         assert_eq!(m.tpath(80, f, &()), "node0[3]");
983         assert_eq!(m.tpath(200, f, &()), "node0[4]");
984         assert_eq!(m.remove(20, f, &()), None);
985         m.verify(f, &());
986 
987         // [ 40 50 60 90 200 ]
988 
989         {
990             let mut c = m.cursor_mut(f, &());
991             assert_eq!(c.goto_first(), Some(4.0));
992             assert_eq!(c.key(), Some(40));
993             assert_eq!(c.value(), Some(4.0));
994             assert_eq!(c.next(), Some((50, 5.5)));
995             assert_eq!(c.next(), Some((60, 6.0)));
996             assert_eq!(c.next(), Some((90, 9.0)));
997             assert_eq!(c.next(), Some((200, 20.0)));
998             c.verify();
999             assert_eq!(c.next(), None);
1000             c.verify();
1001         }
1002 
1003         // Removals from the root leaf node beyond underflow.
1004         assert_eq!(m.remove(200, f, &()), Some(20.0));
1005         assert_eq!(m.remove(40, f, &()), Some(4.0));
1006         assert_eq!(m.remove(60, f, &()), Some(6.0));
1007         m.verify(f, &());
1008         assert_eq!(m.remove(50, f, &()), Some(5.5));
1009         m.verify(f, &());
1010         assert_eq!(m.remove(90, f, &()), Some(9.0));
1011         m.verify(f, &());
1012         assert!(m.is_empty());
1013     }
1014 
1015     #[test]
split_level0_leaf()1016     fn split_level0_leaf() {
1017         // Various ways of splitting a full leaf node at level 0.
1018         let f = &mut MapForest::<u32, f32>::new();
1019 
1020         fn full_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
1021             let mut m = Map::new();
1022             for n in 1..8 {
1023                 m.insert(n * 10, n as f32 * 1.1, f, &());
1024             }
1025             m
1026         }
1027 
1028         // Insert at front of leaf.
1029         let mut m = full_leaf(f);
1030         m.insert(5, 4.2, f, &());
1031         m.verify(f, &());
1032         assert_eq!(m.get(5, f, &()), Some(4.2));
1033 
1034         // Retain even entries, with altered values.
1035         m.retain(f, |k, v| {
1036             *v = (k / 10) as f32;
1037             (k % 20) == 0
1038         });
1039         assert_eq!(
1040             m.iter(f).collect::<Vec<_>>(),
1041             [(20, 2.0), (40, 4.0), (60, 6.0)]
1042         );
1043 
1044         // Insert at back of leaf.
1045         let mut m = full_leaf(f);
1046         m.insert(80, 4.2, f, &());
1047         m.verify(f, &());
1048         assert_eq!(m.get(80, f, &()), Some(4.2));
1049 
1050         // Insert before middle (40).
1051         let mut m = full_leaf(f);
1052         m.insert(35, 4.2, f, &());
1053         m.verify(f, &());
1054         assert_eq!(m.get(35, f, &()), Some(4.2));
1055 
1056         // Insert after middle (40).
1057         let mut m = full_leaf(f);
1058         m.insert(45, 4.2, f, &());
1059         m.verify(f, &());
1060         assert_eq!(m.get(45, f, &()), Some(4.2));
1061 
1062         m.clear(f);
1063         assert!(m.is_empty());
1064     }
1065 
1066     #[test]
split_level1_leaf()1067     fn split_level1_leaf() {
1068         // Various ways of splitting a full leaf node at level 1.
1069         let f = &mut MapForest::<u32, f32>::new();
1070 
1071         // Return a map whose root node is a full inner node, and the leaf nodes are all full
1072         // containing:
1073         //
1074         // 110, 120, ..., 170
1075         // 210, 220, ..., 270
1076         // ...
1077         // 810, 820, ..., 870
1078         fn full(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
1079             let mut m = Map::new();
1080 
1081             // Start by inserting elements in order.
1082             // This should leave 8 leaf nodes with 4 elements in each.
1083             for row in 1..9 {
1084                 for col in 1..5 {
1085                     m.insert(row * 100 + col * 10, row as f32 + col as f32 * 0.1, f, &());
1086                 }
1087             }
1088 
1089             // Then top up the leaf nodes without splitting them.
1090             for row in 1..9 {
1091                 for col in 5..8 {
1092                     m.insert(row * 100 + col * 10, row as f32 + col as f32 * 0.1, f, &());
1093                 }
1094             }
1095 
1096             m
1097         }
1098 
1099         let mut m = full(f);
1100         // Verify geometry. Get get node2 as the root and leaves node0, 1, 3, ...
1101         m.verify(f, &());
1102         assert_eq!(m.tpath(110, f, &()), "node2[0]--node0[0]");
1103         assert_eq!(m.tpath(140, f, &()), "node2[0]--node0[3]");
1104         assert_eq!(m.tpath(210, f, &()), "node2[1]--node1[0]");
1105         assert_eq!(m.tpath(270, f, &()), "node2[1]--node1[6]");
1106         assert_eq!(m.tpath(310, f, &()), "node2[2]--node3[0]");
1107         assert_eq!(m.tpath(810, f, &()), "node2[7]--node8[0]");
1108         assert_eq!(m.tpath(870, f, &()), "node2[7]--node8[6]");
1109 
1110         {
1111             let mut c = m.cursor_mut(f, &());
1112             assert_eq!(c.goto_first(), Some(1.1));
1113             assert_eq!(c.key(), Some(110));
1114         }
1115 
1116         // Front of first leaf.
1117         m.insert(0, 4.2, f, &());
1118         m.verify(f, &());
1119         assert_eq!(m.get(0, f, &()), Some(4.2));
1120 
1121         // First leaf split 4-4 after appending to LHS.
1122         f.clear();
1123         m = full(f);
1124         m.insert(135, 4.2, f, &());
1125         m.verify(f, &());
1126         assert_eq!(m.get(135, f, &()), Some(4.2));
1127 
1128         // First leaf split 4-4 after prepending to RHS.
1129         f.clear();
1130         m = full(f);
1131         m.insert(145, 4.2, f, &());
1132         m.verify(f, &());
1133         assert_eq!(m.get(145, f, &()), Some(4.2));
1134 
1135         // First leaf split 4-4 after appending to RHS.
1136         f.clear();
1137         m = full(f);
1138         m.insert(175, 4.2, f, &());
1139         m.verify(f, &());
1140         assert_eq!(m.get(175, f, &()), Some(4.2));
1141 
1142         // Left-middle leaf split, ins LHS.
1143         f.clear();
1144         m = full(f);
1145         m.insert(435, 4.2, f, &());
1146         m.verify(f, &());
1147         assert_eq!(m.get(435, f, &()), Some(4.2));
1148 
1149         // Left-middle leaf split, ins RHS.
1150         f.clear();
1151         m = full(f);
1152         m.insert(445, 4.2, f, &());
1153         m.verify(f, &());
1154         assert_eq!(m.get(445, f, &()), Some(4.2));
1155 
1156         // Right-middle leaf split, ins LHS.
1157         f.clear();
1158         m = full(f);
1159         m.insert(535, 4.2, f, &());
1160         m.verify(f, &());
1161         assert_eq!(m.get(535, f, &()), Some(4.2));
1162 
1163         // Right-middle leaf split, ins RHS.
1164         f.clear();
1165         m = full(f);
1166         m.insert(545, 4.2, f, &());
1167         m.verify(f, &());
1168         assert_eq!(m.get(545, f, &()), Some(4.2));
1169 
1170         // Last leaf split, ins LHS.
1171         f.clear();
1172         m = full(f);
1173         m.insert(835, 4.2, f, &());
1174         m.verify(f, &());
1175         assert_eq!(m.get(835, f, &()), Some(4.2));
1176 
1177         // Last leaf split, ins RHS.
1178         f.clear();
1179         m = full(f);
1180         m.insert(845, 4.2, f, &());
1181         m.verify(f, &());
1182         assert_eq!(m.get(845, f, &()), Some(4.2));
1183 
1184         // Front of last leaf.
1185         f.clear();
1186         m = full(f);
1187         m.insert(805, 4.2, f, &());
1188         m.verify(f, &());
1189         assert_eq!(m.get(805, f, &()), Some(4.2));
1190 
1191         m.clear(f);
1192         m.verify(f, &());
1193     }
1194 
1195     // Make a tree with two barely healthy leaf nodes:
1196     // [ 10 20 30 40 ] [ 50 60 70 80 ]
two_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32>1197     fn two_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
1198         f.clear();
1199         let mut m = Map::new();
1200         for n in 1..9 {
1201             m.insert(n * 10, n as f32, f, &());
1202         }
1203         m
1204     }
1205 
1206     #[test]
remove_level1()1207     fn remove_level1() {
1208         let f = &mut MapForest::<u32, f32>::new();
1209         let mut m = two_leaf(f);
1210 
1211         // Verify geometry.
1212         m.verify(f, &());
1213         assert_eq!(m.tpath(10, f, &()), "node2[0]--node0[0]");
1214         assert_eq!(m.tpath(40, f, &()), "node2[0]--node0[3]");
1215         assert_eq!(m.tpath(49, f, &()), "node2[0]--node0[4]");
1216         assert_eq!(m.tpath(50, f, &()), "node2[1]--node1[0]");
1217         assert_eq!(m.tpath(80, f, &()), "node2[1]--node1[3]");
1218 
1219         // Remove the front entry from a node that stays healthy.
1220         assert_eq!(m.insert(55, 5.5, f, &()), None);
1221         assert_eq!(m.remove(50, f, &()), Some(5.0));
1222         m.verify(f, &());
1223         assert_eq!(m.tpath(49, f, &()), "node2[0]--node0[4]");
1224         assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[4]");
1225         assert_eq!(m.tpath(55, f, &()), "node2[1]--node1[0]");
1226 
1227         // Remove the front entry from the first leaf node: No critical key to update.
1228         assert_eq!(m.insert(15, 1.5, f, &()), None);
1229         assert_eq!(m.remove(10, f, &()), Some(1.0));
1230         m.verify(f, &());
1231 
1232         // [ 15 20 30 40 ] [ 55 60 70 80 ]
1233 
1234         // Remove the front entry from a right-most node that underflows.
1235         // No rebalancing for the right-most node. Still need critical key update.
1236         assert_eq!(m.remove(55, f, &()), Some(5.5));
1237         m.verify(f, &());
1238         assert_eq!(m.tpath(55, f, &()), "node2[0]--node0[4]");
1239         assert_eq!(m.tpath(60, f, &()), "node2[1]--node1[0]");
1240 
1241         // [ 15 20 30 40 ] [ 60 70 80 ]
1242 
1243         // Replenish the right leaf.
1244         assert_eq!(m.insert(90, 9.0, f, &()), None);
1245         assert_eq!(m.insert(100, 10.0, f, &()), None);
1246         m.verify(f, &());
1247         assert_eq!(m.tpath(55, f, &()), "node2[0]--node0[4]");
1248         assert_eq!(m.tpath(60, f, &()), "node2[1]--node1[0]");
1249 
1250         // [ 15 20 30 40 ] [ 60 70 80 90 100 ]
1251 
1252         // Removing one entry from the left leaf should trigger a rebalancing from the right
1253         // sibling.
1254         assert_eq!(m.remove(20, f, &()), Some(2.0));
1255         m.verify(f, &());
1256 
1257         // [ 15 30 40 60 ] [ 70 80 90 100 ]
1258         // Check that the critical key was updated correctly.
1259         assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[3]");
1260         assert_eq!(m.tpath(60, f, &()), "node2[0]--node0[3]");
1261         assert_eq!(m.tpath(70, f, &()), "node2[1]--node1[0]");
1262 
1263         // Remove front entry from the left-most leaf node, underflowing.
1264         // This should cause two leaf nodes to be merged and the root node to go away.
1265         assert_eq!(m.remove(15, f, &()), Some(1.5));
1266         m.verify(f, &());
1267     }
1268 
1269     #[test]
remove_level1_rightmost()1270     fn remove_level1_rightmost() {
1271         let f = &mut MapForest::<u32, f32>::new();
1272         let mut m = two_leaf(f);
1273 
1274         // [ 10 20 30 40 ] [ 50 60 70 80 ]
1275 
1276         // Remove entries from the right leaf. This doesn't trigger a rebalancing.
1277         assert_eq!(m.remove(60, f, &()), Some(6.0));
1278         assert_eq!(m.remove(80, f, &()), Some(8.0));
1279         assert_eq!(m.remove(50, f, &()), Some(5.0));
1280         m.verify(f, &());
1281 
1282         // [ 10 20 30 40 ] [ 70 ]
1283         assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[4]");
1284         assert_eq!(m.tpath(70, f, &()), "node2[1]--node1[0]");
1285 
1286         // Removing the last entry from the right leaf should cause a collapse.
1287         assert_eq!(m.remove(70, f, &()), Some(7.0));
1288         m.verify(f, &());
1289     }
1290 
1291     // Make a 3-level tree with barely healthy nodes.
1292     // 1 root, 8 inner nodes, 7*4+5=33 leaf nodes, 4 entries each.
level3_sparse(f: &mut MapForest<u32, f32>) -> Map<u32, f32>1293     fn level3_sparse(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
1294         f.clear();
1295         let mut m = Map::new();
1296         for n in 1..133 {
1297             m.insert(n * 10, n as f32, f, &());
1298         }
1299         m
1300     }
1301 
1302     #[test]
level3_removes()1303     fn level3_removes() {
1304         let f = &mut MapForest::<u32, f32>::new();
1305         let mut m = level3_sparse(f);
1306         m.verify(f, &());
1307 
1308         // Check geometry.
1309         // Root: node11
1310         // [ node2 170 node10 330 node16 490 node21 650 node26 810 node31 970 node36 1130 node41 ]
1311         // L1: node11
1312         assert_eq!(m.tpath(0, f, &()), "node11[0]--node2[0]--node0[0]");
1313         assert_eq!(m.tpath(10000, f, &()), "node11[7]--node41[4]--node40[4]");
1314 
1315         // 650 is a critical key in the middle of the root.
1316         assert_eq!(m.tpath(640, f, &()), "node11[3]--node21[3]--node19[3]");
1317         assert_eq!(m.tpath(650, f, &()), "node11[4]--node26[0]--node20[0]");
1318 
1319         // Deleting 640 triggers a rebalance from node19 to node 20, cascading to n21 -> n26.
1320         assert_eq!(m.remove(640, f, &()), Some(64.0));
1321         m.verify(f, &());
1322         assert_eq!(m.tpath(650, f, &()), "node11[3]--node26[3]--node20[3]");
1323 
1324         // 1130 is in the first leaf of the last L1 node. Deleting it triggers a rebalance node35
1325         // -> node37, but no rebalance above where there is no right sibling.
1326         assert_eq!(m.tpath(1130, f, &()), "node11[6]--node41[0]--node35[0]");
1327         assert_eq!(m.tpath(1140, f, &()), "node11[6]--node41[0]--node35[1]");
1328         assert_eq!(m.remove(1130, f, &()), Some(113.0));
1329         m.verify(f, &());
1330         assert_eq!(m.tpath(1140, f, &()), "node11[6]--node41[0]--node37[0]");
1331     }
1332 
1333     #[test]
insert_many()1334     fn insert_many() {
1335         let f = &mut MapForest::<u32, f32>::new();
1336         let mut m = Map::<u32, f32>::new();
1337 
1338         let mm = 4096;
1339         let mut x = 0;
1340 
1341         for n in 0..mm {
1342             assert_eq!(m.insert(x, n as f32, f, &()), None);
1343             m.verify(f, &());
1344 
1345             x = (x + n + 1) % mm;
1346         }
1347 
1348         x = 0;
1349         for n in 0..mm {
1350             assert_eq!(m.get(x, f, &()), Some(n as f32));
1351             x = (x + n + 1) % mm;
1352         }
1353 
1354         x = 0;
1355         for n in 0..mm {
1356             assert_eq!(m.remove(x, f, &()), Some(n as f32));
1357             m.verify(f, &());
1358 
1359             x = (x + n + 1) % mm;
1360         }
1361 
1362         assert!(m.is_empty());
1363     }
1364 
1365     #[test]
immutable_cursor()1366     fn immutable_cursor() {
1367         let f = &mut MapForest::<u32, f32>::new();
1368         let mut m = Map::<u32, f32>::new();
1369 
1370         for i in 100..200 {
1371             m.insert(i, i as f32, f, &());
1372         }
1373 
1374         let mut c = m.cursor(f, &());
1375 
1376         let v = c.goto_first().unwrap();
1377         assert_eq!(v, 100.0);
1378         assert_eq!(c.key(), Some(100));
1379         assert_eq!(c.value(), Some(100.0));
1380 
1381         let (k, v) = c.next().unwrap();
1382         assert_eq!(k, 101);
1383         assert_eq!(v, 101.0);
1384         assert_eq!(c.key(), Some(101));
1385         assert_eq!(c.value(), Some(101.0));
1386 
1387         let (k, v) = c.next().unwrap();
1388         assert_eq!(k, 102);
1389         assert_eq!(v, 102.0);
1390         assert_eq!(c.key(), Some(102));
1391         assert_eq!(c.value(), Some(102.0));
1392 
1393         let (k, v) = c.prev().unwrap();
1394         assert_eq!(k, 101);
1395         assert_eq!(v, 101.0);
1396         assert_eq!(c.key(), Some(101));
1397         assert_eq!(c.value(), Some(101.0));
1398 
1399         let v = c.goto(175).unwrap();
1400         assert_eq!(v, 175.0);
1401         assert_eq!(c.key(), Some(175));
1402         assert_eq!(c.value(), Some(175.0));
1403 
1404         let (k, v) = c.next().unwrap();
1405         assert_eq!(k, 176);
1406         assert_eq!(v, 176.0);
1407         assert_eq!(c.key(), Some(176));
1408         assert_eq!(c.value(), Some(176.0));
1409 
1410         let (k, v) = c.prev().unwrap();
1411         assert_eq!(k, 175);
1412         assert_eq!(v, 175.0);
1413         assert_eq!(c.key(), Some(175));
1414         assert_eq!(c.value(), Some(175.0));
1415 
1416         let v = c.goto(200);
1417         assert!(v.is_none());
1418         assert!(c.key().is_none());
1419         assert!(c.value().is_none());
1420 
1421         for i in (100..200).rev() {
1422             let (k, v) = c.prev().unwrap();
1423             assert_eq!(k, i);
1424             assert_eq!(v, i as f32);
1425             assert_eq!(c.key(), Some(i));
1426             assert_eq!(c.value(), Some(i as f32));
1427         }
1428         assert!(c.prev().is_none());
1429 
1430         assert_eq!(c.key(), Some(100));
1431         assert_eq!(c.value(), Some(100.0));
1432         for i in 101..200 {
1433             let (k, v) = c.next().unwrap();
1434             assert_eq!(k, i);
1435             assert_eq!(v, i as f32);
1436             assert_eq!(c.key(), Some(i));
1437             assert_eq!(c.value(), Some(i as f32));
1438         }
1439         assert!(c.next().is_none());
1440         assert!(c.key().is_none());
1441         assert!(c.value().is_none());
1442     }
1443 
1444     #[test]
into_iter()1445     fn into_iter() {
1446         let mut f = MapForest::<u32, f32>::new();
1447         let mut m = Map::<u32, f32>::new();
1448 
1449         for i in 10..20 {
1450             m.insert(i, i as f32, &mut f, &());
1451         }
1452 
1453         assert_eq!(
1454             m.into_iter(f).collect::<Vec<_>>(),
1455             vec![
1456                 (10, 10.0),
1457                 (11, 11.0),
1458                 (12, 12.0),
1459                 (13, 13.0),
1460                 (14, 14.0),
1461                 (15, 15.0),
1462                 (16, 16.0),
1463                 (17, 17.0),
1464                 (18, 18.0),
1465                 (19, 19.0),
1466             ],
1467         );
1468     }
1469 
1470     #[test]
range()1471     fn range() {
1472         let mut f = MapForest::<u32, f32>::new();
1473         let mut m = Map::<u32, f32>::new();
1474 
1475         for i in 10..20 {
1476             m.insert(i, i as f32, &mut f, &());
1477         }
1478 
1479         assert_eq!(
1480             m.range(.., &f, &()).collect::<Vec<_>>(),
1481             vec![
1482                 (10, 10.0),
1483                 (11, 11.0),
1484                 (12, 12.0),
1485                 (13, 13.0),
1486                 (14, 14.0),
1487                 (15, 15.0),
1488                 (16, 16.0),
1489                 (17, 17.0),
1490                 (18, 18.0),
1491                 (19, 19.0),
1492             ],
1493         );
1494 
1495         assert_eq!(
1496             m.range(5..12, &f, &()).collect::<Vec<_>>(),
1497             vec![(10, 10.0), (11, 11.0)],
1498         );
1499 
1500         assert_eq!(
1501             m.range(18..30, &f, &()).collect::<Vec<_>>(),
1502             vec![(18, 18.0), (19, 19.0)],
1503         );
1504 
1505         assert_eq!(
1506             m.range(..13, &f, &()).collect::<Vec<_>>(),
1507             vec![(10, 10.0), (11, 11.0), (12, 12.0)],
1508         );
1509 
1510         assert_eq!(
1511             m.range(18.., &f, &()).collect::<Vec<_>>(),
1512             vec![(18, 18.0), (19, 19.0)],
1513         );
1514 
1515         assert_eq!(
1516             m.range(12..=15, &f, &()).collect::<Vec<_>>(),
1517             vec![(12, 12.0), (13, 13.0), (14, 14.0), (15, 15.0)],
1518         );
1519 
1520         // Check when the query range is outside the entry range.
1521         assert_eq!(m.range(0..5, &f, &()).collect::<Vec<_>>(), vec![]);
1522         assert_eq!(m.range(30..40, &f, &()).collect::<Vec<_>>(), vec![]);
1523 
1524         // Check when the query range's start and end land in between entries.
1525         for i in 30..40 {
1526             if i % 2 == 0 {
1527                 m.insert(i, i as f32, &mut f, &());
1528             }
1529         }
1530         assert_eq!(
1531             m.range(31..35, &f, &()).collect::<Vec<_>>(),
1532             vec![(32, 32.0), (34, 34.0)],
1533         );
1534         assert_eq!(
1535             m.range(29..33, &f, &()).collect::<Vec<_>>(),
1536             vec![(30, 30.0), (32, 32.0)],
1537         );
1538         assert_eq!(
1539             m.range(35..40, &f, &()).collect::<Vec<_>>(),
1540             vec![(36, 36.0), (38, 38.0)],
1541         );
1542     }
1543 
1544     #[test]
range_next_back()1545     fn range_next_back() {
1546         let mut f = MapForest::<u32, f32>::new();
1547         let mut m = Map::<u32, f32>::new();
1548 
1549         for i in 10..20 {
1550             m.insert(i, i as f32, &mut f, &());
1551         }
1552 
1553         assert_eq!(
1554             m.range(12..16, &f, &()).rev().collect::<Vec<_>>(),
1555             vec![(15, 15.0), (14, 14.0), (13, 13.0), (12, 12.0)],
1556         );
1557 
1558         assert_eq!(
1559             m.range(18.., &f, &()).rev().collect::<Vec<_>>(),
1560             vec![(19, 19.0), (18, 18.0)],
1561         );
1562 
1563         assert_eq!(
1564             m.range(..12, &f, &()).rev().collect::<Vec<_>>(),
1565             vec![(11, 11.0), (10, 10.0)],
1566         );
1567 
1568         assert_eq!(
1569             m.range(11..=12, &f, &()).rev().collect::<Vec<_>>(),
1570             vec![(12, 12.0), (11, 11.0)],
1571         );
1572 
1573         let mut iter = m.range(13..=16, &f, &());
1574         assert_eq!(iter.next(), Some((13, 13.0)));
1575         assert_eq!(iter.next_back(), Some((16, 16.0)));
1576         assert_eq!(iter.next(), Some((14, 14.0)));
1577         assert_eq!(iter.next_back(), Some((15, 15.0)));
1578         assert!(iter.next().is_none());
1579         assert!(iter.next_back().is_none());
1580 
1581         let mut iter = m.range(13..=16, &f, &());
1582         assert_eq!(iter.next_back(), Some((16, 16.0)));
1583         assert_eq!(iter.next(), Some((13, 13.0)));
1584         assert_eq!(iter.next_back(), Some((15, 15.0)));
1585         assert_eq!(iter.next(), Some((14, 14.0)));
1586         assert!(iter.next_back().is_none());
1587         assert!(iter.next().is_none());
1588     }
1589 }
1590