range_map/
lib.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use arrayvec::ArrayVec;
6use std::cmp::{Eq, PartialEq};
7use std::fmt;
8use std::fmt::{Debug, Formatter};
9use std::ops::{Bound, Range, RangeBounds};
10use std::sync::Arc;
11
12/// The `B` constant for the btree.
13///
14/// Controls the size of the nodes inside the tree.
15const B: usize = 6;
16
17/// The capacity of nodes in the btree.
18const NODE_CAPACITY: usize = 2 * B;
19
20/// A trait for types that can calculate the gap between two values.
21pub trait Gap {
22    /// Measure the gap between two values.
23    fn measure_gap(&self, other: &Self) -> u64;
24}
25
26impl Gap for u32 {
27    fn measure_gap(&self, other: &Self) -> u64 {
28        if *self > *other { (*self - *other) as u64 } else { (*other - *self) as u64 }
29    }
30}
31
32/// A location inside the btree.
33#[derive(Debug, Default, Clone, Copy)]
34struct Cursor {
35    /// The number of valid indices in the `indices` array.
36    depth: u8,
37
38    /// The indices of the entry, ordered from leaf to root.
39    indices: [u8; 7],
40}
41
42impl Cursor {
43    /// Create a cursor with a single index.
44    fn with_index(index: usize) -> Self {
45        let mut cursor = Self::default();
46        cursor.push(index);
47        cursor
48    }
49
50    /// Whether the cursor is empty.
51    ///
52    /// A cursor is empty if it contains no more indices. This happens when a traversal has reached
53    /// a leaf node.
54    fn is_empty(&self) -> bool {
55        self.depth == 0
56    }
57
58    /// Push an index onto the front of the cursor.
59    ///
60    /// The front of the cursor is towards the root of the tree.
61    fn push(&mut self, index: usize) {
62        self.indices[self.depth as usize] = index as u8;
63        self.depth += 1;
64    }
65
66    /// Push an index onto the back of the cursor.
67    ///
68    /// The back of the cursor is towards the leaves of the tree.
69    fn push_back(&mut self, index: usize) {
70        self.indices.rotate_right(1);
71        self.indices[0] = index as u8;
72        self.depth += 1;
73    }
74
75    /// Pop an index off the front of the cursor.
76    ///
77    /// The front of the cursor is towards the root of the tree.
78    fn pop(&mut self) -> Option<usize> {
79        if self.depth == 0 {
80            None
81        } else {
82            self.depth -= 1;
83            Some(self.indices[self.depth as usize] as usize)
84        }
85    }
86
87    /// Pop an index off the back of the cursor.
88    ///
89    /// The back of the cursor is towards the leaves of the tree.
90    fn pop_back(&mut self) -> Option<usize> {
91        if self.depth == 0 {
92            None
93        } else {
94            self.depth -= 1;
95            let index = self.indices[0] as usize;
96            self.indices.rotate_left(1);
97            Some(index)
98        }
99    }
100
101    /// The backmost index in the cursor.
102    ///
103    /// The back of the cursor is towards the leaves of the tree.
104    ///
105    /// Assumes the cursor is non-empty.
106    fn back(&self) -> usize {
107        self.indices[0] as usize
108    }
109
110    /// Increment the backmost index in the cursor.
111    ///
112    /// The back of the cursor is towards the leaves of the tree.
113    ///
114    /// Assumes the cursor is non-empty.
115    fn increment_back(&mut self) {
116        self.indices[0] += 1;
117    }
118
119    /// Decrement the backmost index in the cursor.
120    ///
121    /// The back of the cursor is towards the leaves of the tree.
122    ///
123    /// Assumes the cursor is non-empty.
124    fn decrement_back(&mut self) {
125        self.indices[0] -= 1;
126    }
127}
128
129impl PartialEq for Cursor {
130    fn eq(&self, other: &Self) -> bool {
131        if self.depth != other.depth {
132            return false;
133        }
134        for i in 0..self.depth {
135            if self.indices[i as usize] != other.indices[i as usize] {
136                return false;
137            }
138        }
139        true
140    }
141}
142
143impl Eq for Cursor {}
144
145/// Where to place the cursor relative to the given key.
146enum CursorPosition {
147    /// The given key represents a left edge of a range.
148    ///
149    /// Place the cursor to the left of a range containing the cursor.
150    Left,
151
152    /// The given key represents a right edge of a range.
153    ///
154    /// Place the cursor to the right of a range containing the cursor.
155    Right,
156}
157
158/// Search of the given key in the given array of ranges.
159///
160/// If the array contains a range that contains the key, returns the index of that range.
161/// Otherwise, returns the index at which the given key could be inserted into the array to
162/// maintain the ordering.
163fn binary_search<K: Ord>(key: &K, keys: &ArrayVec<Range<K>, NODE_CAPACITY>) -> usize {
164    let mut left = 0usize;
165    let mut right = keys.len();
166    while left < right {
167        let mid = left + (right - left) / 2;
168        // TODO: Consider `get_unchecked`.
169        let range = &keys[mid];
170        if key < &range.start {
171            // This range is too large.
172            right = mid;
173        } else if key < &range.end {
174            // We found the range that contains this key.
175            return mid;
176        } else {
177            // The key might be found in the next range.
178            left = mid + 1;
179        }
180    }
181    // The key falls between two ranges. Return the index at which this key could be inserted to
182    // maintain the ordering.
183    left
184}
185
186/// A leaf node in the btree.
187///
188/// Stores a flat map of keys to values, with the `i`th entry in the keys array corresponding to
189/// the `i`th entry in the values array. The balancing rules of the btree ensure that every
190/// non-root leaf has between N and N/2 entries populated.
191#[derive(Clone)]
192struct NodeLeaf<K: Ord + Copy + Gap, V: Clone> {
193    /// The maximum gap between keys in this leaf node.
194    max_gap: u64,
195
196    /// The keys stored in this leaf node.
197    ///
198    /// We store the key in a dense array to improve cache performance during lookups. We often
199    /// need to binary-search the keys in a given leaf node, which means having those keys close
200    /// together improves cache performance.
201    keys: ArrayVec<Range<K>, NODE_CAPACITY>,
202
203    /// The value stored in this leaf node.
204    values: ArrayVec<V, NODE_CAPACITY>,
205}
206
207/// Shows the map structure of the leaf node.
208impl<K, V> Debug for NodeLeaf<K, V>
209where
210    K: Debug + Ord + Copy + Gap,
211    V: Debug + Clone,
212{
213    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
214        f.debug_map().entries(self.keys.iter().zip(self.values.iter())).finish()
215    }
216}
217
218/// The result of performing an insertion into a btree node.
219enum InsertResult<K: Ord + Copy + Gap, V: Clone> {
220    /// The value was successfully inserted into an empty slot.
221    Inserted,
222
223    /// The value was inserted into an empty slot in a leaf node but that insertion caused the
224    /// leaf node to exceed its capacity and split into two leaf nodes. The existing leaf node
225    /// now holds the entries to the left of the split and the entries to the right of the split
226    /// are returned. The split occurred at the returned key.
227    SplitLeaf(K, Arc<NodeLeaf<K, V>>),
228
229    /// The value was inserted into an empty slot in a subtree but that insertion caused the
230    /// internal node to exceed its capacity and split into two internal nodes. The internal node
231    /// now holds the entries to the left of the split and the entries to the right of the split
232    /// are returned. The split occurred at the returned key.
233    SplitInternal(K, Arc<NodeInternal<K, V>>),
234}
235
236/// State information returned from the removal operation.
237struct RemoveState<K: Ord + Copy + Gap, V: Clone> {
238    /// The minimal key for this subtree, if available.
239    ///
240    /// If the minimal key for this subtree is not available, then the minimal key is guaranteed
241    /// to be unchanged by this operation.
242    maybe_split_key: Option<K>,
243
244    /// The value previously stored at this key.
245    removed_value: V,
246}
247
248/// The intermediate result of a remove operation.
249///
250/// The root of the CowTree converts this result value into `Option<T>`, per the usual map
251/// interface.
252enum RemoveResult<K: Ord + Copy + Gap, V: Clone> {
253    /// The key the client asked to remove was not found in the map.
254    NotFound,
255
256    /// The key was successfully removed from the map.
257    ///
258    Removed(RemoveState<K, V>),
259
260    /// The key was removed from the map but the node that previously contained that node no longer
261    /// has sufficient children.
262    ///
263    /// The caller is responsible for rebalancing its children to ensure that each node has at
264    /// least this minimum number of children. If the balance invariant can be resolved locally,
265    /// the caller should return `Removed` to its caller. If rebalancing the local children
266    /// causes this node to have fewer than the minimum number of children, the caller should
267    /// return `Underflow` to its caller.
268    Underflow(RemoveState<K, V>),
269}
270
271impl<K, V> NodeLeaf<K, V>
272where
273    K: Ord + Copy + Gap,
274    V: Clone,
275{
276    /// Create a new leaf node.
277    fn new(keys: ArrayVec<Range<K>, NODE_CAPACITY>, values: ArrayVec<V, NODE_CAPACITY>) -> Self {
278        let mut node = Self { max_gap: 0, keys, values };
279        node.update_max_gap();
280        node
281    }
282
283    /// Create an empty leaf node.
284    ///
285    /// Empty leaf nodes are used only at the root of the tree.
286    fn empty() -> Self {
287        Self { max_gap: 0, keys: ArrayVec::new(), values: ArrayVec::new() }
288    }
289
290    /// Gets the index in this leaf that corresponds to the given cursor.
291    ///
292    /// Assumes the cursor contains exactly one index.
293    ///
294    /// Returns `None` if the cursor points beyond the end if this node.
295    fn get_index(&self, mut cursor: Cursor) -> Option<usize> {
296        let index = cursor.pop().expect("Cursor has sufficient depth");
297        assert!(cursor.is_empty(), "Cursor has excess depth");
298        if index >= self.keys.len() {
299            return None;
300        }
301        Some(index)
302    }
303
304    /// Search this leaf for the given key and return both the key and the value found.
305    fn get_key_value(&self, cursor: Cursor) -> Option<(&Range<K>, &V)> {
306        if let Some(index) = self.get_index(cursor) {
307            let key = &self.keys[index];
308            let value = &self.values[index];
309            Some((key, value))
310        } else {
311            None
312        }
313    }
314
315    /// The last key/value pair stored in this leaf.
316    fn last_key_value(&self) -> Option<(&Range<K>, &V)> {
317        let key = self.keys.last()?;
318        let value = self.values.last()?;
319        Some((key, value))
320    }
321
322    /// Find the given key in this node.
323    ///
324    /// Updates `cursor` to point to the position indicated by `position`.
325    fn find(&self, key: &K, position: CursorPosition, cursor: &mut Cursor) {
326        let index = binary_search(key, &self.keys);
327        match position {
328            CursorPosition::Left => {
329                cursor.push(index);
330            }
331            CursorPosition::Right => {
332                if let Some(range) = self.keys.get(index) {
333                    if *key > range.start {
334                        cursor.push(index + 1);
335                        return;
336                    }
337                }
338                cursor.push(index);
339            }
340        }
341    }
342
343    /// Compute the maximum gap for this node.
344    fn compute_max_gap(&self) -> u64 {
345        let mut max_gap = 0;
346        for i in 0..self.keys.len() {
347            if i + 1 < self.keys.len() {
348                max_gap = max_gap.max(self.keys[i].end.measure_gap(&self.keys[i + 1].start));
349            }
350        }
351        max_gap
352    }
353
354    /// Validates that the cached max_gap value matches what we would compute now.
355    #[cfg(test)]
356    fn validate_max_gap(&self) {
357        let computed = self.compute_max_gap();
358        assert_eq!(computed, self.max_gap);
359    }
360
361    /// Update the maximum gap for this node.
362    fn update_max_gap(&mut self) {
363        self.max_gap = self.compute_max_gap();
364    }
365
366    /// Measure the gap between the last key in this node and the first key in the other node.
367    fn measure_gap(&self, other: &Self) -> u64 {
368        // We know that `self.keys` is not empty because this function is only called when there is
369        // more than one leaf node. The only situation in which `self.keys` is empty is when the
370        // tree is empty, in which case there is only a single empty leaf node.
371        self.keys[self.keys.len() - 1].end.measure_gap(&other.keys[0].start)
372    }
373
374    /// Get the range of keys in this node.
375    fn key_range(&self) -> Range<K> {
376        self.keys[0].start..self.keys[self.keys.len() - 1].end
377    }
378
379    /// Insert the given entry at the location indicated by `cursor`.
380    ///
381    /// Inserting a value into a leaf node might cause this node to split into two leaf nodes.
382    fn insert(&mut self, mut cursor: Cursor, range: Range<K>, value: V) -> InsertResult<K, V> {
383        let index = cursor.pop().expect("valid cursor");
384        let result = if self.keys.len() == NODE_CAPACITY {
385            if index == NODE_CAPACITY {
386                let mut keys = ArrayVec::new();
387                let mut values = ArrayVec::new();
388                let key = range.start;
389                keys.push(range);
390                values.push(value);
391                return InsertResult::SplitLeaf(key, Arc::new(Self::new(keys, values)));
392            }
393            let middle = NODE_CAPACITY / 2;
394            assert!(middle > 0);
395            let mut right = Self::new(
396                self.keys.drain(middle..).collect(),
397                self.values.drain(middle..).collect(),
398            );
399            if index <= middle {
400                self.keys.insert(index, range);
401                self.values.insert(index, value);
402            } else {
403                right.keys.insert(index - middle, range);
404                right.values.insert(index - middle, value);
405                right.update_max_gap();
406            }
407            InsertResult::SplitLeaf(right.keys[0].start, Arc::new(right))
408        } else {
409            self.keys.insert(index, range);
410            self.values.insert(index, value);
411            InsertResult::Inserted
412        };
413        self.update_max_gap();
414        result
415    }
416
417    /// Remove the entry indicated by `cursor`.
418    fn remove(&mut self, cursor: Cursor) -> RemoveResult<K, V> {
419        if let Some(index) = self.get_index(cursor) {
420            self.keys.remove(index);
421            let removed_value = self.values.remove(index);
422            let maybe_split_key = self.keys.first().map(|range| range.start);
423            self.update_max_gap();
424            if self.keys.len() < NODE_CAPACITY / 2 {
425                RemoveResult::Underflow(RemoveState { maybe_split_key, removed_value })
426            } else {
427                RemoveResult::Removed(RemoveState { maybe_split_key, removed_value })
428            }
429        } else {
430            RemoveResult::NotFound
431        }
432    }
433
434    /// Find a gap that is at least as large as the given gap and is less than the given upper bound.
435    ///
436    /// Returns the end of the gap, which might above or below the node.
437    fn find_gap_end(&self, gap: u64, upper_bound: &K) -> K {
438        if self.keys.is_empty() {
439            return *upper_bound;
440        }
441
442        let node_end = &self.keys[self.keys.len() - 1].end;
443        if node_end <= upper_bound && node_end.measure_gap(upper_bound) >= gap {
444            return *upper_bound;
445        }
446
447        if self.max_gap >= gap {
448            for (i, _key) in self.keys.iter().enumerate().rev() {
449                if i > 0 {
450                    let start = &self.keys[i - 1].end;
451                    if start >= upper_bound {
452                        continue;
453                    }
454                    let end = upper_bound.min(&self.keys[i].start);
455                    if start.measure_gap(end) >= gap {
456                        return *end;
457                    }
458                }
459            }
460        }
461
462        let node_start = &self.keys[0].start;
463        *upper_bound.min(node_start)
464    }
465}
466
467/// The children of an internal node in the btree.
468#[derive(Clone, Debug)]
469enum ChildList<K: Ord + Copy + Gap, V: Clone> {
470    /// Used when an internal node has leaf nodes as children.
471    Leaf(ArrayVec<Arc<NodeLeaf<K, V>>, NODE_CAPACITY>),
472
473    /// Used when an internal node has other internal nodes as children.
474    Internal(ArrayVec<Arc<NodeInternal<K, V>>, NODE_CAPACITY>),
475}
476
477impl<K, V> ChildList<K, V>
478where
479    K: Ord + Copy + Gap,
480    V: Clone,
481{
482    /// Create a child list that has no children.
483    fn new_empty(&self) -> Self {
484        match self {
485            ChildList::Leaf(_) => ChildList::Leaf(ArrayVec::new()),
486            ChildList::Internal(_) => ChildList::Internal(ArrayVec::new()),
487        }
488    }
489
490    /// The number of children for this node.
491    fn len(&self) -> usize {
492        match self {
493            ChildList::Leaf(children) => children.len(),
494            ChildList::Internal(children) => children.len(),
495        }
496    }
497
498    /// The number of gradchildren located at the child with the given index.
499    fn size_at(&self, i: usize) -> usize {
500        match self {
501            ChildList::Leaf(children) => children[i].keys.len(),
502            ChildList::Internal(children) => children[i].children.len(),
503        }
504    }
505
506    /// Obtain the child located at the given index.
507    fn get(&self, i: usize) -> Node<K, V> {
508        match self {
509            ChildList::Leaf(children) => Node::Leaf(children[i].clone()),
510            ChildList::Internal(children) => Node::Internal(children[i].clone()),
511        }
512    }
513
514    /// Get a reference to the child located at the given index.
515    fn get_ref(&self, i: usize) -> NodeRef<'_, K, V> {
516        match self {
517            ChildList::Leaf(children) => NodeRef::Leaf(&children[i]),
518            ChildList::Internal(children) => NodeRef::Internal(&children[i]),
519        }
520    }
521
522    /// Get the range of keys in the subtree rooted at this node.
523    fn subtree_key_range(&self) -> Range<K> {
524        match self {
525            ChildList::Leaf(children) => {
526                children[0].key_range().start..children[children.len() - 1].key_range().end
527            }
528            ChildList::Internal(children) => {
529                children[0].subtree_key_range.start
530                    ..children[children.len() - 1].subtree_key_range.end
531            }
532        }
533    }
534
535    /// Removes all the entries starting at the given index from the child list.
536    ///
537    /// The removed entries are returned in a new child list.
538    fn split_off(&mut self, index: usize) -> Self {
539        match self {
540            ChildList::Leaf(children) => ChildList::Leaf(children.drain(index..).collect()),
541            ChildList::Internal(children) => ChildList::Internal(children.drain(index..).collect()),
542        }
543    }
544
545    /// Removes all the entries prior to the given index from the child list.
546    ///
547    /// The removed entries are returned in a new child list.
548    fn split_off_front(&mut self, index: usize) -> Self {
549        match self {
550            ChildList::Leaf(children) => ChildList::Leaf(children.drain(..index).collect()),
551            ChildList::Internal(children) => ChildList::Internal(children.drain(..index).collect()),
552        }
553    }
554
555    /// Insert a child into the child list.
556    ///
557    /// The type of child node must match the type of the child list.
558    fn insert(&mut self, index: usize, child: Node<K, V>) {
559        match self {
560            ChildList::Leaf(children) => {
561                let Node::Leaf(leaf) = child else {
562                    unreachable!("Inserting a non-leaf into an internal node for leaf nodes.");
563                };
564                children.insert(index, leaf);
565            }
566            ChildList::Internal(children) => {
567                let Node::Internal(internal) = child else {
568                    unreachable!(
569                        "Inserting a non-internal into an internal node for internal nodes."
570                    );
571                };
572                children.insert(index, internal);
573            }
574        }
575    }
576
577    /// Remove the child at the given index from the child list.
578    fn remove(&mut self, index: usize) {
579        match self {
580            ChildList::Leaf(children) => {
581                children.remove(index);
582            }
583            ChildList::Internal(children) => {
584                children.remove(index);
585            }
586        }
587    }
588
589    /// Add the children from the given `ChildList` to this child list.
590    fn extend(&mut self, other: &Self) {
591        match (self, other) {
592            (ChildList::Leaf(self_children), ChildList::Leaf(other_children)) => {
593                self_children.extend(other_children.iter().cloned());
594            }
595            (ChildList::Internal(self_children), ChildList::Internal(other_children)) => {
596                self_children.extend(other_children.iter().cloned());
597            }
598            _ => unreachable!("Type mismatch while extending a child list."),
599        }
600    }
601}
602
603/// An internal node in the btree.
604#[derive(Clone, Debug)]
605struct NodeInternal<K: Ord + Copy + Gap, V: Clone> {
606    /// The maximum gap between keys in this internal node.
607    max_gap: u64,
608
609    /// The range of keys stored in the subtree rooted at this node.
610    subtree_key_range: Range<K>,
611
612    /// A cache of the keys that partition the keys in the children.
613    /// The key at index `i` is the smallest key stored in the subtree
614    /// of the `i`+1 child.
615    ///
616    /// We only ever store CAPACITY - 1 keys in this array.
617    keys: ArrayVec<K, NODE_CAPACITY>,
618
619    /// The children of this node.
620    children: ChildList<K, V>,
621}
622
623/// Get mutable references to two entries in a slice.
624///
625/// When rebalancing nodes, we need to mutate two nodes at the same time. Normally, if you take a
626/// mutable reference to one element of an array, the borrow checker prevents you from taking a
627/// mutable reference to a second element of the same array.
628///
629/// The nightly version of `std::primitive::slice` has `get_many_mut` to let you get mutable
630/// references to multiple elements. However, without that interface, the recommended approach for
631/// avoiding `unsafe` is to use `split_at_mut`.
632fn get_two_mut<T>(slice: &mut [T], i: usize, j: usize) -> (&mut T, &mut T) {
633    if i < j {
634        let (a, b) = slice.split_at_mut(j);
635        (&mut a[i], &mut b[0])
636    } else {
637        let (a, b) = slice.split_at_mut(i);
638        (&mut b[0], &mut a[j])
639    }
640}
641
642impl<K, V> NodeInternal<K, V>
643where
644    K: Ord + Copy + Gap,
645    V: Clone,
646{
647    /// Create a new internal node.
648    fn new(keys: ArrayVec<K, NODE_CAPACITY>, children: ChildList<K, V>) -> Self {
649        let mut node =
650            Self { max_gap: 0, subtree_key_range: children.subtree_key_range(), keys, children };
651        node.update_max_gap();
652        node
653    }
654
655    /// The index of the child that might contain the given key.
656    ///
657    /// Searches the cached keys at this node to determine which child node might contain the given
658    /// key.
659    fn get_child_index(&self, key: &K) -> usize {
660        let p = self.keys.partition_point(|k| k < key);
661        if self.keys.get(p) == Some(key) {
662            // If the query key exactly matches the split key, then we need to look for this key to
663            // the right of the split.
664            p + 1
665        } else {
666            // Otherwise, we look to the left of the split.
667            p
668        }
669    }
670
671    /// Search this subtree for the given key and return both the key and the value found.
672    fn get_key_value(&self, mut cursor: Cursor) -> Option<(&Range<K>, &V)> {
673        let index = cursor.pop().expect("valid cursor");
674        match &self.children {
675            ChildList::Leaf(children) => children[index].get_key_value(cursor),
676            ChildList::Internal(children) => children[index].get_key_value(cursor),
677        }
678    }
679
680    /// Returns a reference to the node that contains the entry indicated by the cursor.
681    ///
682    /// Assumes the cursor points a descendant of this node.
683    fn get_containing_node(&self, mut cursor: Cursor) -> NodeRef<'_, K, V> {
684        debug_assert!(cursor.depth >= 2);
685        let index = cursor.pop().expect("valid cursor");
686        if cursor.depth == 1 {
687            return self.children.get_ref(index);
688        }
689        match &self.children {
690            ChildList::Leaf(_) => unreachable!("leaf nodes do not have children"),
691            ChildList::Internal(children) => children[index].get_containing_node(cursor),
692        }
693    }
694
695    /// The last key/value pair stored in this subtree.
696    fn last_key_value(&self) -> Option<(&Range<K>, &V)> {
697        match &self.children {
698            ChildList::Leaf(children) => {
699                children.last().expect("child lists are always non-empty").last_key_value()
700            }
701            ChildList::Internal(children) => {
702                children.last().expect("child lists are always non-empty").last_key_value()
703            }
704        }
705    }
706
707    /// Find the given key in this node.
708    ///
709    /// Updates `cursor` to point to the position indicated by `position`.
710    fn find(&self, key: &K, position: CursorPosition, cursor: &mut Cursor) {
711        let index = self.get_child_index(&key);
712        match &self.children {
713            ChildList::Leaf(children) => children[index].find(key, position, cursor),
714            ChildList::Internal(children) => children[index].find(key, position, cursor),
715        }
716        cursor.push(index);
717    }
718
719    /// Compute the maximum gap for this node.
720    fn compute_max_gap(&self) -> u64 {
721        let mut max_gap = 0;
722        match &self.children {
723            ChildList::Leaf(children) => {
724                for i in 0..children.len() {
725                    max_gap = max_gap.max(children[i].max_gap);
726                    if i < children.len() - 1 {
727                        max_gap = max_gap.max(children[i].measure_gap(&children[i + 1]));
728                    }
729                }
730            }
731            ChildList::Internal(children) => {
732                for i in 0..children.len() {
733                    max_gap = max_gap.max(children[i].max_gap);
734                    if i < children.len() - 1 {
735                        max_gap = max_gap.max(children[i].measure_gap(&children[i + 1]));
736                    }
737                }
738            }
739        }
740        max_gap
741    }
742
743    /// Validates that the cached max_gap value matches what we would compute now,
744    /// and recursively validates all children.
745    #[cfg(test)]
746    fn validate_max_gap(&self) {
747        let computed = self.compute_max_gap();
748        assert_eq!(computed, self.max_gap);
749
750        // Recursively validate children
751        match &self.children {
752            ChildList::Leaf(children) => {
753                for child in children {
754                    child.validate_max_gap();
755                }
756            }
757            ChildList::Internal(children) => {
758                for child in children {
759                    child.validate_max_gap();
760                }
761            }
762        }
763    }
764
765    /// Checks that the keys stored in this node are actually the leftmost edge of the
766    /// corresponding children. Specifically, key `i` should be the `start` of the range of the
767    /// zeroth entry in of child `i+1`.
768    ///
769    /// Panics if the node (or a descendant) does not validate.
770    #[cfg(test)]
771    fn validate_keys(&self) -> K {
772        let mut first_key = None;
773        match &self.children {
774            ChildList::Leaf(children) => {
775                for (i, child) in children.iter().enumerate() {
776                    let child_key = child.keys[0].start;
777                    if i == 0 {
778                        first_key = Some(child_key);
779                    } else {
780                        assert!(child_key == self.keys[i - 1]);
781                    }
782                }
783            }
784            ChildList::Internal(children) => {
785                for (i, child) in children.iter().enumerate() {
786                    let child_key = child.validate_keys();
787                    if i == 0 {
788                        first_key = Some(child_key);
789                    } else {
790                        assert!(child_key == self.keys[i - 1]);
791                    }
792                }
793            }
794        }
795
796        first_key.expect("internal nodes must be non-empty")
797    }
798
799    /// Update the maximum gap for this node.
800    fn update_max_gap(&mut self) {
801        self.subtree_key_range = self.children.subtree_key_range();
802        self.max_gap = self.compute_max_gap();
803    }
804
805    /// Measure the gap between the last key in this node and the first key in the other node.
806    fn measure_gap(&self, other: &Self) -> u64 {
807        self.subtree_key_range.end.measure_gap(&other.subtree_key_range.start)
808    }
809
810    /// Insert the given child node at `index` in this node.
811    ///
812    /// `key` must be the smallest key that occurs in the `child` subtree.
813    ///
814    /// The caller must ensure that the child is inserted in the correct location.
815    fn insert_child(&mut self, index: usize, key: K, child: Node<K, V>) -> InsertResult<K, V> {
816        let n = self.children.len();
817        if n == NODE_CAPACITY {
818            // TODO: This branch is incorrect. We need to check for NODE_CAPACITY - 1 because the
819            // index indicates the location of the existing child that split, which will never be
820            // beyond the end of the array.
821            if index == NODE_CAPACITY {
822                let mut children = self.children.new_empty();
823                children.insert(0, child);
824                let right = Self::new(ArrayVec::new(), children);
825                return InsertResult::SplitInternal(key, Arc::new(right));
826            }
827            let middle = NODE_CAPACITY / 2;
828            assert!(middle > 0);
829            let mut internal =
830                Self::new(self.keys.drain(middle..).collect(), self.children.split_off(middle));
831            let split_key = self.keys.pop().unwrap();
832            if index < middle {
833                self.keys.insert(index, key);
834                self.children.insert(index + 1, child);
835            } else {
836                internal.keys.insert(index - middle, key);
837                internal.children.insert(index - middle + 1, child);
838            }
839            debug_assert!(self.keys.len() + 1 == self.children.len());
840            debug_assert!(internal.keys.len() + 1 == internal.children.len());
841            internal.update_max_gap();
842            InsertResult::SplitInternal(split_key, Arc::new(internal))
843        } else {
844            self.keys.insert(index, key);
845            self.children.insert(index + 1, child);
846            debug_assert!(self.keys.len() + 1 == self.children.len());
847            InsertResult::Inserted
848        }
849    }
850
851    /// Insert the given entry at the location indicated by `cursor`.
852    ///
853    /// Inserting a value into an internal node might cause this node to split into two internal
854    /// nodes.
855    fn insert(&mut self, mut cursor: Cursor, range: Range<K>, value: V) -> InsertResult<K, V> {
856        let index = cursor.pop().expect("valid cursor");
857        let result = match &mut self.children {
858            ChildList::Leaf(children) => {
859                Arc::make_mut(&mut children[index]).insert(cursor, range, value)
860            }
861            ChildList::Internal(children) => {
862                Arc::make_mut(&mut children[index]).insert(cursor, range, value)
863            }
864        };
865        let result = match result {
866            InsertResult::Inserted => InsertResult::Inserted,
867            InsertResult::SplitLeaf(key, right) => self.insert_child(index, key, Node::Leaf(right)),
868            InsertResult::SplitInternal(key, right) => {
869                self.insert_child(index, key, Node::Internal(right))
870            }
871        };
872        self.update_max_gap();
873        result
874    }
875
876    /// Determine whether to rebalance the child with the given index to the left or to the right.
877    ///
878    /// Given a choice, we will rebalance the child with its larger neighbor.
879    ///
880    /// The indices returned are always sequential.
881    fn select_children_to_rebalance(&self, index: usize) -> (usize, usize) {
882        if index == 0 {
883            (index, index + 1)
884        } else if index == self.children.len() - 1 {
885            (index - 1, index)
886        } else {
887            let left_index = index - 1;
888            let left_size = self.children.size_at(left_index);
889            let right_index = index + 1;
890            let right_size = self.children.size_at(right_index);
891            if left_size > right_size { (left_index, index) } else { (index, right_index) }
892        }
893    }
894    /// Rebalance the child at the given index.
895    ///
896    /// If the child and its neighbor are sufficiently small, this function will merge them into a
897    /// single node.
898    fn rebalance_child(&mut self, index: usize) {
899        // Cannot rebalance if we have fewer than two children. This situation occurs only at the
900        // root of the tree.
901        if self.children.len() < 2 {
902            return;
903        }
904        let (left, right) = self.select_children_to_rebalance(index);
905        let left_size = self.children.size_at(left);
906        debug_assert!(left_size > 0);
907        let right_size = self.children.size_at(right);
908        debug_assert!(right_size > 0);
909        let n = left_size + right_size;
910        match &mut self.children {
911            ChildList::Leaf(children) => {
912                let (left_shared_node, right_shared_node) = get_two_mut(children, left, right);
913                let left_node = Arc::make_mut(left_shared_node);
914                if n <= NODE_CAPACITY {
915                    // Merge the right node into the left node.
916                    left_node.keys.extend(right_shared_node.keys.iter().cloned());
917                    left_node.values.extend(right_shared_node.values.iter().cloned());
918                    left_node.update_max_gap();
919                    self.keys.remove(left);
920                    self.children.remove(right);
921                } else {
922                    // Rebalance the elements between the nodes.
923                    let split = n / 2;
924                    let right_node = Arc::make_mut(right_shared_node);
925                    if left_node.values.len() < split {
926                        // Move elements from right to left.
927                        let move_count = split - left_node.values.len();
928                        left_node.keys.extend(right_node.keys.drain(..move_count));
929                        left_node.values.extend(right_node.values.drain(..move_count));
930                    } else {
931                        // Move elements from left to right.
932                        let mut keys = ArrayVec::new();
933                        keys.extend(left_node.keys.drain(split..));
934                        keys.extend(right_node.keys.iter().cloned());
935                        right_node.keys = keys;
936
937                        let mut values = ArrayVec::new();
938                        values.extend(left_node.values.drain(split..));
939                        values.extend(right_node.values.iter().cloned());
940                        right_node.values = values;
941                    }
942                    left_node.update_max_gap();
943                    right_node.update_max_gap();
944                    // Update the split key to reflect the new division between the nodes.
945                    self.keys[left] = right_node.keys[0].start;
946                }
947            }
948            ChildList::Internal(children) => {
949                let (left_shard_node, right_shared_node) = get_two_mut(children, left, right);
950                let left_node = Arc::make_mut(left_shard_node);
951                let old_split_key = &self.keys[left];
952                if n <= NODE_CAPACITY {
953                    // Merge the right node into the left node.
954                    left_node.keys.push(old_split_key.clone());
955                    left_node.keys.extend(right_shared_node.keys.iter().cloned());
956                    left_node.children.extend(&right_shared_node.children);
957                    left_node.update_max_gap();
958                    debug_assert!(left_node.keys.len() + 1 == left_node.children.len());
959                    self.keys.remove(left);
960                    self.children.remove(right);
961                } else {
962                    // Rebalance the elements between the nodes.
963                    let split = n / 2;
964                    let split_key;
965                    let right_node = Arc::make_mut(right_shared_node);
966                    if left_node.children.len() < split {
967                        // Move elements from right to left.
968                        let move_count = split - left_node.children.len();
969                        left_node.keys.push(old_split_key.clone());
970                        left_node.keys.extend(right_node.keys.drain(..move_count));
971                        split_key =
972                            left_node.keys.pop().expect("must have moved at least one element");
973
974                        left_node.children.extend(&right_node.children.split_off_front(move_count));
975                        debug_assert!(left_node.keys.len() + 1 == left_node.children.len());
976                    } else {
977                        // Move elements from left to right.
978                        let mut it = left_node.keys.drain((split - 1)..);
979                        split_key = it.next().expect("must be moving at least one element");
980                        let mut keys = ArrayVec::new();
981                        keys.extend(it);
982                        keys.push(old_split_key.clone());
983                        keys.extend(right_node.keys.iter().cloned());
984                        right_node.keys = keys;
985
986                        let mut children = right_node.children.new_empty();
987                        children.extend(&left_node.children.split_off(split));
988                        children.extend(&right_node.children);
989                        right_node.children = children;
990                        debug_assert!(left_node.keys.len() + 1 == left_node.children.len());
991                        debug_assert!(right_node.keys.len() + 1 == right_node.children.len());
992                    }
993                    left_node.update_max_gap();
994                    right_node.update_max_gap();
995                    // Update the split key to reflect the new division between the nodes.
996                    self.keys[left] = split_key;
997                }
998            }
999        }
1000    }
1001
1002    /// Remove the entry indicated by `cursor`.
1003    fn remove(&mut self, mut cursor: Cursor) -> RemoveResult<K, V> {
1004        let index = cursor.pop().expect("valid cursor");
1005        let result = match &mut self.children {
1006            ChildList::Leaf(children) => Arc::make_mut(&mut children[index]).remove(cursor),
1007            ChildList::Internal(children) => Arc::make_mut(&mut children[index]).remove(cursor),
1008        };
1009        match result {
1010            RemoveResult::NotFound => RemoveResult::NotFound,
1011            RemoveResult::Removed(RemoveState { mut maybe_split_key, removed_value }) => {
1012                if let Some(key) = maybe_split_key {
1013                    if index > 0 {
1014                        self.keys[index - 1] = key;
1015                        maybe_split_key = None;
1016                    }
1017                }
1018                self.update_max_gap();
1019                RemoveResult::Removed(RemoveState { maybe_split_key, removed_value })
1020            }
1021            RemoveResult::Underflow(RemoveState { mut maybe_split_key, removed_value }) => {
1022                if let Some(key) = maybe_split_key {
1023                    if index > 0 {
1024                        self.keys[index - 1] = key;
1025                        maybe_split_key = None;
1026                    }
1027                }
1028                // If the child underflowed to zero we can simply remove the child rather than
1029                // rebalancing existing nodes.
1030                if self.children.size_at(index) == 0 {
1031                    // If the child underflowed to zero elements, the child cannot have a split key
1032                    // because the child has no keys.
1033                    debug_assert!(maybe_split_key.is_none());
1034                    self.children.remove(index);
1035                    if index == 0 {
1036                        // When we remove the first child, we need to return the split key for the
1037                        // new first child to our parent.
1038                        maybe_split_key = Some(self.keys.remove(0));
1039                    } else {
1040                        // Otherwise, we can just remove the split key associated with the removed
1041                        // child.
1042                        self.keys.remove(index - 1);
1043                    }
1044                } else {
1045                    self.rebalance_child(index);
1046                }
1047                self.update_max_gap();
1048                if self.children.len() < NODE_CAPACITY / 2 {
1049                    RemoveResult::Underflow(RemoveState { maybe_split_key, removed_value })
1050                } else {
1051                    RemoveResult::Removed(RemoveState { maybe_split_key, removed_value })
1052                }
1053            }
1054        }
1055    }
1056
1057    /// Find a gap that is at least as large as the given gap and is less than the given upper bound.
1058    ///
1059    /// Returns the end of the gap, which might above or below the node.
1060    fn find_gap_end(&self, gap: u64, upper_bound: &K) -> K {
1061        let node_start = &self.subtree_key_range.start;
1062        if node_start > upper_bound {
1063            return *upper_bound;
1064        }
1065
1066        let node_end = &self.subtree_key_range.end;
1067        if node_end <= upper_bound && node_end.measure_gap(upper_bound) >= gap {
1068            return *upper_bound;
1069        }
1070
1071        if self.max_gap >= gap {
1072            match &self.children {
1073                ChildList::Leaf(children) => {
1074                    let mut child_upper_bound = *upper_bound;
1075                    for (i, child) in children.iter().enumerate().rev() {
1076                        if child.key_range().start >= *upper_bound {
1077                            continue;
1078                        }
1079                        let end = child.find_gap_end(gap, &child_upper_bound);
1080                        if i > 0 {
1081                            let start = children[i - 1].key_range().end;
1082                            if start.measure_gap(&end) < gap {
1083                                child_upper_bound = start;
1084                                continue;
1085                            }
1086                        }
1087                        return end;
1088                    }
1089                }
1090                ChildList::Internal(children) => {
1091                    let mut child_upper_bound = *upper_bound;
1092                    for (i, child) in children.iter().enumerate().rev() {
1093                        if child.subtree_key_range.start >= *upper_bound {
1094                            continue;
1095                        }
1096                        let end = child.find_gap_end(gap, &child_upper_bound);
1097                        if i > 0 {
1098                            let start = children[i - 1].subtree_key_range.end;
1099                            if start.measure_gap(&end) < gap {
1100                                child_upper_bound = start;
1101                                continue;
1102                            }
1103                        }
1104                        return end;
1105                    }
1106                }
1107            }
1108        }
1109
1110        *node_start
1111    }
1112}
1113
1114/// A node in the btree.
1115#[derive(Clone, Debug)]
1116enum Node<K: Ord + Copy + Gap, V: Clone> {
1117    /// An internal node.
1118    Internal(Arc<NodeInternal<K, V>>),
1119
1120    /// A leaf node.
1121    Leaf(Arc<NodeLeaf<K, V>>),
1122}
1123
1124impl<K, V> Node<K, V>
1125where
1126    K: Ord + Copy + Gap,
1127    V: Clone,
1128{
1129    /// The number of children stored at this node.
1130    fn len(&self) -> usize {
1131        match self {
1132            Node::Internal(node) => node.children.len(),
1133            Node::Leaf(node) => node.keys.len(),
1134        }
1135    }
1136
1137    /// Search this node for the given key and return both the key and the value found.
1138    fn get_key_value(&self, cursor: Cursor) -> Option<(&Range<K>, &V)> {
1139        match self {
1140            Node::Leaf(node) => node.get_key_value(cursor),
1141            Node::Internal(node) => node.get_key_value(cursor),
1142        }
1143    }
1144
1145    /// The last key/value pair stored in this node.
1146    fn last_key_value(&self) -> Option<(&Range<K>, &V)> {
1147        match self {
1148            Node::Leaf(node) => node.last_key_value(),
1149            Node::Internal(node) => node.last_key_value(),
1150        }
1151    }
1152
1153    /// Converts a reference into a Node into a NodeRef.
1154    fn as_ref(&self) -> NodeRef<'_, K, V> {
1155        match self {
1156            Node::Internal(node) => NodeRef::Internal(node),
1157            Node::Leaf(node) => NodeRef::Leaf(node),
1158        }
1159    }
1160
1161    /// Returns a reference to the node that contains the entry indicated by the cursor.
1162    ///
1163    /// Assumes the cursor is non-empty.
1164    fn get_containing_node(&self, cursor: Cursor) -> NodeRef<'_, K, V> {
1165        assert!(cursor.depth > 0);
1166        if cursor.depth == 1 {
1167            return self.as_ref();
1168        }
1169        match self {
1170            Node::Internal(node) => node.get_containing_node(cursor),
1171            Node::Leaf(_) => unreachable!("leaf nodes do not have children"),
1172        }
1173    }
1174
1175    /// Insert the given value at the location indicated by `cursor`.
1176    ///
1177    /// If the insertion causes this node to split, the node will always split into two instances
1178    /// of the same type of node.
1179    fn insert(&mut self, cursor: Cursor, range: Range<K>, value: V) -> InsertResult<K, V> {
1180        match self {
1181            Node::Internal(node) => Arc::make_mut(node).insert(cursor, range, value),
1182            Node::Leaf(node) => Arc::make_mut(node).insert(cursor, range, value),
1183        }
1184    }
1185
1186    /// Remove the entry indicated by `cursor`.
1187    fn remove(&mut self, cursor: Cursor) -> RemoveResult<K, V> {
1188        match self {
1189            Node::Internal(node) => Arc::make_mut(node).remove(cursor),
1190            Node::Leaf(node) => Arc::make_mut(node).remove(cursor),
1191        }
1192    }
1193
1194    /// Find the given key in this node.
1195    ///
1196    /// Updates `cursor` to point to the position indicated by `position`.
1197    fn find(&self, key: &K, position: CursorPosition, cursor: &mut Cursor) {
1198        match self {
1199            Node::Internal(node) => node.find(key, position, cursor),
1200            Node::Leaf(node) => node.find(key, position, cursor),
1201        }
1202    }
1203
1204    /// Find a gap that is at least as large as the given gap and is less than the given upper bound.
1205    ///
1206    /// Returns the end of the gap, which might above or below the node.
1207    fn find_gap_end(&self, gap: u64, upper_bound: &K) -> K {
1208        match self {
1209            Node::Leaf(node) => node.find_gap_end(gap, upper_bound),
1210            Node::Internal(node) => node.find_gap_end(gap, upper_bound),
1211        }
1212    }
1213
1214    #[cfg(test)]
1215    fn validate_max_gap(&self) {
1216        match self {
1217            Node::Leaf(node) => node.validate_max_gap(),
1218            Node::Internal(node) => node.validate_max_gap(),
1219        }
1220    }
1221
1222    /// If this node is an internal node, validates that the keys held by this node are correct.
1223    ///
1224    /// Panics if the node (or a descendant) does not validate.
1225    #[cfg(test)]
1226    fn validate_keys(&self) {
1227        if let Node::Internal(node) = self {
1228            node.validate_keys();
1229        }
1230    }
1231}
1232
1233/// A node in the btree.
1234#[derive(Clone, Debug)]
1235enum NodeRef<'a, K: Ord + Copy + Gap, V: Clone> {
1236    /// An internal node.
1237    Internal(&'a Arc<NodeInternal<K, V>>),
1238
1239    /// A leaf node.
1240    Leaf(&'a Arc<NodeLeaf<K, V>>),
1241}
1242
1243impl<'a, K, V> NodeRef<'a, K, V>
1244where
1245    K: Ord + Copy + Gap,
1246    V: Clone,
1247{
1248    /// The number of children stored at this node.
1249    fn len(&self) -> usize {
1250        match self {
1251            NodeRef::Internal(node) => node.children.len(),
1252            NodeRef::Leaf(node) => node.keys.len(),
1253        }
1254    }
1255}
1256
1257/// An iterator over the key-value pairs stored in a RangeMap.
1258#[derive(Debug)]
1259pub struct Iter<'a, K: Ord + Copy + Gap, V: Clone> {
1260    /// The state of the forward iteration.
1261    ///
1262    /// The cursor points to the next entry to enumerate.
1263    forward: Cursor,
1264
1265    /// The state of the backward iteration.
1266    ///
1267    /// The cursor points to the child that was most recently iterated or just past the end of the
1268    /// entry list if no entries have been enumerated from this leaf yet.
1269    backward: Cursor,
1270
1271    /// The root node of the tree.
1272    root: &'a Node<K, V>,
1273}
1274
1275impl<'a, K, V> Iter<'a, K, V>
1276where
1277    K: Ord + Copy + Gap,
1278    V: Clone,
1279{
1280    /// Whether the iterator is complete.
1281    ///
1282    /// Iteration stops when the forward and backward cursors meet.
1283    fn is_done(&self) -> bool {
1284        self.forward == self.backward
1285    }
1286}
1287
1288impl<'a, K, V> Iterator for Iter<'a, K, V>
1289where
1290    K: Ord + Copy + Gap,
1291    V: Clone,
1292{
1293    type Item = (&'a Range<K>, &'a V);
1294
1295    fn next(&mut self) -> Option<Self::Item> {
1296        while !self.is_done() {
1297            match self.root.get_containing_node(self.forward) {
1298                NodeRef::Leaf(leaf) => {
1299                    let index = self.forward.back();
1300                    if index < leaf.keys.len() {
1301                        let key = &leaf.keys[index];
1302                        let value = &leaf.values[index];
1303                        self.forward.increment_back();
1304                        return Some((key, value));
1305                    } else {
1306                        self.forward.pop_back();
1307                        self.forward.increment_back();
1308                    }
1309                }
1310                NodeRef::Internal(internal) => {
1311                    let index = self.forward.back();
1312                    if index < internal.children.len() {
1313                        self.forward.push_back(0);
1314                    } else {
1315                        self.forward.pop_back();
1316                        self.forward.increment_back();
1317                    }
1318                }
1319            }
1320        }
1321        None
1322    }
1323}
1324
1325impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V>
1326where
1327    K: Ord + Copy + Gap,
1328    V: Clone,
1329{
1330    fn next_back(&mut self) -> Option<Self::Item> {
1331        while !self.is_done() {
1332            match self.root.get_containing_node(self.backward) {
1333                NodeRef::Leaf(leaf) => {
1334                    let index = self.backward.back();
1335                    if index > 0 {
1336                        let key = &leaf.keys[index - 1];
1337                        let value = &leaf.values[index - 1];
1338                        self.backward.decrement_back();
1339                        return Some((key, value));
1340                    } else {
1341                        self.backward.pop_back();
1342                    }
1343                }
1344                NodeRef::Internal(internal) => {
1345                    let index = self.backward.back();
1346                    if index > 0 {
1347                        let child = internal.children.get_ref(index - 1);
1348                        self.backward.decrement_back();
1349                        self.backward.push_back(child.len());
1350                    } else {
1351                        self.backward.pop_back();
1352                    }
1353                }
1354            }
1355        }
1356        None
1357    }
1358}
1359
1360/// A map from ranges to values.
1361///
1362/// This map can be cloned efficiently. If the map is modified after being cloned, the relevant
1363/// parts of the map's internal structure will be copied lazily.
1364#[derive(Clone, Debug)]
1365pub struct RangeMap<K: Ord + Copy + Gap, V: Clone + Eq> {
1366    /// The root node of the tree.
1367    ///
1368    /// The root node is either a leaf of an internal node, depending on the number of entries in
1369    /// the map.
1370    node: Node<K, V>,
1371}
1372
1373impl<K, V> Default for RangeMap<K, V>
1374where
1375    K: Ord + Copy + Gap,
1376    V: Clone + Eq,
1377{
1378    fn default() -> Self {
1379        Self { node: Node::Leaf(Arc::new(NodeLeaf::empty())) }
1380    }
1381}
1382
1383impl<K, V> RangeMap<K, V>
1384where
1385    K: Ord + Copy + Gap,
1386    V: Clone + Eq,
1387{
1388    /// Whether this map contains any entries.
1389    pub fn is_empty(&self) -> bool {
1390        match &self.node {
1391            Node::Leaf(node) => node.keys.is_empty(),
1392            Node::Internal(_) => false,
1393        }
1394    }
1395
1396    /// Find the given key in this node.
1397    ///
1398    /// Returns a Cursor that points to the position indicated by `position`.
1399    fn find(&self, key: &K, position: CursorPosition) -> Cursor {
1400        let mut cursor = Cursor::default();
1401        self.node.find(key, position, &mut cursor);
1402        cursor
1403    }
1404
1405    /// If the entry indicated by the cursor contains `key`, returns the range and value stored at
1406    /// that entry.
1407    fn get_if_contains_key(&self, key: &K, cursor: Cursor) -> Option<(&Range<K>, &V)> {
1408        if let Some((range, value)) = self.node.get_key_value(cursor) {
1409            if range.contains(key) {
1410                return Some((range, value));
1411            }
1412        }
1413        None
1414    }
1415
1416    /// Searches the map for a range that contains the given key.
1417    ///
1418    /// Returns the range and value if such a range is found.
1419    pub fn get(&self, key: K) -> Option<(&Range<K>, &V)> {
1420        self.get_if_contains_key(&key, self.find(&key, CursorPosition::Left))
1421    }
1422
1423    /// Searches the map for all keys that overlap the given range.
1424    ///
1425    /// Returns an iterator over the keys.
1426    pub fn get_keys(&self, range: Range<K>) -> impl Iterator<Item = &Range<K>> {
1427        self.range(range).map(|(key, _)| key)
1428    }
1429
1430    /// The last range stored in this map.
1431    pub fn last_range(&self) -> Option<&Range<K>> {
1432        self.node.last_key_value().map(|(key, _)| key)
1433    }
1434
1435    /// Searches the map for a range that contains the given key.
1436    ///
1437    /// If such a range is found, returns a cursor to that entry, the range, and the value.
1438    fn get_cursor_key_value(&mut self, key: &K) -> Option<(Cursor, Range<K>, V)> {
1439        let cursor = self.find(key, CursorPosition::Left);
1440        self.get_if_contains_key(key, cursor)
1441            .map(|(range, value)| (cursor, range.clone(), value.clone()))
1442    }
1443
1444    /// Find a gap that is at least as large as the given gap and is less than the given upper bound.
1445    ///
1446    /// Returns the end of the gap, which might above or below the entries in the map.
1447    pub fn find_gap_end(&self, gap: u64, upper_bound: &K) -> K {
1448        self.node.find_gap_end(gap, upper_bound)
1449    }
1450
1451    /// Remove the entry with the given key from the map.
1452    ///
1453    /// If the key was present in the map, returns the value previously stored at the given key.
1454    pub fn remove(&mut self, range: Range<K>) -> Vec<V> {
1455        let mut removed_values = vec![];
1456
1457        if range.end <= range.start {
1458            return removed_values;
1459        }
1460
1461        if let Some((cursor, old_range, v)) = self.get_cursor_key_value(&range.start) {
1462            // Remove that range from the map.
1463            removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1464
1465            // If the removed range extends after the end of the given range,
1466            // re-insert the part of the old range that extends beyond the end
1467            // of the given range.
1468            if old_range.end > range.end {
1469                self.insert_range_internal(range.end..old_range.end, v.clone());
1470            }
1471
1472            // If the removed range extends before the start of the given
1473            // range, re-insert the part of the old range that extends before
1474            // the start of the given range.
1475            if old_range.start < range.start {
1476                self.insert_range_internal(old_range.start..range.start, v);
1477            }
1478
1479            // Notice that we can end up splitting the old range into two
1480            // separate ranges if the old range extends both beyond the given
1481            // range and before the given range.
1482        }
1483
1484        if let Some((cursor, old_range, v)) = self.get_cursor_key_value(&range.end) {
1485            // If the old range starts before the removed range, we need to trim the old range.
1486            // TODO: Optimize with replace once available.
1487            if old_range.start < range.end {
1488                // Remove that range from the map.
1489                removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1490
1491                // If the removed range extends after the end of the given range,
1492                // re-insert the part of the old range that extends beyond the end
1493                // of the given range.
1494                if old_range.end > range.end {
1495                    self.insert_range_internal(range.end..old_range.end, v);
1496                }
1497            }
1498        }
1499
1500        let doomed = self.range(range.start..range.end).map(|(r, _)| r.start).collect::<Vec<_>>();
1501        for key in doomed {
1502            let cursor = self.find(&key, CursorPosition::Left);
1503            removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1504        }
1505
1506        #[cfg(test)]
1507        self.node.validate_max_gap();
1508
1509        #[cfg(test)]
1510        self.node.validate_keys();
1511
1512        removed_values
1513    }
1514
1515    /// Insert the given range and value at the location indicated by the cursor.
1516    fn insert_at(&mut self, cursor: Cursor, range: Range<K>, value: V) -> Option<V> {
1517        let result = self.node.insert(cursor, range, value);
1518        match result {
1519            InsertResult::Inserted => None,
1520            InsertResult::SplitLeaf(key, right) => {
1521                let mut keys = ArrayVec::new();
1522                let mut children = ArrayVec::new();
1523                keys.push(key);
1524                let Node::Leaf(left) = self.node.clone() else {
1525                    unreachable!("non-leaf node split into leaf node");
1526                };
1527                children.push(left);
1528                children.push(right);
1529                self.node =
1530                    Node::Internal(Arc::new(NodeInternal::new(keys, ChildList::Leaf(children))));
1531                None
1532            }
1533            InsertResult::SplitInternal(key, right) => {
1534                let mut keys = ArrayVec::new();
1535                let mut children = ArrayVec::new();
1536                keys.push(key);
1537                let Node::Internal(left) = self.node.clone() else {
1538                    unreachable!("non-internal node split into internal node");
1539                };
1540                children.push(left);
1541                children.push(right);
1542                let mut new_internal = NodeInternal::new(keys, ChildList::Internal(children));
1543                new_internal.update_max_gap();
1544                self.node = Node::Internal(Arc::new(new_internal));
1545                None
1546            }
1547        }
1548    }
1549
1550    /// Insert the given range and value.
1551    ///
1552    /// Assumes the range is empty and that adjacent ranges have different values.
1553    fn insert_range_internal(&mut self, range: Range<K>, value: V) -> Option<V> {
1554        assert!(range.end > range.start);
1555        let cursor = self.find(&range.start, CursorPosition::Left);
1556        self.insert_at(cursor, range, value)
1557    }
1558
1559    /// Inserts a range with the given value.
1560    ///
1561    /// The keys included in the given range are now associated with the given
1562    /// value. If those keys were previously associated with another value,
1563    /// are no longer associated with that previous value.
1564    ///
1565    /// This method can cause one or more values in the map to be dropped if
1566    /// the all of the keys associated with those values are contained within
1567    /// the given range.
1568    ///
1569    /// If the inserted range is directly adjacent to another range with an equal value, the
1570    /// inserted range will be merged with the adjacent ranges.
1571    pub fn insert(&mut self, mut range: Range<K>, value: V) {
1572        if range.end <= range.start {
1573            return;
1574        }
1575        self.remove(range.clone());
1576
1577        // Check for a range directly before this one. If it exists, it will be the last range with
1578        // start < range.start.
1579        if let Some((prev_range, prev_value)) = self.range(..range.start).next_back() {
1580            if prev_range.end == range.start && value == *prev_value {
1581                let cursor = self.find(&prev_range.start, CursorPosition::Left);
1582                range.start = prev_range.start;
1583                self.remove_at(cursor);
1584            }
1585        }
1586
1587        // Check for a range directly after. If it exists, we can look it up by exact start value
1588        // of range.end.
1589        if let Some((cursor, next_range, next_value)) = self.get_cursor_key_value(&range.end) {
1590            if next_range.start == range.end && value == next_value {
1591                range.end = next_range.end;
1592                self.remove_at(cursor);
1593            }
1594        }
1595
1596        self.insert_range_internal(range, value);
1597
1598        #[cfg(test)]
1599        self.node.validate_max_gap();
1600
1601        #[cfg(test)]
1602        self.node.validate_keys();
1603    }
1604
1605    /// Remove the entry with the given cursor from the map.
1606    fn remove_at(&mut self, cursor: Cursor) -> Option<V> {
1607        let result = self.node.remove(cursor);
1608        match result {
1609            RemoveResult::NotFound => None,
1610            RemoveResult::Removed(RemoveState { removed_value, .. }) => Some(removed_value),
1611            RemoveResult::Underflow(RemoveState { removed_value, .. }) => {
1612                match &mut self.node {
1613                    Node::Leaf(_) => {
1614                        // Nothing we can do about an underflow of a single leaf node at the root.
1615                    }
1616                    Node::Internal(node) => {
1617                        // If the root has underflown to a trivial list, we can shrink the tree.
1618                        if node.children.len() == 1 {
1619                            self.node = node.children.get(0);
1620                        }
1621                    }
1622                }
1623                Some(removed_value)
1624            }
1625        }
1626    }
1627
1628    /// Iterate through the keys and values stored in the map.
1629    pub fn iter(&self) -> Iter<'_, K, V> {
1630        Iter {
1631            forward: Cursor::with_index(0),
1632            backward: Cursor::with_index(self.node.len()),
1633            root: &self.node,
1634        }
1635    }
1636
1637    /// Create the cursor stack for the start bound of the given range.
1638    fn find_start_bound(&self, bounds: &impl RangeBounds<K>) -> Cursor {
1639        let key = match bounds.start_bound() {
1640            Bound::Included(key) => key,
1641            Bound::Excluded(key) => key,
1642            Bound::Unbounded => {
1643                return Cursor::with_index(0);
1644            }
1645        };
1646        self.find(key, CursorPosition::Left)
1647    }
1648
1649    /// Create the cursor stack for the end bound of the given range.
1650    fn find_end_bound(&self, bounds: &impl RangeBounds<K>) -> Cursor {
1651        let key = match bounds.end_bound() {
1652            Bound::Included(key) => key,
1653            Bound::Excluded(key) => key,
1654            Bound::Unbounded => {
1655                return Cursor::with_index(self.node.len());
1656            }
1657        };
1658        self.find(key, CursorPosition::Right)
1659    }
1660
1661    /// Iterate through the keys and values stored in the given range in the map.
1662    pub fn range(&self, bounds: impl RangeBounds<K>) -> Iter<'_, K, V> {
1663        let forward = self.find_start_bound(&bounds);
1664        let backward = self.find_end_bound(&bounds);
1665        Iter { forward, backward, root: &self.node }
1666    }
1667
1668    /// Find the first range in the map that starts at or after the given key.
1669    pub fn find_at_or_after(&self, key: K) -> Option<(&Range<K>, &V)> {
1670        let mut iter = self.range(key..).filter(move |(range, _)| key <= range.start);
1671        iter.next()
1672    }
1673}
1674
1675#[cfg(test)]
1676mod test {
1677    use super::*;
1678
1679    #[::fuchsia::test]
1680    fn test_empty() {
1681        let mut map = RangeMap::<u32, i32>::default();
1682
1683        assert!(map.get(12).is_none());
1684        map.remove(10..34);
1685        // This is a test to make sure we can handle reversed ranges
1686        #[allow(clippy::reversed_empty_ranges)]
1687        map.remove(34..10);
1688    }
1689
1690    #[::fuchsia::test]
1691    fn test_is_empty() {
1692        let mut map = RangeMap::<u32, i32>::default();
1693
1694        // Test empty map
1695        assert!(map.is_empty());
1696
1697        // Test map with 5 entries
1698        for i in 0..5 {
1699            let start = i * 10;
1700            let end = start + 5;
1701            map.insert(start..end, i as i32);
1702        }
1703        assert!(!map.is_empty());
1704
1705        // Remove all entries
1706        for i in 0..5 {
1707            let start = i * 10;
1708            let end = start + 5;
1709            map.remove(start..end);
1710        }
1711        assert!(map.is_empty());
1712
1713        // Test map with 50 entries
1714        for i in 0..50 {
1715            let start = i * 10;
1716            let end = start + 5;
1717            map.insert(start..end, i as i32);
1718        }
1719        assert!(!map.is_empty());
1720
1721        // Remove all entries
1722        for i in 0..50 {
1723            let start = i * 10;
1724            let end = start + 5;
1725            map.remove(start..end);
1726        }
1727        assert!(map.is_empty());
1728    }
1729
1730    #[::fuchsia::test]
1731    fn test_insert_into_empty() {
1732        let mut map = RangeMap::<u32, i32>::default();
1733
1734        map.insert(10..34, -14);
1735
1736        assert_eq!((&(10..34), &-14), map.get(12).unwrap());
1737        assert_eq!((&(10..34), &-14), map.get(10).unwrap());
1738        assert!(map.get(9).is_none());
1739        assert_eq!((&(10..34), &-14), map.get(33).unwrap());
1740        assert!(map.get(34).is_none());
1741    }
1742
1743    #[::fuchsia::test]
1744    fn test_iter() {
1745        let mut map = RangeMap::<u32, i32>::default();
1746
1747        map.insert(10..34, -14);
1748        map.insert(74..92, -12);
1749
1750        let mut iter = map.iter();
1751
1752        assert_eq!(iter.next().expect("missing elem"), (&(10..34), &-14));
1753        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
1754
1755        assert!(iter.next().is_none());
1756
1757        let entry = map.find_at_or_after(10);
1758        assert_eq!(entry.expect("missing elem"), (&(10..34), &-14));
1759        let entry = map.find_at_or_after(11);
1760        assert_eq!(entry.expect("missing elem"), (&(74..92), &-12));
1761        let entry = map.find_at_or_after(74);
1762        assert_eq!(entry.expect("missing elem"), (&(74..92), &-12));
1763        let entry = map.find_at_or_after(75);
1764        assert_eq!(entry, None);
1765
1766        assert_eq!(map.range(..9).collect::<Vec<_>>(), vec![]);
1767        assert_eq!(map.range(..34).collect::<Vec<_>>(), vec![(&(10..34), &-14)]);
1768        assert_eq!(map.range(..74).collect::<Vec<_>>(), vec![(&(10..34), &-14)]);
1769        assert_eq!(map.range(..75).collect::<Vec<_>>(), vec![(&(10..34), &-14), (&(74..92), &-12)]);
1770        assert_eq!(map.range(..91).collect::<Vec<_>>(), vec![(&(10..34), &-14), (&(74..92), &-12)]);
1771        assert_eq!(map.range(..92).collect::<Vec<_>>(), vec![(&(10..34), &-14), (&(74..92), &-12)]);
1772    }
1773
1774    #[::fuchsia::test]
1775    fn test_remove_overlapping_edge() {
1776        let mut map = RangeMap::<u32, i32>::default();
1777
1778        map.insert(10..34, -14);
1779
1780        map.remove(2..11);
1781        assert_eq!((&(11..34), &-14), map.get(11).unwrap());
1782
1783        map.remove(33..42);
1784        assert_eq!((&(11..33), &-14), map.get(12).unwrap());
1785    }
1786
1787    #[::fuchsia::test]
1788    fn test_remove_middle_splits_range() {
1789        let mut map = RangeMap::<u32, i32>::default();
1790
1791        map.insert(10..34, -14);
1792        map.remove(15..18);
1793
1794        assert_eq!((&(10..15), &-14), map.get(12).unwrap());
1795        assert_eq!((&(18..34), &-14), map.get(20).unwrap());
1796    }
1797
1798    #[::fuchsia::test]
1799    fn test_remove_upper_half_of_split_range_leaves_lower_range() {
1800        let mut map = RangeMap::<u32, i32>::default();
1801
1802        map.insert(10..34, -14);
1803        map.remove(15..18);
1804        map.insert(2..7, -21);
1805        map.remove(20..42);
1806
1807        assert_eq!((&(2..7), &-21), map.get(5).unwrap());
1808        assert_eq!((&(10..15), &-14), map.get(12).unwrap());
1809    }
1810
1811    #[::fuchsia::test]
1812    fn test_range_map_overlapping_insert() {
1813        let mut map = RangeMap::<u32, i32>::default();
1814
1815        map.insert(2..7, -21);
1816        map.insert(5..9, -42);
1817        map.insert(1..3, -43);
1818        map.insert(6..8, -44);
1819
1820        assert_eq!((&(1..3), &-43), map.get(2).unwrap());
1821        assert_eq!((&(3..5), &-21), map.get(4).unwrap());
1822        assert_eq!((&(5..6), &-42), map.get(5).unwrap());
1823        assert_eq!((&(6..8), &-44), map.get(7).unwrap());
1824    }
1825
1826    #[::fuchsia::test]
1827    fn test_intersect_single() {
1828        let mut map = RangeMap::<u32, i32>::default();
1829
1830        map.insert(2..7, -10);
1831
1832        let mut iter = map.range(3..4);
1833        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1834        assert_eq!(iter.next(), None);
1835
1836        let mut iter = map.range(2..3);
1837        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1838        assert_eq!(iter.next(), None);
1839
1840        let mut iter = map.range(1..4);
1841        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1842        assert_eq!(iter.next(), None);
1843
1844        let mut iter = map.range(1..2);
1845        assert_eq!(iter.next(), None);
1846
1847        let mut iter = map.range(6..7);
1848        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1849        assert_eq!(iter.next(), None);
1850    }
1851
1852    #[::fuchsia::test]
1853    fn test_intersect_multiple() {
1854        let mut map = RangeMap::<u32, i32>::default();
1855
1856        map.insert(2..7, -10);
1857        map.insert(7..9, -20);
1858        map.insert(10..11, -30);
1859
1860        let mut iter = map.range(3..8);
1861        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1862        assert_eq!(iter.next(), Some((&(7..9), &-20)));
1863        assert_eq!(iter.next(), None);
1864
1865        let mut iter = map.range(3..11);
1866        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1867        assert_eq!(iter.next(), Some((&(7..9), &-20)));
1868        assert_eq!(iter.next(), Some((&(10..11), &-30)));
1869        assert_eq!(iter.next(), None);
1870    }
1871
1872    #[::fuchsia::test]
1873    fn test_intersect_no_gaps() {
1874        let mut map = RangeMap::<u32, i32>::default();
1875
1876        map.insert(0..1, -10);
1877        map.insert(1..2, -20);
1878        map.insert(2..3, -30);
1879
1880        let mut iter = map.range(0..3);
1881        assert_eq!(iter.next(), Some((&(0..1), &-10)));
1882        assert_eq!(iter.next(), Some((&(1..2), &-20)));
1883        assert_eq!(iter.next(), Some((&(2..3), &-30)));
1884        assert_eq!(iter.next(), None);
1885    }
1886
1887    #[test]
1888    fn test_merging() {
1889        let mut map = RangeMap::<u32, i32>::default();
1890
1891        map.insert(1..2, -10);
1892        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..2), &-10)]);
1893        map.insert(3..4, -10);
1894        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..2), &-10), (&(3..4), &-10)]);
1895        map.insert(2..3, -10);
1896        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..4), &-10)]);
1897        map.insert(0..1, -10);
1898        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(0..4), &-10)]);
1899        map.insert(4..5, -10);
1900        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(0..5), &-10)]);
1901        map.insert(2..3, -20);
1902        assert_eq!(
1903            map.iter().collect::<Vec<_>>(),
1904            vec![(&(0..2), &-10), (&(2..3), &-20), (&(3..5), &-10)]
1905        );
1906    }
1907
1908    #[test]
1909    fn test_remove_multiple_ranges_exact_query() {
1910        let first = 15..21;
1911        let second = first.end..29;
1912
1913        let mut map = RangeMap::default();
1914        map.insert(first.clone(), 1);
1915        map.insert(second.clone(), 2);
1916
1917        assert_eq!(map.remove(first.start..second.end), &[1, 2]);
1918    }
1919
1920    #[::fuchsia::test]
1921    fn test_large_insert_and_remove() {
1922        let mut map = RangeMap::<u32, i32>::default();
1923        let num_entries = 1000;
1924
1925        // Insert a large number of entries
1926        for i in 0..num_entries {
1927            let start = i as u32 * 10;
1928            let end = start + 5;
1929            let value = i as i32;
1930            map.insert(start..end, value);
1931        }
1932
1933        // Verify that all inserted entries can be retrieved
1934        for i in 0..num_entries {
1935            let start = i as u32 * 10;
1936            let end = start + 5;
1937            let point = start + 2;
1938            if let Some((range, value)) = map.get(point) {
1939                assert!(range.start <= point && point < range.end);
1940                assert_eq!(*range, start..end);
1941                assert_eq!(*value, i as i32);
1942            } else {
1943                panic!("Expected to find a range for point {}", point);
1944            }
1945        }
1946
1947        // Remove a large number of entries
1948        for i in 0..num_entries {
1949            let start = i as u32 * 10;
1950            let end = start + 5;
1951            map.remove(start..end);
1952        }
1953
1954        // Verify that the map is empty after removing all entries
1955        assert!(map.is_empty());
1956    }
1957
1958    #[::fuchsia::test]
1959    fn test_large_insert_and_remove_overlapping() {
1960        let mut map = RangeMap::<u32, i32>::default();
1961        let num_entries = 1000;
1962
1963        // Insert a large number of entries with overlapping ranges
1964        for i in 0..num_entries {
1965            let start = i as u32 * 5;
1966            let end = start + 20;
1967            let value = i as i32;
1968            map.insert(start..end, value);
1969        }
1970
1971        // Verify that all inserted entries can be retrieved
1972        for i in 0..num_entries {
1973            let point = i as u32 * 5 + 1;
1974            if let Some((range, value)) = map.get(point) {
1975                assert!(range.start <= point && point < range.end);
1976                assert_eq!(*value, i as i32);
1977            } else {
1978                panic!("Expected to find a range for point {}", point);
1979            }
1980        }
1981
1982        // Remove a large number of entries with overlapping ranges
1983        for i in 0..num_entries {
1984            let start = i as u32 * 5;
1985            let end = start + 20;
1986            map.remove(start..end);
1987        }
1988
1989        // Verify that the map is empty after removing all entries
1990        assert!(map.is_empty());
1991    }
1992
1993    #[::fuchsia::test]
1994    fn test_large_insert_and_get_specific_points() {
1995        let mut map = RangeMap::<u32, i32>::default();
1996        let num_entries = 1000;
1997        let mut inserted_ranges = Vec::new();
1998
1999        // Insert a large number of entries
2000        for i in 0..num_entries {
2001            let start = i as u32 * 10;
2002            let end = start + 5;
2003            let value = i as i32;
2004            map.insert(start..end, value);
2005            inserted_ranges.push((start..end, value));
2006        }
2007
2008        // Verify that specific points can be retrieved correctly
2009        for (range, value) in &inserted_ranges {
2010            let point = range.start + 2;
2011            let (retrieved_range, retrieved_value) = map.get(point).unwrap();
2012            assert_eq!(retrieved_range, range);
2013            assert_eq!(retrieved_value, value);
2014        }
2015    }
2016
2017    #[::fuchsia::test]
2018    fn test_large_insert_and_iter() {
2019        let mut map = RangeMap::<u32, i32>::default();
2020        let num_entries = 1000;
2021        let mut inserted_ranges = Vec::new();
2022
2023        // Insert a large number of entries
2024        for i in 0..num_entries {
2025            let start = i as u32 * 10;
2026            let end = start + 5;
2027            let value = i as i32;
2028            map.insert(start..end, value);
2029            inserted_ranges.push((start..end, value));
2030        }
2031
2032        // Verify that iter() returns all inserted entries
2033        let mut iter_ranges: Vec<(&Range<u32>, &i32)> = map.iter().collect();
2034        iter_ranges.sort_by_key(|(range, _)| range.start);
2035        let mut inserted_ranges_sorted: Vec<(Range<u32>, i32)> = inserted_ranges.clone();
2036        inserted_ranges_sorted.sort_by_key(|(range, _)| range.start);
2037
2038        assert_eq!(iter_ranges.len(), inserted_ranges_sorted.len());
2039        for (i, (range, value)) in iter_ranges.iter().enumerate() {
2040            assert_eq!(*range, &inserted_ranges_sorted[i].0);
2041            assert_eq!(*value, &inserted_ranges_sorted[i].1);
2042        }
2043    }
2044
2045    #[::fuchsia::test]
2046    fn test_large_insert_and_range() {
2047        let mut map = RangeMap::<u32, i32>::default();
2048        let num_entries = 1000;
2049
2050        // Insert a large number of entries
2051        for i in 0..num_entries {
2052            let start = i as u32 * 10;
2053            let end = start + 5;
2054            let value = i as i32;
2055            map.insert(start..end, value);
2056        }
2057
2058        // Verify range(start_point..)
2059        let start_point = 5000;
2060        let mut iter = map.range(start_point..);
2061        while let Some((range, _)) = iter.next() {
2062            assert!(range.start >= start_point);
2063        }
2064    }
2065
2066    #[::fuchsia::test]
2067    fn test_large_insert_and_iter_ending_at() {
2068        let mut map = RangeMap::<u32, i32>::default();
2069        let num_entries = 1000;
2070
2071        // Insert a large number of entries
2072        for i in 0..num_entries {
2073            let start = i as u32 * 10;
2074            let end = start + 5;
2075            let value = i as i32;
2076            map.insert(start..end, value);
2077        }
2078
2079        // Verify range(..end_point)
2080        let end_point = 5000;
2081        let mut iter = map.range(..end_point);
2082        while let Some((range, _)) = iter.next() {
2083            assert!(range.start < end_point);
2084        }
2085    }
2086
2087    #[::fuchsia::test]
2088    fn test_large_insert_and_intersection() {
2089        let mut map = RangeMap::<u32, i32>::default();
2090        let num_entries = 1000;
2091
2092        // Insert a large number of entries
2093        for i in 0..num_entries {
2094            let start = i as u32 * 10;
2095            let end = start + 5;
2096            let value = i as i32;
2097            map.insert(start..end, value);
2098        }
2099
2100        // Verify intersection()
2101        let intersect_start = 4000;
2102        let intersect_end = 4050;
2103        let mut iter = map.range(intersect_start..intersect_end);
2104        while let Some((range, _)) = iter.next() {
2105            assert!(range.start < intersect_end && range.end > intersect_start);
2106        }
2107    }
2108
2109    #[::fuchsia::test]
2110    fn test_large_insert_and_last_range() {
2111        let mut map = RangeMap::<u32, i32>::default();
2112        let num_entries = 1000;
2113        let mut last_range = None;
2114
2115        // Insert a large number of entries
2116        for i in 0..num_entries {
2117            let start = i as u32 * 10;
2118            let end = start + 5;
2119            let value = i as i32;
2120            map.insert(start..end, value);
2121            last_range = Some(start..end);
2122        }
2123
2124        // Verify last_range()
2125        if let Some(expected_last_range) = last_range {
2126            assert_eq!(map.last_range().unwrap(), &expected_last_range);
2127        }
2128    }
2129
2130    #[::fuchsia::test]
2131    fn test_large_insert_and_remove_alternating() {
2132        let mut map = RangeMap::<u32, i32>::default();
2133        let num_entries = 1000;
2134
2135        // Insert and remove entries in an alternating pattern
2136        for i in 0..num_entries {
2137            let start = i as u32 * 10;
2138            let end = start + 5;
2139            let value = i as i32;
2140            map.insert(start..end, value);
2141
2142            if i % 2 == 0 {
2143                // Remove every other entry
2144                map.remove(start..end);
2145            }
2146        }
2147
2148        // Verify that only the non-removed entries remain
2149        for i in 0..num_entries {
2150            let start = i as u32 * 10;
2151            let end = start + 5;
2152            let point = start + 2;
2153            if i % 2 != 0 {
2154                if let Some((range, value)) = map.get(point) {
2155                    assert!(range.start <= point && point < range.end);
2156                    assert_eq!(*range, start..end);
2157                    assert_eq!(*value, i as i32);
2158                } else {
2159                    panic!("Expected to find a range for point {}", point);
2160                }
2161            } else {
2162                assert!(map.get(point).is_none());
2163            }
2164        }
2165    }
2166
2167    #[::fuchsia::test]
2168    fn test_large_insert_and_remove_multiple_times() {
2169        let mut map = RangeMap::<u32, i32>::default();
2170        let num_entries = 200;
2171        let num_iterations = 5;
2172
2173        for _ in 0..num_iterations {
2174            // Insert a large number of entries
2175            for i in 0..num_entries {
2176                let start = i as u32 * 10;
2177                let end = start + 5;
2178                let value = i as i32;
2179                map.insert(start..end, value);
2180            }
2181
2182            // Remove a large number of entries
2183            for i in 0..num_entries {
2184                let start = i as u32 * 10;
2185                let end = start + 5;
2186                map.remove(start..end);
2187            }
2188        }
2189
2190        // Verify that the map is empty after multiple insert/remove cycles
2191        assert!(map.is_empty());
2192    }
2193
2194    #[::fuchsia::test]
2195    fn test_leaf_node_split() {
2196        let mut map = RangeMap::<u32, i32>::default();
2197
2198        // Fill the leaf node to capacity
2199        for i in 0..NODE_CAPACITY {
2200            let start = i as u32 * 10;
2201            let end = start + 5;
2202            map.insert(start..end, i as i32);
2203        }
2204
2205        // Verify the initial state
2206        assert_eq!(map.node.len(), NODE_CAPACITY);
2207
2208        {
2209            let mut map = map.clone();
2210
2211            // Insert a value that causes a split in the first half
2212            let split_start = 15;
2213            let split_end = 20;
2214            map.insert(split_start..split_end, -1);
2215
2216            // Verify the split
2217            assert!(matches!(map.node, Node::Internal(_)));
2218            if let Node::Internal(internal) = &map.node {
2219                assert_eq!(internal.children.len(), 2);
2220                assert!(internal.children.size_at(0) >= NODE_CAPACITY / 2);
2221                assert!(internal.children.size_at(1) >= NODE_CAPACITY / 2);
2222            }
2223        }
2224
2225        {
2226            let mut map = map.clone();
2227
2228            // Insert a value that causes a split in the second half
2229            let split_start = (NODE_CAPACITY as u32 * 10) + 5;
2230            let split_end = split_start + 5;
2231            map.insert(split_start..split_end, -2);
2232
2233            // Verify the second split
2234            if let Node::Internal(internal) = &map.node {
2235                assert_eq!(internal.children.len(), 2);
2236                assert!(internal.children.size_at(0) >= NODE_CAPACITY / 2);
2237            }
2238        }
2239    }
2240
2241    #[::fuchsia::test]
2242    fn test_merging_ranges_with_equal_values() {
2243        let mut map = RangeMap::<u32, i32>::default();
2244
2245        // Insert some initial ranges
2246        map.insert(10..20, 100);
2247        map.insert(30..40, 200);
2248        map.insert(50..60, 100);
2249
2250        // Merge adjacent ranges with equal values
2251        map.insert(20..30, 100);
2252        assert_eq!(map.get(15).unwrap(), (&(10..30), &100));
2253        assert_eq!(map.get(35).unwrap(), (&(30..40), &200));
2254        assert_eq!(map.get(55).unwrap(), (&(50..60), &100));
2255
2256        // Merge non-adjacent ranges with equal values
2257        map.insert(40..50, 100);
2258        assert_eq!(map.get(15).unwrap(), (&(10..30), &100));
2259        assert_eq!(map.get(35).unwrap(), (&(30..40), &200));
2260        assert_eq!(map.get(45).unwrap(), (&(40..60), &100));
2261
2262        // Insert a range with a different value
2263        map.insert(60..70, 300);
2264        assert_eq!(map.get(65).unwrap(), (&(60..70), &300));
2265
2266        // Merge a range with a different value
2267        map.insert(70..80, 300);
2268        assert_eq!(map.get(65).unwrap(), (&(60..80), &300));
2269
2270        // Merge a range with a different value at the beginning
2271        map.insert(0..10, 400);
2272        assert_eq!(map.get(5).unwrap(), (&(0..10), &400));
2273
2274        // Merge a range with a different value at the end
2275        map.insert(80..90, 400);
2276        assert_eq!(map.get(85).unwrap(), (&(80..90), &400));
2277
2278        // Merge a range with a different value in the middle
2279        map.insert(90..100, 400);
2280        assert_eq!(map.get(95).unwrap(), (&(80..100), &400));
2281        map.insert(100..110, 400);
2282        assert_eq!(map.get(95).unwrap(), (&(80..110), &400));
2283        map.insert(110..120, 400);
2284        assert_eq!(map.get(95).unwrap(), (&(80..120), &400));
2285
2286        // Merge a range with a different value in the middle
2287        map.insert(10..90, 400);
2288        assert_eq!(map.get(5).unwrap(), (&(0..120), &400));
2289    }
2290
2291    #[::fuchsia::test]
2292    fn test_large_insert_and_remove_with_gaps() {
2293        let mut map = RangeMap::<u32, i32>::default();
2294        let num_entries = 500;
2295
2296        // Insert entries with gaps
2297        for i in 0..num_entries {
2298            let start = i as u32 * 20;
2299            let end = start + 5;
2300            let value = i as i32;
2301            map.insert(start..end, value);
2302        }
2303
2304        // Remove entries with gaps
2305        for i in 0..num_entries {
2306            if i % 2 == 0 {
2307                let start = i as u32 * 20;
2308                let end = start + 5;
2309                map.remove(start..end);
2310            }
2311        }
2312
2313        // Verify the state of the map
2314        for i in 0..num_entries {
2315            let start = i as u32 * 20;
2316            let end = start + 5;
2317            let point = start + 2;
2318
2319            if i % 2 != 0 {
2320                if let Some((range, value)) = map.get(point) {
2321                    assert!(range.start <= point && point < range.end);
2322                    assert_eq!(*range, start..end);
2323                    assert_eq!(*value, i as i32);
2324                } else {
2325                    panic!("Expected to find a range for point {}", point);
2326                }
2327            } else {
2328                assert!(map.get(point).is_none());
2329            }
2330        }
2331    }
2332
2333    #[::fuchsia::test]
2334    fn test_critical_removal() {
2335        fn range_for(i: u32) -> Range<u32> {
2336            let start = i * 20;
2337            let end = start + 5;
2338            start..end
2339        }
2340
2341        // Index 0 cannot be the critical index.
2342        for critial_index in 1..100 {
2343            let mut map = RangeMap::<u32, i32>::default();
2344
2345            // Insert enough entries for force a split somewhere.
2346            for i in 0..100 {
2347                let value = i as i32;
2348                map.insert(range_for(i), value);
2349            }
2350
2351            // We don't know whether the split will occur at the critical index, but we try all the
2352            // possible indices to ensure that we test an index at which a split occurred.
2353            let critical_range = range_for(critial_index);
2354            map.remove(critical_range.clone());
2355
2356            // We now insert a range that spans the critical point. This range will be inserted to
2357            // left of the split.
2358            let value = -10 as i32;
2359            let spanning_range = (critical_range.start - 2)..(critical_range.start + 2);
2360            map.insert(spanning_range.clone(), value);
2361
2362            // Check to see if we can find the range by looking before the critical point.
2363            let key_before_split = critical_range.start - 1;
2364            assert_eq!(map.get(key_before_split), Some((&spanning_range, &value)));
2365
2366            // Check to see if we can find the range by looking after the critical point.
2367            let key_after_split = critical_range.start + 1;
2368            assert_eq!(map.get(key_after_split), Some((&spanning_range, &value)));
2369        }
2370    }
2371
2372    #[::fuchsia::test]
2373    fn test_find_gap_end_empty() {
2374        let map = RangeMap::<u32, i32>::default();
2375        assert_eq!(map.find_gap_end(10, &100), 100);
2376    }
2377
2378    #[::fuchsia::test]
2379    fn test_find_gap_end_single_range() {
2380        let mut map = RangeMap::<u32, i32>::default();
2381        map.insert(10..20, 1);
2382        assert_eq!(map.find_gap_end(10, &100), 100);
2383    }
2384
2385    #[::fuchsia::test]
2386    fn test_find_gap_end_multiple_ranges() {
2387        let mut map = RangeMap::<u32, i32>::default();
2388        map.insert(10..20, 1);
2389        map.insert(40..50, 2);
2390        map.insert(60..70, 3);
2391
2392        // Test finding gaps of various sizes
2393        assert_eq!(map.find_gap_end(5, &80), 80);
2394        assert_eq!(map.find_gap_end(10, &80), 80);
2395        assert_eq!(map.find_gap_end(11, &80), 40);
2396        assert_eq!(map.find_gap_end(20, &80), 40);
2397        assert_eq!(map.find_gap_end(21, &80), 10);
2398
2399        // Test finding gaps of various sizes with a lower bound
2400        assert_eq!(map.find_gap_end(5, &10), 10);
2401        assert_eq!(map.find_gap_end(5, &20), 10);
2402        assert_eq!(map.find_gap_end(5, &30), 30);
2403        assert_eq!(map.find_gap_end(5, &40), 40);
2404        assert_eq!(map.find_gap_end(5, &50), 40);
2405        assert_eq!(map.find_gap_end(5, &60), 60);
2406        assert_eq!(map.find_gap_end(5, &70), 60);
2407        assert_eq!(map.find_gap_end(5, &80), 80);
2408        assert_eq!(map.find_gap_end(5, &90), 90);
2409        assert_eq!(map.find_gap_end(5, &100), 100);
2410    }
2411
2412    #[::fuchsia::test]
2413    fn test_find_gap_end_rightmost() {
2414        let mut map = RangeMap::<u32, i32>::default();
2415        map.insert(10..20, 1);
2416        map.insert(30..40, 2);
2417        map.insert(50..60, 3);
2418        map.insert(70..80, 4);
2419
2420        // All gaps are size 10, should find the rightmost one
2421        assert_eq!(map.find_gap_end(10, &10), 10);
2422        assert_eq!(map.find_gap_end(10, &20), 10);
2423        assert_eq!(map.find_gap_end(10, &30), 30);
2424        assert_eq!(map.find_gap_end(10, &40), 30);
2425        assert_eq!(map.find_gap_end(10, &50), 50);
2426        assert_eq!(map.find_gap_end(10, &60), 50);
2427        assert_eq!(map.find_gap_end(10, &70), 70);
2428        assert_eq!(map.find_gap_end(10, &80), 70);
2429        assert_eq!(map.find_gap_end(10, &90), 90);
2430        assert_eq!(map.find_gap_end(10, &100), 100);
2431    }
2432
2433    #[::fuchsia::test]
2434    fn test_find_gap_end_large_map() {
2435        let mut map = RangeMap::<u32, i32>::default();
2436        let num_entries = 1000;
2437
2438        fn range_for(i: u32) -> Range<u32> {
2439            let start = i * (8000 - i) as u32;
2440            let end = start + 10;
2441            start..end
2442        }
2443
2444        // Insert ranges with decreasing gap sizes
2445        for i in 0..num_entries {
2446            map.insert(range_for(i), i as i32);
2447        }
2448
2449        let upper_bound = range_for(num_entries - 1).end;
2450
2451        for i in 0..num_entries - 1 {
2452            let current_range = range_for(i);
2453            let next_range = range_for(i + 1);
2454            let gap_size = current_range.end.measure_gap(&next_range.start);
2455            assert_eq!(map.find_gap_end(gap_size, &upper_bound), next_range.start);
2456        }
2457    }
2458
2459    #[::fuchsia::test]
2460    fn test_find_gap_end_after_removal() {
2461        let mut map = RangeMap::<u32, i32>::default();
2462        map.insert(10..20, 1);
2463        map.insert(30..40, 2);
2464        map.insert(50..60, 3);
2465
2466        assert_eq!(map.find_gap_end(12, &60), 10);
2467
2468        map.remove(30..35);
2469        assert_eq!(map.find_gap_end(12, &60), 35);
2470    }
2471
2472    #[::fuchsia::test]
2473    fn test_find_gap_end_adjacent_ranges() {
2474        let mut map = RangeMap::<u32, i32>::default();
2475        map.insert(10..20, 1);
2476        map.insert(20..30, 2);
2477        map.insert(30..40, 3);
2478
2479        // No gaps between ranges
2480        assert_eq!(map.find_gap_end(1, &100), 100);
2481    }
2482
2483    #[::fuchsia::test]
2484    fn test_find_gap_end_merging() {
2485        let mut map = RangeMap::<u32, i32>::default();
2486        map.insert(10..20, 1);
2487        map.insert(30..40, 2);
2488        map.insert(50..60, 2); // Same value as previous range
2489
2490        // Initially should find the last gap
2491        assert_eq!(map.find_gap_end(10, &100), 100);
2492
2493        // Merge the last two ranges
2494        map.insert(40..50, 1);
2495
2496        // Now should find the first gap
2497        assert_eq!(map.find_gap_end(10, &100), 100);
2498    }
2499
2500    #[::fuchsia::test]
2501    fn test_remove_empty_node() {
2502        let mut map = RangeMap::<u32, i32>::default();
2503
2504        // Fill two leaf nodes to capacity.
2505        for i in 0..(NODE_CAPACITY * 2) {
2506            let start = i as u32 * 10;
2507            let end = start + 1;
2508            map.insert(start..end, i as i32 * 100);
2509        }
2510
2511        // Insert at the right edge of the left leaf node, causing us to create a middle leaf node
2512        // with a single element.
2513        let i = NODE_CAPACITY - 1;
2514        let start = i as u32 * 10 + 2;
2515        let end = start + 1;
2516        let critical_range = start..end;
2517        map.insert(critical_range.clone(), i as i32 * 1000);
2518
2519        {
2520            // Remove the lone entry from the middle leaf node, ensuring that we pass our internal
2521            // validation checks in that scenario.
2522            let mut map = map.clone();
2523            map.remove(critical_range.clone());
2524        }
2525
2526        {
2527            let mut map = map.clone();
2528
2529            // Remove a bunch of nodes from the left side to encourage rebalancing from the right.
2530            for i in 0..(NODE_CAPACITY / 2 - 1) {
2531                let start = i as u32 * 10;
2532                let end = start + 1;
2533                map.remove(start..end);
2534            }
2535
2536            // Again, we remove the lone entry from the middle leaf node. This check is similar to
2537            // the previous check, but now the left leaf node has fewer elements, which would have
2538            // caused us to rebalance to the right in a previous version of this data structure.
2539            map.remove(critical_range);
2540        }
2541    }
2542}