range_map/
btree.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::borrow::Borrow;
7use std::cmp::{Eq, PartialEq};
8use std::fmt;
9use std::fmt::{Debug, Formatter};
10use std::ops::{Bound, Range, RangeBounds};
11use std::sync::Arc;
12
13/// The `B` constant for the btree.
14///
15/// Controls the size of the nodes inside the tree.
16const B: usize = 6;
17
18/// The capacity of nodes in the btree.
19const NODE_CAPACITY: usize = 2 * B;
20
21/// A location inside the btree.
22#[derive(Debug, Default, Clone, Copy)]
23struct Cursor {
24    /// The number of valid indices in the `indices` array.
25    depth: u8,
26
27    /// The indices of the entry, ordered from leaf to root.
28    indices: [u8; 7],
29}
30
31impl Cursor {
32    /// Create a cursor with a single index.
33    fn with_index(index: usize) -> Self {
34        let mut cursor = Self::default();
35        cursor.push(index);
36        cursor
37    }
38
39    /// Whether the cursor is empty.
40    ///
41    /// A cursor is empty if it contains no more indices. This happens when a traversal has reached
42    /// a leaf node.
43    fn is_empty(&self) -> bool {
44        self.depth == 0
45    }
46
47    /// Push an index onto the front of the cursor.
48    ///
49    /// The front of the cursor is towards the root of the tree.
50    fn push(&mut self, index: usize) {
51        self.indices[self.depth as usize] = index as u8;
52        self.depth += 1;
53    }
54
55    /// Push an index onto the back of the cursor.
56    ///
57    /// The back of the cursor is towards the leaves of the tree.
58    fn push_back(&mut self, index: usize) {
59        self.indices.rotate_right(1);
60        self.indices[0] = index as u8;
61        self.depth += 1;
62    }
63
64    /// Pop an index off the front of the cursor.
65    ///
66    /// The front of the cursor is towards the root of the tree.
67    fn pop(&mut self) -> Option<usize> {
68        if self.depth == 0 {
69            None
70        } else {
71            self.depth -= 1;
72            Some(self.indices[self.depth as usize] as usize)
73        }
74    }
75
76    /// Pop an index off the back of the cursor.
77    ///
78    /// The back of the cursor is towards the leaves of the tree.
79    fn pop_back(&mut self) -> Option<usize> {
80        if self.depth == 0 {
81            None
82        } else {
83            self.depth -= 1;
84            let index = self.indices[0] as usize;
85            self.indices.rotate_left(1);
86            Some(index)
87        }
88    }
89
90    /// The backmost index in the cursor.
91    ///
92    /// The back of the cursor is towards the leaves of the tree.
93    ///
94    /// Assumes the cursor is non-empty.
95    fn back(&self) -> usize {
96        self.indices[0] as usize
97    }
98
99    /// Increment the backmost index in the cursor.
100    ///
101    /// The back of the cursor is towards the leaves of the tree.
102    ///
103    /// Assumes the cursor is non-empty.
104    fn increment_back(&mut self) {
105        self.indices[0] += 1;
106    }
107
108    /// Decrement the backmost index in the cursor.
109    ///
110    /// The back of the cursor is towards the leaves of the tree.
111    ///
112    /// Assumes the cursor is non-empty.
113    fn decrement_back(&mut self) {
114        self.indices[0] -= 1;
115    }
116}
117
118impl PartialEq for Cursor {
119    fn eq(&self, other: &Self) -> bool {
120        if self.depth != other.depth {
121            return false;
122        }
123        for i in 0..self.depth {
124            if self.indices[i as usize] != other.indices[i as usize] {
125                return false;
126            }
127        }
128        true
129    }
130}
131
132impl Eq for Cursor {}
133
134/// Where to place the cursor relative to the given key.
135enum CursorPosition {
136    /// The given key represents a left edge of a range.
137    ///
138    /// Place the cursor to the left of a range containing the cursor.
139    Left,
140
141    /// The given key represents a right edge of a range.
142    ///
143    /// Place the cursor to the right of a range containing the cursor.
144    Right,
145}
146
147/// Search of the given key in the given array of ranges.
148///
149/// If the array contains a range that contains the key, returns the index of that range.
150/// Otherwise, returns the index at which the given key could be inserted into the array to
151/// maintain the ordering.
152fn binary_search<K: Ord>(key: &K, keys: &ArrayVec<Range<K>, NODE_CAPACITY>) -> usize {
153    let mut left = 0usize;
154    let mut right = keys.len();
155    while left < right {
156        let mid = left + (right - left) / 2;
157        // TODO: Consider `get_unchecked`.
158        let range = &keys[mid];
159        if key < &range.start {
160            // This range is too large.
161            right = mid;
162        } else if key < &range.end {
163            // We found the range that contains this key.
164            return mid;
165        } else {
166            // The key might be found in the next range.
167            left = mid + 1;
168        }
169    }
170    // The key falls between two ranges. Return the index at which this key could be inserted to
171    // maintain the ordering.
172    left
173}
174
175/// A leaf node in the btree.
176///
177/// Stores a flat map of keys to values, with the `i`th entry in the keys array corresponding to
178/// the `i`th entry in the values array. The balancing rules of the btree ensure that every
179/// non-root leaf has between N and N/2 entries populated.
180#[derive(Clone)]
181struct NodeLeaf<K: Ord + Copy, V: Clone> {
182    /// The keys stored in this leaf node.
183    ///
184    /// We store the key in a dense array to improve cache performance during lookups. We often
185    /// need to binary-search the keys in a given leaf node, which means having those keys close
186    /// together improves cache performance.
187    keys: ArrayVec<Range<K>, NODE_CAPACITY>,
188
189    /// The value stored in this leaf node.
190    values: ArrayVec<V, NODE_CAPACITY>,
191}
192
193/// Shows the map structure of the leaf node.
194impl<K, V> Debug for NodeLeaf<K, V>
195where
196    K: Debug + Ord + Copy,
197    V: Debug + Clone,
198{
199    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
200        f.debug_map().entries(self.keys.iter().zip(self.values.iter())).finish()
201    }
202}
203
204/// The result of performing an insertion into a btree node.
205enum InsertResult<K: Ord + Copy, V: Clone> {
206    /// The value was successfully inserted into an empty slot.
207    Inserted,
208
209    /// The value was inserted into an empty slot in a leaf node but that insertion caused the
210    /// leaf node to exceed its capacity and split into two leaf nodes. The existing leaf node
211    /// now holds the entries to the left of the split and the entries to the right of the split
212    /// are returned. The split occurred at the returned key.
213    SplitLeaf(K, Arc<NodeLeaf<K, V>>),
214
215    /// The value was inserted into an empty slot in a subtree but that insertion caused the
216    /// internal node to exceed its capacity and split into two internal nodes. The internal node
217    /// now holds the entries to the left of the split and the entries to the right of the split
218    /// are returned. The split occurred at the returned key.
219    SplitInternal(K, Arc<NodeInternal<K, V>>),
220}
221
222/// State information returned from the removal operation.
223struct RemoveState<K: Ord + Copy, V: Clone> {
224    /// The minimal key for this subtree, if available.
225    ///
226    /// If the minimal key for this subtree is not available, then the minimal key is guaranteed
227    /// to be unchanged by this operation.
228    maybe_split_key: Option<K>,
229
230    /// The value previously stored at this key.
231    removed_value: V,
232}
233
234/// The intermediate result of a remove operation.
235///
236/// The root of the CowTree converts this result value into `Option<T>`, per the usual map
237/// interface.
238enum RemoveResult<K: Ord + Copy, V: Clone> {
239    /// The key the client asked to remove was not found in the map.
240    NotFound,
241
242    /// The key was successfully removed from the map.
243    ///
244    Removed(RemoveState<K, V>),
245
246    /// The key was removed from the map but the node that previously contained that node no longer
247    /// has sufficient children.
248    ///
249    /// The caller is responsible for rebalancing its children to ensure that each node has at
250    /// least this minimum number of children. If the balance invariant can be resolved locally,
251    /// the caller should return `Removed` to its caller. If rebalancing the local children
252    /// causes this node to have fewer than the minimum number of children, the caller should
253    /// return `Underflow` to its caller.
254    Underflow(RemoveState<K, V>),
255}
256
257impl<K, V> NodeLeaf<K, V>
258where
259    K: Ord + Copy,
260    V: Clone,
261{
262    /// Create an empty leaf node.
263    ///
264    /// Empty leaf nodes are used only at the root of the tree.
265    fn empty() -> Self {
266        Self { keys: ArrayVec::new(), values: ArrayVec::new() }
267    }
268
269    /// Gets the index in this leaf that corresponds to the given cursor.
270    ///
271    /// Assumes the cursor contains exactly one index.
272    ///
273    /// Returns `None` if the cursor points beyond the end if this node.
274    fn get_index(&self, mut cursor: Cursor) -> Option<usize> {
275        let index = cursor.pop().expect("Cursor has sufficient depth");
276        assert!(cursor.is_empty(), "Cursor has excess depth");
277        if index >= self.keys.len() {
278            return None;
279        }
280        Some(index)
281    }
282
283    /// Search this leaf for the given key and return both the key and the value found.
284    fn get_key_value(&self, cursor: Cursor) -> Option<(&Range<K>, &V)> {
285        if let Some(index) = self.get_index(cursor) {
286            let key = &self.keys[index];
287            let value = &self.values[index];
288            Some((key, value))
289        } else {
290            None
291        }
292    }
293
294    /// The last key/value pair stored in this leaf.
295    fn last_key_value(&self) -> Option<(&Range<K>, &V)> {
296        let key = self.keys.last()?;
297        let value = self.values.last()?;
298        Some((key, value))
299    }
300
301    /// Find the given key in this node.
302    ///
303    /// Updates `cursor` to point to the position indicated by `position`.
304    fn find(&self, key: &K, position: CursorPosition, cursor: &mut Cursor) {
305        let index = binary_search(key, &self.keys);
306        match position {
307            CursorPosition::Left => {
308                cursor.push(index);
309            }
310            CursorPosition::Right => {
311                if let Some(range) = self.keys.get(index) {
312                    if *key > range.start {
313                        cursor.push(index + 1);
314                        return;
315                    }
316                }
317                cursor.push(index);
318            }
319        }
320    }
321
322    /// Insert the given entry at the location indicated by `cursor`.
323    ///
324    /// Inserting a value into a leaf node might cause this node to split into two leaf nodes.
325    fn insert(&mut self, mut cursor: Cursor, range: Range<K>, value: V) -> InsertResult<K, V> {
326        let index = cursor.pop().expect("valid cursor");
327        if self.keys.len() == NODE_CAPACITY {
328            if index == NODE_CAPACITY {
329                let mut keys = ArrayVec::new();
330                let mut values = ArrayVec::new();
331                let key = range.start;
332                keys.push(range);
333                values.push(value);
334                return InsertResult::SplitLeaf(key, Arc::new(Self { keys, values }));
335            }
336            let middle = NODE_CAPACITY / 2;
337            assert!(middle > 0);
338            let mut right = Self {
339                keys: self.keys.drain(middle..).collect(),
340                values: self.values.drain(middle..).collect(),
341            };
342            if index <= middle {
343                self.keys.insert(index, range);
344                self.values.insert(index, value);
345            } else {
346                right.keys.insert(index - middle, range);
347                right.values.insert(index - middle, value);
348            }
349            InsertResult::SplitLeaf(right.keys[0].start, Arc::new(right))
350        } else {
351            self.keys.insert(index, range);
352            self.values.insert(index, value);
353            InsertResult::Inserted
354        }
355    }
356
357    /// Remove the entry indicated by `cursor`.
358    fn remove(&mut self, cursor: Cursor) -> RemoveResult<K, V> {
359        if let Some(index) = self.get_index(cursor) {
360            self.keys.remove(index);
361            let removed_value = self.values.remove(index);
362            let maybe_split_key = self.keys.first().map(|range| range.start);
363            if self.keys.len() < NODE_CAPACITY / 2 {
364                RemoveResult::Underflow(RemoveState { maybe_split_key, removed_value })
365            } else {
366                RemoveResult::Removed(RemoveState { maybe_split_key, removed_value })
367            }
368        } else {
369            RemoveResult::NotFound
370        }
371    }
372}
373
374/// The children of an internal node in the btree.
375#[derive(Clone, Debug)]
376enum ChildList<K: Ord + Copy, V: Clone> {
377    /// Used when an internal node has leaf nodes as children.
378    Leaf(ArrayVec<Arc<NodeLeaf<K, V>>, NODE_CAPACITY>),
379
380    /// Used when an internal node has other internal nodes as children.
381    Internal(ArrayVec<Arc<NodeInternal<K, V>>, NODE_CAPACITY>),
382}
383
384impl<K, V> ChildList<K, V>
385where
386    K: Ord + Copy,
387    V: Clone,
388{
389    /// Create a child list that has no children.
390    fn new_empty(&self) -> Self {
391        match self {
392            ChildList::Leaf(_) => ChildList::Leaf(ArrayVec::new()),
393            ChildList::Internal(_) => ChildList::Internal(ArrayVec::new()),
394        }
395    }
396
397    /// The number of children for this node.
398    fn len(&self) -> usize {
399        match self {
400            ChildList::Leaf(children) => children.len(),
401            ChildList::Internal(children) => children.len(),
402        }
403    }
404
405    /// The number of gradchildren located at the child with the given index.
406    fn size_at(&self, i: usize) -> usize {
407        match self {
408            ChildList::Leaf(children) => children[i].keys.len(),
409            ChildList::Internal(children) => children[i].children.len(),
410        }
411    }
412
413    /// Obtain the child located at the given index.
414    fn get(&self, i: usize) -> Node<K, V> {
415        match self {
416            ChildList::Leaf(children) => Node::Leaf(children[i].clone()),
417            ChildList::Internal(children) => Node::Internal(children[i].clone()),
418        }
419    }
420
421    /// Get a reference to the child located at the given index.
422    fn get_ref(&self, i: usize) -> NodeRef<'_, K, V> {
423        match self {
424            ChildList::Leaf(children) => NodeRef::Leaf(&children[i]),
425            ChildList::Internal(children) => NodeRef::Internal(&children[i]),
426        }
427    }
428
429    /// Removes all the entries starting at the given index from the child list.
430    ///
431    /// The removed entries are returned in a new child list.
432    fn split_off(&mut self, index: usize) -> Self {
433        match self {
434            ChildList::Leaf(children) => ChildList::Leaf(children.drain(index..).collect()),
435            ChildList::Internal(children) => ChildList::Internal(children.drain(index..).collect()),
436        }
437    }
438
439    /// Removes all the entries prior to the given index from the child list.
440    ///
441    /// The removed entries are returned in a new child list.
442    fn split_off_front(&mut self, index: usize) -> Self {
443        match self {
444            ChildList::Leaf(children) => ChildList::Leaf(children.drain(..index).collect()),
445            ChildList::Internal(children) => ChildList::Internal(children.drain(..index).collect()),
446        }
447    }
448
449    /// Insert a child into the child list.
450    ///
451    /// The type of child node must match the type of the child list.
452    fn insert(&mut self, index: usize, child: Node<K, V>) {
453        match self {
454            ChildList::Leaf(children) => {
455                let Node::Leaf(leaf) = child else {
456                    unreachable!("Inserting a non-leaf into an internal node for leaf nodes.");
457                };
458                children.insert(index, leaf);
459            }
460            ChildList::Internal(children) => {
461                let Node::Internal(internal) = child else {
462                    unreachable!(
463                        "Inserting a non-internal into an internal node for internal nodes."
464                    );
465                };
466                children.insert(index, internal);
467            }
468        }
469    }
470
471    /// Remove the child at the given index from the child list.
472    fn remove(&mut self, index: usize) {
473        match self {
474            ChildList::Leaf(children) => {
475                children.remove(index);
476            }
477            ChildList::Internal(children) => {
478                children.remove(index);
479            }
480        }
481    }
482
483    /// Add the children from the given `ChildList` to this child list.
484    fn extend(&mut self, other: &Self) {
485        match (self, other) {
486            (ChildList::Leaf(self_children), ChildList::Leaf(other_children)) => {
487                self_children.extend(other_children.iter().cloned());
488            }
489            (ChildList::Internal(self_children), ChildList::Internal(other_children)) => {
490                self_children.extend(other_children.iter().cloned());
491            }
492            _ => unreachable!("Type mismatch while extending a child list."),
493        }
494    }
495}
496
497/// An internal node in the btree.
498#[derive(Clone, Debug)]
499struct NodeInternal<K: Ord + Copy, V: Clone> {
500    /// A cache of the keys that partition the keys in the children.
501    /// The key at index `i` is the smallest key stored in the subtree
502    /// of the `i`+1 child.
503    ///
504    /// We only ever store CAPACITY - 1 keys in this array.
505    keys: ArrayVec<K, NODE_CAPACITY>,
506
507    /// The children of this node.
508    children: ChildList<K, V>,
509}
510
511/// Get mutable references to two entries in a slice.
512///
513/// When rebalancing nodes, we need to mutate two nodes at the same time. Normally, if you take a
514/// mutable reference to one element of an array, the borrow checker prevents you from taking a
515/// mutable reference to a second element of the same array.
516///
517/// The nightly version of `std::primitive::slice` has `get_many_mut` to let you get mutable
518/// references to multiple elements. However, without that interface, the recommended approach for
519/// avoiding `unsafe` is to use `split_at_mut`.
520fn get_two_mut<T>(slice: &mut [T], i: usize, j: usize) -> (&mut T, &mut T) {
521    if i < j {
522        let (a, b) = slice.split_at_mut(j);
523        (&mut a[i], &mut b[0])
524    } else {
525        let (a, b) = slice.split_at_mut(i);
526        (&mut b[0], &mut a[j])
527    }
528}
529
530impl<K, V> NodeInternal<K, V>
531where
532    K: Ord + Copy,
533    V: Clone,
534{
535    /// The index of the child that might contain the given key.
536    ///
537    /// Searches the cached keys at this node to determine which child node might contain the given
538    /// key.
539    fn get_child_index(&self, key: &K) -> usize {
540        let p = self.keys.partition_point(|k| k < key);
541        if self.keys.get(p) == Some(key) {
542            // If the query key exactly matches the split key, then we need to look for this key to
543            // the right of the split.
544            p + 1
545        } else {
546            // Otherwise, we look to the left of the split.
547            p
548        }
549    }
550
551    /// Search this subtree for the given key and return both the key and the value found.
552    fn get_key_value(&self, mut cursor: Cursor) -> Option<(&Range<K>, &V)> {
553        let index = cursor.pop().expect("valid cursor");
554        match &self.children {
555            ChildList::Leaf(children) => children[index].get_key_value(cursor),
556            ChildList::Internal(children) => children[index].get_key_value(cursor),
557        }
558    }
559
560    /// Returns a reference to the node that contains the entry indicated by the cursor.
561    ///
562    /// Assumes the cursor points a descendant of this node.
563    fn get_containing_node(&self, mut cursor: Cursor) -> NodeRef<'_, K, V> {
564        debug_assert!(cursor.depth >= 2);
565        let index = cursor.pop().expect("valid cursor");
566        if cursor.depth == 1 {
567            return self.children.get_ref(index);
568        }
569        match &self.children {
570            ChildList::Leaf(_) => unreachable!("leaf nodes do not have children"),
571            ChildList::Internal(children) => children[index].get_containing_node(cursor),
572        }
573    }
574
575    /// The last key/value pair stored in this subtree.
576    fn last_key_value(&self) -> Option<(&Range<K>, &V)> {
577        match &self.children {
578            ChildList::Leaf(children) => {
579                children.last().expect("child lists are always non-empty").last_key_value()
580            }
581            ChildList::Internal(children) => {
582                children.last().expect("child lists are always non-empty").last_key_value()
583            }
584        }
585    }
586
587    /// Find the given key in this node.
588    ///
589    /// Updates `cursor` to point to the position indicated by `position`.
590    fn find(&self, key: &K, position: CursorPosition, cursor: &mut Cursor) {
591        let index = self.get_child_index(&key);
592        match &self.children {
593            ChildList::Leaf(children) => children[index].find(key, position, cursor),
594            ChildList::Internal(children) => children[index].find(key, position, cursor),
595        }
596        cursor.push(index);
597    }
598
599    /// Insert the given child node at `index` in this node.
600    ///
601    /// `key` must be the smallest key that occurs in the `child` subtree.
602    ///
603    /// The caller must ensure that the child is inserted in the correct location.
604    fn insert_child(&mut self, index: usize, key: K, child: Node<K, V>) -> InsertResult<K, V> {
605        let n = self.children.len();
606        if n == NODE_CAPACITY {
607            if index == NODE_CAPACITY {
608                let mut children = self.children.new_empty();
609                children.insert(0, child);
610                let right = Self { keys: ArrayVec::new(), children };
611                return InsertResult::SplitInternal(key, Arc::new(right));
612            }
613            let middle = NODE_CAPACITY / 2;
614            assert!(middle > 0);
615            let mut internal = Self {
616                keys: self.keys.drain(middle..).collect(),
617                children: self.children.split_off(middle),
618            };
619            let split_key = self.keys.pop().unwrap();
620            if index < middle {
621                self.keys.insert(index, key);
622                self.children.insert(index + 1, child);
623            } else {
624                internal.keys.insert(index - middle, key);
625                internal.children.insert(index - middle + 1, child);
626            }
627            debug_assert!(self.keys.len() + 1 == self.children.len());
628            debug_assert!(internal.keys.len() + 1 == internal.children.len());
629            InsertResult::SplitInternal(split_key, Arc::new(internal))
630        } else {
631            self.keys.insert(index, key);
632            self.children.insert(index + 1, child);
633            debug_assert!(self.keys.len() + 1 == self.children.len());
634            InsertResult::Inserted
635        }
636    }
637
638    /// Insert the given entry at the location indicated by `cursor`.
639    ///
640    /// Inserting a value into an internal node might cause this node to split into two internal
641    /// nodes.
642    fn insert(&mut self, mut cursor: Cursor, range: Range<K>, value: V) -> InsertResult<K, V> {
643        let index = cursor.pop().expect("valid cursor");
644        let result = match &mut self.children {
645            ChildList::Leaf(children) => {
646                Arc::make_mut(&mut children[index]).insert(cursor, range, value)
647            }
648            ChildList::Internal(children) => {
649                Arc::make_mut(&mut children[index]).insert(cursor, range, value)
650            }
651        };
652        match result {
653            InsertResult::Inserted => InsertResult::Inserted,
654            InsertResult::SplitLeaf(key, right) => self.insert_child(index, key, Node::Leaf(right)),
655            InsertResult::SplitInternal(key, right) => {
656                self.insert_child(index, key, Node::Internal(right))
657            }
658        }
659    }
660
661    /// Determine whether to rebalance the child with the given index to the left or to the right.
662    ///
663    /// Given a choice, we will rebalance the child with its larger neighbor.
664    ///
665    /// The indices returned are always sequential.
666    fn select_children_to_rebalance(&self, index: usize) -> (usize, usize) {
667        if index == 0 {
668            (index, index + 1)
669        } else if index == self.children.len() - 1 {
670            (index - 1, index)
671        } else {
672            let left_index = index - 1;
673            let left_size = self.children.size_at(left_index);
674            let right_index = index + 1;
675            let right_size = self.children.size_at(right_index);
676            if left_size > right_size {
677                (left_index, index)
678            } else {
679                (index, right_index)
680            }
681        }
682    }
683
684    /// Rebalance the child at the given index.
685    ///
686    /// If the child and its neighbor are sufficiently small, this function will merge them into a
687    /// single node.
688    fn rebalance_child(&mut self, index: usize) {
689        // Cannot rebalance if we have fewer than two children. This situation occurs only at the
690        // root of the tree.
691        if self.children.len() < 2 {
692            return;
693        }
694        let (left, right) = self.select_children_to_rebalance(index);
695        let n = self.children.size_at(left) + self.children.size_at(right);
696        match &mut self.children {
697            ChildList::Leaf(children) => {
698                let (left_shard_node, right_shared_node) = get_two_mut(children, left, right);
699                let left_node = Arc::make_mut(left_shard_node);
700                if n <= NODE_CAPACITY {
701                    // Merge the right node into the left node.
702                    left_node.keys.extend(right_shared_node.keys.iter().cloned());
703                    left_node.values.extend(right_shared_node.values.iter().cloned());
704                    self.keys.remove(left);
705                    self.children.remove(right);
706                } else {
707                    // Rebalance the elements between the nodes.
708                    let split = n / 2;
709                    let right_node = Arc::make_mut(right_shared_node);
710                    if left_node.values.len() < split {
711                        // Move elements from right to left.
712                        let move_count = split - left_node.values.len();
713                        left_node.keys.extend(right_node.keys.drain(..move_count));
714                        left_node.values.extend(right_node.values.drain(..move_count));
715                    } else {
716                        // Move elements from left to right.
717                        let mut keys = ArrayVec::new();
718                        keys.extend(left_node.keys.drain(split..));
719                        keys.extend(right_node.keys.iter().cloned());
720                        right_node.keys = keys;
721
722                        let mut values = ArrayVec::new();
723                        values.extend(left_node.values.drain(split..));
724                        values.extend(right_node.values.iter().cloned());
725                        right_node.values = values;
726                    }
727                    // Update the split key to reflect the new division between the nodes.
728                    self.keys[left] = right_node.keys[0].start;
729                }
730            }
731            ChildList::Internal(children) => {
732                let (left_shard_node, right_shared_node) = get_two_mut(children, left, right);
733                let left_node = Arc::make_mut(left_shard_node);
734                let old_split_key = &self.keys[left];
735                if n <= NODE_CAPACITY {
736                    // Merge the right node into the left node.
737                    left_node.keys.push(old_split_key.clone());
738                    left_node.keys.extend(right_shared_node.keys.iter().cloned());
739                    left_node.children.extend(&right_shared_node.children);
740                    debug_assert!(left_node.keys.len() + 1 == left_node.children.len());
741                    self.keys.remove(left);
742                    self.children.remove(right);
743                } else {
744                    // Rebalance the elements between the nodes.
745                    let split = n / 2;
746                    let split_key;
747                    let right_node = Arc::make_mut(right_shared_node);
748                    if left_node.children.len() < split {
749                        // Move elements from right to left.
750                        let move_count = split - left_node.children.len();
751                        left_node.keys.push(old_split_key.clone());
752                        left_node.keys.extend(right_node.keys.drain(..move_count));
753                        split_key =
754                            left_node.keys.pop().expect("must have moved at least one element");
755
756                        left_node.children.extend(&right_node.children.split_off_front(move_count));
757                        debug_assert!(left_node.keys.len() + 1 == left_node.children.len());
758                    } else {
759                        // Move elements from left to right.
760                        let mut it = left_node.keys.drain((split - 1)..);
761                        split_key = it.next().expect("must be moving at least one element");
762                        let mut keys = ArrayVec::new();
763                        keys.extend(it);
764                        keys.push(old_split_key.clone());
765                        keys.extend(right_node.keys.iter().cloned());
766                        right_node.keys = keys;
767
768                        let mut children = right_node.children.new_empty();
769                        children.extend(&left_node.children.split_off(split));
770                        children.extend(&right_node.children);
771                        right_node.children = children;
772                        debug_assert!(left_node.keys.len() + 1 == left_node.children.len());
773                        debug_assert!(right_node.keys.len() + 1 == right_node.children.len());
774                    }
775                    // Update the split key to reflect the new division between the nodes.
776                    self.keys[left] = split_key;
777                }
778            }
779        }
780    }
781
782    /// Remove the entry indicated by `cursor`.
783    fn remove(&mut self, mut cursor: Cursor) -> RemoveResult<K, V> {
784        let index = cursor.pop().expect("valid cursor");
785        let result = match &mut self.children {
786            ChildList::Leaf(children) => Arc::make_mut(&mut children[index]).remove(cursor),
787            ChildList::Internal(children) => Arc::make_mut(&mut children[index]).remove(cursor),
788        };
789        match result {
790            RemoveResult::NotFound => RemoveResult::NotFound,
791            RemoveResult::Removed(RemoveState { mut maybe_split_key, removed_value }) => {
792                if let Some(key) = maybe_split_key {
793                    if index > 0 {
794                        self.keys[index - 1] = key;
795                        maybe_split_key = None;
796                    }
797                }
798                RemoveResult::Removed(RemoveState { maybe_split_key, removed_value })
799            }
800            RemoveResult::Underflow(RemoveState { mut maybe_split_key, removed_value }) => {
801                if let Some(key) = maybe_split_key {
802                    if index > 0 {
803                        self.keys[index - 1] = key;
804                        maybe_split_key = None;
805                    }
806                }
807                self.rebalance_child(index);
808                if self.children.len() < NODE_CAPACITY / 2 {
809                    RemoveResult::Underflow(RemoveState { maybe_split_key, removed_value })
810                } else {
811                    RemoveResult::Removed(RemoveState { maybe_split_key, removed_value })
812                }
813            }
814        }
815    }
816}
817
818/// A node in the btree.
819#[derive(Clone, Debug)]
820enum Node<K: Ord + Copy, V: Clone> {
821    /// An internal node.
822    Internal(Arc<NodeInternal<K, V>>),
823
824    /// A leaf node.
825    Leaf(Arc<NodeLeaf<K, V>>),
826}
827
828impl<K, V> Node<K, V>
829where
830    K: Ord + Copy,
831    V: Clone,
832{
833    /// The number of children stored at this node.
834    fn len(&self) -> usize {
835        match self {
836            Node::Internal(node) => node.children.len(),
837            Node::Leaf(node) => node.keys.len(),
838        }
839    }
840
841    /// Search this node for the given key and return both the key and the value found.
842    fn get_key_value(&self, cursor: Cursor) -> Option<(&Range<K>, &V)> {
843        match self {
844            Node::Leaf(node) => node.get_key_value(cursor),
845            Node::Internal(node) => node.get_key_value(cursor),
846        }
847    }
848
849    /// The last key/value pair stored in this node.
850    fn last_key_value(&self) -> Option<(&Range<K>, &V)> {
851        match self {
852            Node::Leaf(node) => node.last_key_value(),
853            Node::Internal(node) => node.last_key_value(),
854        }
855    }
856
857    /// Converts a reference into a Node into a NodeRef.
858    fn as_ref(&self) -> NodeRef<'_, K, V> {
859        match self {
860            Node::Internal(node) => NodeRef::Internal(node),
861            Node::Leaf(node) => NodeRef::Leaf(node),
862        }
863    }
864
865    /// Returns a reference to the node that contains the entry indicated by the cursor.
866    ///
867    /// Assumes the cursor is non-empty.
868    fn get_containing_node(&self, cursor: Cursor) -> NodeRef<'_, K, V> {
869        assert!(cursor.depth > 0);
870        if cursor.depth == 1 {
871            return self.as_ref();
872        }
873        match self {
874            Node::Internal(node) => node.get_containing_node(cursor),
875            Node::Leaf(_) => unreachable!("leaf nodes do not have children"),
876        }
877    }
878
879    /// Insert the given value at the location indicated by `cursor`.
880    ///
881    /// If the insertion causes this node to split, the node will always split into two instances
882    /// of the same type of node.
883    fn insert(&mut self, cursor: Cursor, range: Range<K>, value: V) -> InsertResult<K, V> {
884        match self {
885            Node::Internal(node) => Arc::make_mut(node).insert(cursor, range, value),
886            Node::Leaf(node) => Arc::make_mut(node).insert(cursor, range, value),
887        }
888    }
889
890    /// Remove the entry indicated by `cursor`.
891    fn remove(&mut self, cursor: Cursor) -> RemoveResult<K, V> {
892        match self {
893            Node::Internal(node) => Arc::make_mut(node).remove(cursor),
894            Node::Leaf(node) => Arc::make_mut(node).remove(cursor),
895        }
896    }
897
898    /// Find the given key in this node.
899    ///
900    /// Updates `cursor` to point to the position indicated by `position`.
901    fn find(&self, key: &K, position: CursorPosition, cursor: &mut Cursor) {
902        match self {
903            Node::Internal(node) => node.find(key, position, cursor),
904            Node::Leaf(node) => node.find(key, position, cursor),
905        }
906    }
907}
908
909/// A node in the btree.
910#[derive(Clone, Debug)]
911enum NodeRef<'a, K: Ord + Copy, V: Clone> {
912    /// An internal node.
913    Internal(&'a Arc<NodeInternal<K, V>>),
914
915    /// A leaf node.
916    Leaf(&'a Arc<NodeLeaf<K, V>>),
917}
918
919impl<'a, K, V> NodeRef<'a, K, V>
920where
921    K: Ord + Copy,
922    V: Clone,
923{
924    /// The number of children stored at this node.
925    fn len(&self) -> usize {
926        match self {
927            NodeRef::Internal(node) => node.children.len(),
928            NodeRef::Leaf(node) => node.keys.len(),
929        }
930    }
931}
932
933/// An iterator over the key-value pairs stored in a RangeMap2.
934#[derive(Debug)]
935pub struct Iter<'a, K: Ord + Copy, V: Clone> {
936    /// The state of the forward iteration.
937    ///
938    /// The cursor points to the next entry to enumerate.
939    forward: Cursor,
940
941    /// The state of the backward iteration.
942    ///
943    /// The cursor points to the child that was most recently iterated or just past the end of the
944    /// entry list if no entries have been enumerated from this leaf yet.
945    backward: Cursor,
946
947    /// The root node of the tree.
948    root: &'a Node<K, V>,
949}
950
951impl<'a, K, V> Iter<'a, K, V>
952where
953    K: Ord + Copy,
954    V: Clone,
955{
956    /// Whether the iterator is complete.
957    ///
958    /// Iteration stops when the forward and backward cursors meet.
959    fn is_done(&self) -> bool {
960        self.forward == self.backward
961    }
962}
963
964impl<'a, K, V> Iterator for Iter<'a, K, V>
965where
966    K: Ord + Copy,
967    V: Clone,
968{
969    type Item = (&'a Range<K>, &'a V);
970
971    fn next(&mut self) -> Option<Self::Item> {
972        while !self.is_done() {
973            match self.root.get_containing_node(self.forward) {
974                NodeRef::Leaf(leaf) => {
975                    let index = self.forward.back();
976                    if index < leaf.keys.len() {
977                        let key = &leaf.keys[index];
978                        let value = &leaf.values[index];
979                        self.forward.increment_back();
980                        return Some((key, value));
981                    } else {
982                        self.forward.pop_back();
983                        self.forward.increment_back();
984                    }
985                }
986                NodeRef::Internal(internal) => {
987                    let index = self.forward.back();
988                    if index < internal.children.len() {
989                        self.forward.push_back(0);
990                    } else {
991                        self.forward.pop_back();
992                        self.forward.increment_back();
993                    }
994                }
995            }
996        }
997        None
998    }
999}
1000
1001impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V>
1002where
1003    K: Ord + Copy,
1004    V: Clone,
1005{
1006    fn next_back(&mut self) -> Option<Self::Item> {
1007        while !self.is_done() {
1008            match self.root.get_containing_node(self.backward) {
1009                NodeRef::Leaf(leaf) => {
1010                    let index = self.backward.back();
1011                    if index > 0 {
1012                        let key = &leaf.keys[index - 1];
1013                        let value = &leaf.values[index - 1];
1014                        self.backward.decrement_back();
1015                        return Some((key, value));
1016                    } else {
1017                        self.backward.pop_back();
1018                    }
1019                }
1020                NodeRef::Internal(internal) => {
1021                    let index = self.backward.back();
1022                    if index > 0 {
1023                        let child = internal.children.get_ref(index - 1);
1024                        self.backward.decrement_back();
1025                        self.backward.push_back(child.len());
1026                    } else {
1027                        self.backward.pop_back();
1028                    }
1029                }
1030            }
1031        }
1032        None
1033    }
1034}
1035
1036/// A map from ranges to values.
1037///
1038/// This map can be cloned efficiently. If the map is modified after being cloned, the relevant
1039/// parts of the map's internal structure will be copied lazily.
1040#[derive(Clone, Debug)]
1041pub struct RangeMap2<K: Ord + Copy, V: Clone + Eq> {
1042    /// The root node of the tree.
1043    ///
1044    /// The root node is either a leaf of an internal node, depending on the number of entries in
1045    /// the map.
1046    node: Node<K, V>,
1047}
1048
1049impl<K, V> Default for RangeMap2<K, V>
1050where
1051    K: Ord + Copy,
1052    V: Clone + Eq,
1053{
1054    fn default() -> Self {
1055        Self { node: Node::Leaf(Arc::new(NodeLeaf::empty())) }
1056    }
1057}
1058
1059impl<K, V> RangeMap2<K, V>
1060where
1061    K: Ord + Copy,
1062    V: Clone + Eq,
1063{
1064    /// Whether this map contains any entries.
1065    pub fn is_empty(&self) -> bool {
1066        match &self.node {
1067            Node::Leaf(node) => node.keys.is_empty(),
1068            Node::Internal(_) => false,
1069        }
1070    }
1071
1072    /// Find the given key in this node.
1073    ///
1074    /// Returns a Cursor that points to the position indicated by `position`.
1075    fn find(&self, key: &K, position: CursorPosition) -> Cursor {
1076        let mut cursor = Cursor::default();
1077        self.node.find(key, position, &mut cursor);
1078        cursor
1079    }
1080
1081    /// If the entry indicated by the cursor contains `key`, returns the range and value stored at
1082    /// that entry.
1083    fn get_if_contains_key(&self, key: &K, cursor: Cursor) -> Option<(&Range<K>, &V)> {
1084        if let Some((range, value)) = self.node.get_key_value(cursor) {
1085            if range.contains(key) {
1086                return Some((range, value));
1087            }
1088        }
1089        None
1090    }
1091
1092    /// Searches the map for a range that contains the given key.
1093    ///
1094    /// Returns the range and value if such a range is found.
1095    pub fn get(&self, key: K) -> Option<(&Range<K>, &V)> {
1096        self.get_if_contains_key(&key, self.find(&key, CursorPosition::Left))
1097    }
1098
1099    /// The last range stored in this map.
1100    pub fn last_range(&self) -> Option<&Range<K>> {
1101        self.node.last_key_value().map(|(key, _)| key)
1102    }
1103
1104    /// Searches the map for a range that contains the given key.
1105    ///
1106    /// If such a range is found, returns a cursor to that entry, the range, and the value.
1107    fn get_cursor_key_value(&mut self, key: &K) -> Option<(Cursor, Range<K>, V)> {
1108        let cursor = self.find(key, CursorPosition::Left);
1109        self.get_if_contains_key(key, cursor)
1110            .map(|(range, value)| (cursor, range.clone(), value.clone()))
1111    }
1112
1113    /// Remove the entry with the given key from the map.
1114    ///
1115    /// If the key was present in the map, returns the value previously stored at the given key.
1116    pub fn remove(&mut self, range: Range<K>) -> Vec<V> {
1117        let mut removed_values = vec![];
1118
1119        if range.end <= range.start {
1120            return removed_values;
1121        }
1122
1123        if let Some((cursor, old_range, v)) = self.get_cursor_key_value(&range.start) {
1124            // Remove that range from the map.
1125            removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1126
1127            // If the removed range extends after the end of the given range,
1128            // re-insert the part of the old range that extends beyond the end
1129            // of the given range.
1130            if old_range.end > range.end {
1131                self.insert_range_internal(range.end..old_range.end, v.clone());
1132            }
1133
1134            // If the removed range extends before the start of the given
1135            // range, re-insert the part of the old range that extends before
1136            // the start of the given range.
1137            if old_range.start < range.start {
1138                self.insert_range_internal(old_range.start..range.start, v);
1139            }
1140
1141            // Notice that we can end up splitting the old range into two
1142            // separate ranges if the old range extends both beyond the given
1143            // range and before the given range.
1144        }
1145
1146        if let Some((cursor, old_range, v)) = self.get_cursor_key_value(&range.end) {
1147            // If the old range starts before the removed range, we need to trim the old range.
1148            // TODO: Optimize with replace once available.
1149            if old_range.start < range.end {
1150                // Remove that range from the map.
1151                removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1152
1153                // If the removed range extends after the end of the given range,
1154                // re-insert the part of the old range that extends beyond the end
1155                // of the given range.
1156                if old_range.end > range.end {
1157                    self.insert_range_internal(range.end..old_range.end, v);
1158                }
1159            }
1160        }
1161
1162        let doomed = self.range(range.start..range.end).map(|(r, _)| r.start).collect::<Vec<_>>();
1163        for key in doomed {
1164            let cursor = self.find(&key, CursorPosition::Left);
1165            removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1166        }
1167
1168        removed_values
1169    }
1170
1171    /// Insert the given range and value at the location indicated by the cursor.
1172    fn insert_at(&mut self, cursor: Cursor, range: Range<K>, value: V) -> Option<V> {
1173        let result = self.node.insert(cursor, range, value);
1174        match result {
1175            InsertResult::Inserted => None,
1176            InsertResult::SplitLeaf(key, right) => {
1177                let mut keys = ArrayVec::new();
1178                let mut children = ArrayVec::new();
1179                keys.push(key);
1180                let Node::Leaf(left) = self.node.clone() else {
1181                    unreachable!("non-leaf node split into leaf node");
1182                };
1183                children.push(left);
1184                children.push(right);
1185                self.node = Node::Internal(Arc::new(NodeInternal {
1186                    keys,
1187                    children: ChildList::Leaf(children),
1188                }));
1189                None
1190            }
1191            InsertResult::SplitInternal(key, right) => {
1192                let mut keys = ArrayVec::new();
1193                let mut children = ArrayVec::new();
1194                keys.push(key);
1195                let Node::Internal(left) = self.node.clone() else {
1196                    unreachable!("non-internal node split into internal node");
1197                };
1198                children.push(left);
1199                children.push(right);
1200                self.node = Node::Internal(Arc::new(NodeInternal {
1201                    keys,
1202                    children: ChildList::Internal(children),
1203                }));
1204                None
1205            }
1206        }
1207    }
1208
1209    /// Insert the given range and value.
1210    ///
1211    /// Assumes the range is empty and that adjacent ranges have different values.
1212    fn insert_range_internal(&mut self, range: Range<K>, value: V) -> Option<V> {
1213        assert!(range.end > range.start);
1214        let cursor = self.find(&range.start, CursorPosition::Left);
1215        self.insert_at(cursor, range, value)
1216    }
1217
1218    /// Inserts a range with the given value.
1219    ///
1220    /// The keys included in the given range are now associated with the given
1221    /// value. If those keys were previously associated with another value,
1222    /// are no longer associated with that previous value.
1223    ///
1224    /// This method can cause one or more values in the map to be dropped if
1225    /// the all of the keys associated with those values are contained within
1226    /// the given range.
1227    ///
1228    /// If the inserted range is directly adjacent to another range with an equal value, the
1229    /// inserted range will be merged with the adjacent ranges.
1230    pub fn insert(&mut self, mut range: Range<K>, value: V) {
1231        if range.end <= range.start {
1232            return;
1233        }
1234        self.remove(range.clone());
1235
1236        // Check for a range directly before this one. If it exists, it will be the last range with
1237        // start < range.start.
1238        if let Some((prev_range, prev_value)) = self.range(..range.start).next_back() {
1239            if prev_range.end == range.start && value == *prev_value {
1240                let cursor = self.find(&prev_range.start, CursorPosition::Left);
1241                range.start = prev_range.start;
1242                self.remove_at(cursor);
1243            }
1244        }
1245
1246        // Check for a range directly after. If it exists, we can look it up by exact start value
1247        // of range.end.
1248        if let Some((cursor, next_range, next_value)) = self.get_cursor_key_value(&range.end) {
1249            if next_range.start == range.end && value == next_value {
1250                range.end = next_range.end;
1251                self.remove_at(cursor);
1252            }
1253        }
1254
1255        self.insert_range_internal(range, value);
1256    }
1257
1258    /// Remove the entry with the given cursor from the map.
1259    fn remove_at(&mut self, cursor: Cursor) -> Option<V> {
1260        let result = self.node.remove(cursor);
1261        match result {
1262            RemoveResult::NotFound => None,
1263            RemoveResult::Removed(RemoveState { removed_value, .. }) => Some(removed_value),
1264            RemoveResult::Underflow(RemoveState { removed_value, .. }) => {
1265                match &mut self.node {
1266                    Node::Leaf(_) => {
1267                        // Nothing we can do about an underflow of a single leaf node at the root.
1268                    }
1269                    Node::Internal(node) => {
1270                        // If the root has underflown to a trivial list, we can shrink the tree.
1271                        if node.children.len() == 1 {
1272                            self.node = node.children.get(0);
1273                        }
1274                    }
1275                }
1276                Some(removed_value)
1277            }
1278        }
1279    }
1280
1281    /// Iterate through the keys and values stored in the map.
1282    pub fn iter(&self) -> Iter<'_, K, V> {
1283        Iter {
1284            forward: Cursor::with_index(0),
1285            backward: Cursor::with_index(self.node.len()),
1286            root: &self.node,
1287        }
1288    }
1289
1290    /// Create the cursor stack for the start bound of the given range.
1291    fn find_start_bound(&self, bounds: &impl RangeBounds<K>) -> Cursor {
1292        let key = match bounds.start_bound() {
1293            Bound::Included(key) => key,
1294            Bound::Excluded(key) => key,
1295            Bound::Unbounded => {
1296                return Cursor::with_index(0);
1297            }
1298        };
1299        self.find(key, CursorPosition::Left)
1300    }
1301
1302    /// Create the cursor stack for the end bound of the given range.
1303    fn find_end_bound(&self, bounds: &impl RangeBounds<K>) -> Cursor {
1304        let key = match bounds.end_bound() {
1305            Bound::Included(key) => key,
1306            Bound::Excluded(key) => key,
1307            Bound::Unbounded => {
1308                return Cursor::with_index(self.node.len());
1309            }
1310        };
1311        self.find(key, CursorPosition::Right)
1312    }
1313
1314    /// Iterate through the keys and values stored in the given range in the map.
1315    fn range(&self, bounds: impl RangeBounds<K>) -> Iter<'_, K, V> {
1316        let forward = self.find_start_bound(&bounds);
1317        let backward = self.find_end_bound(&bounds);
1318        Iter { forward, backward, root: &self.node }
1319    }
1320
1321    /// Iterate over the ranges in the map, starting at the first range starting after or at the given point.
1322    pub fn iter_starting_at(&self, key: K) -> impl Iterator<Item = (&Range<K>, &V)> {
1323        self.range(key..).filter(move |(range, _)| key <= range.start)
1324    }
1325
1326    /// Iterate over the ranges in the map, starting at the last range starting before or at the given point.
1327    pub fn iter_ending_at(&self, key: K) -> impl DoubleEndedIterator<Item = (&Range<K>, &V)> {
1328        self.range(..key)
1329    }
1330
1331    pub fn intersection(
1332        &self,
1333        range: impl Borrow<Range<K>>,
1334    ) -> impl DoubleEndedIterator<Item = (&Range<K>, &V)> {
1335        let range = range.borrow();
1336        self.range(range.start..range.end)
1337    }
1338}
1339
1340#[cfg(test)]
1341mod test {
1342    use super::*;
1343
1344    #[::fuchsia::test]
1345    fn test_empty() {
1346        let mut map = RangeMap2::<u32, i32>::default();
1347
1348        assert!(map.get(12).is_none());
1349        map.remove(10..34);
1350        // This is a test to make sure we can handle reversed ranges
1351        #[allow(clippy::reversed_empty_ranges)]
1352        map.remove(34..10);
1353    }
1354
1355    #[::fuchsia::test]
1356    fn test_insert_into_empty() {
1357        let mut map = RangeMap2::<u32, i32>::default();
1358
1359        map.insert(10..34, -14);
1360
1361        assert_eq!((&(10..34), &-14), map.get(12).unwrap());
1362        assert_eq!((&(10..34), &-14), map.get(10).unwrap());
1363        assert!(map.get(9).is_none());
1364        assert_eq!((&(10..34), &-14), map.get(33).unwrap());
1365        assert!(map.get(34).is_none());
1366    }
1367
1368    #[::fuchsia::test]
1369    fn test_iter() {
1370        let mut map = RangeMap2::<u32, i32>::default();
1371
1372        map.insert(10..34, -14);
1373        map.insert(74..92, -12);
1374
1375        let mut iter = map.iter();
1376
1377        assert_eq!(iter.next().expect("missing elem"), (&(10..34), &-14));
1378        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
1379
1380        assert!(iter.next().is_none());
1381
1382        let mut iter = map.iter_starting_at(10);
1383        assert_eq!(iter.next().expect("missing elem"), (&(10..34), &-14));
1384        let mut iter = map.iter_starting_at(11);
1385        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
1386        let mut iter = map.iter_starting_at(74);
1387        assert_eq!(iter.next().expect("missing elem"), (&(74..92), &-12));
1388        let mut iter = map.iter_starting_at(75);
1389        assert_eq!(iter.next(), None);
1390
1391        assert_eq!(map.iter_ending_at(9).collect::<Vec<_>>(), vec![]);
1392        assert_eq!(map.iter_ending_at(34).collect::<Vec<_>>(), vec![(&(10..34), &-14)]);
1393        assert_eq!(map.iter_ending_at(74).collect::<Vec<_>>(), vec![(&(10..34), &-14)]);
1394        assert_eq!(
1395            map.iter_ending_at(75).collect::<Vec<_>>(),
1396            vec![(&(10..34), &-14), (&(74..92), &-12)]
1397        );
1398        assert_eq!(
1399            map.iter_ending_at(91).collect::<Vec<_>>(),
1400            vec![(&(10..34), &-14), (&(74..92), &-12)]
1401        );
1402        assert_eq!(
1403            map.iter_ending_at(92).collect::<Vec<_>>(),
1404            vec![(&(10..34), &-14), (&(74..92), &-12)]
1405        );
1406    }
1407
1408    #[::fuchsia::test]
1409    fn test_remove_overlapping_edge() {
1410        let mut map = RangeMap2::<u32, i32>::default();
1411
1412        map.insert(10..34, -14);
1413
1414        map.remove(2..11);
1415        assert_eq!((&(11..34), &-14), map.get(11).unwrap());
1416
1417        map.remove(33..42);
1418        assert_eq!((&(11..33), &-14), map.get(12).unwrap());
1419    }
1420
1421    #[::fuchsia::test]
1422    fn test_remove_middle_splits_range() {
1423        let mut map = RangeMap2::<u32, i32>::default();
1424
1425        map.insert(10..34, -14);
1426        map.remove(15..18);
1427
1428        assert_eq!((&(10..15), &-14), map.get(12).unwrap());
1429        assert_eq!((&(18..34), &-14), map.get(20).unwrap());
1430    }
1431
1432    #[::fuchsia::test]
1433    fn test_remove_upper_half_of_split_range_leaves_lower_range() {
1434        let mut map = RangeMap2::<u32, i32>::default();
1435
1436        map.insert(10..34, -14);
1437        map.remove(15..18);
1438        map.insert(2..7, -21);
1439        map.remove(20..42);
1440
1441        assert_eq!((&(2..7), &-21), map.get(5).unwrap());
1442        assert_eq!((&(10..15), &-14), map.get(12).unwrap());
1443    }
1444
1445    #[::fuchsia::test]
1446    fn test_range_map_overlapping_insert() {
1447        let mut map = RangeMap2::<u32, i32>::default();
1448
1449        map.insert(2..7, -21);
1450        map.insert(5..9, -42);
1451        map.insert(1..3, -43);
1452        map.insert(6..8, -44);
1453
1454        assert_eq!((&(1..3), &-43), map.get(2).unwrap());
1455        assert_eq!((&(3..5), &-21), map.get(4).unwrap());
1456        assert_eq!((&(5..6), &-42), map.get(5).unwrap());
1457        assert_eq!((&(6..8), &-44), map.get(7).unwrap());
1458    }
1459
1460    #[::fuchsia::test]
1461    fn test_intersect_single() {
1462        let mut map = RangeMap2::<u32, i32>::default();
1463
1464        map.insert(2..7, -10);
1465
1466        let mut iter = map.intersection(3..4);
1467        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1468        assert_eq!(iter.next(), None);
1469
1470        let mut iter = map.intersection(2..3);
1471        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1472        assert_eq!(iter.next(), None);
1473
1474        let mut iter = map.intersection(1..4);
1475        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1476        assert_eq!(iter.next(), None);
1477
1478        let mut iter = map.intersection(1..2);
1479        assert_eq!(iter.next(), None);
1480
1481        let mut iter = map.intersection(6..7);
1482        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1483        assert_eq!(iter.next(), None);
1484    }
1485
1486    #[::fuchsia::test]
1487    fn test_intersect_multiple() {
1488        let mut map = RangeMap2::<u32, i32>::default();
1489
1490        map.insert(2..7, -10);
1491        map.insert(7..9, -20);
1492        map.insert(10..11, -30);
1493
1494        let mut iter = map.intersection(3..8);
1495        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1496        assert_eq!(iter.next(), Some((&(7..9), &-20)));
1497        assert_eq!(iter.next(), None);
1498
1499        let mut iter = map.intersection(3..11);
1500        assert_eq!(iter.next(), Some((&(2..7), &-10)));
1501        assert_eq!(iter.next(), Some((&(7..9), &-20)));
1502        assert_eq!(iter.next(), Some((&(10..11), &-30)));
1503        assert_eq!(iter.next(), None);
1504    }
1505
1506    #[::fuchsia::test]
1507    fn test_intersect_no_gaps() {
1508        let mut map = RangeMap2::<u32, i32>::default();
1509
1510        map.insert(0..1, -10);
1511        map.insert(1..2, -20);
1512        map.insert(2..3, -30);
1513
1514        let mut iter = map.intersection(0..3);
1515        assert_eq!(iter.next(), Some((&(0..1), &-10)));
1516        assert_eq!(iter.next(), Some((&(1..2), &-20)));
1517        assert_eq!(iter.next(), Some((&(2..3), &-30)));
1518        assert_eq!(iter.next(), None);
1519    }
1520
1521    #[test]
1522    fn test_merging() {
1523        let mut map = RangeMap2::<u32, i32>::default();
1524
1525        map.insert(1..2, -10);
1526        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..2), &-10)]);
1527        map.insert(3..4, -10);
1528        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..2), &-10), (&(3..4), &-10)]);
1529        map.insert(2..3, -10);
1530        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(1..4), &-10)]);
1531        map.insert(0..1, -10);
1532        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(0..4), &-10)]);
1533        map.insert(4..5, -10);
1534        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&(0..5), &-10)]);
1535        map.insert(2..3, -20);
1536        assert_eq!(
1537            map.iter().collect::<Vec<_>>(),
1538            vec![(&(0..2), &-10), (&(2..3), &-20), (&(3..5), &-10)]
1539        );
1540    }
1541
1542    #[test]
1543    fn test_remove_multiple_ranges_exact_query() {
1544        let first = 15..21;
1545        let second = first.end..29;
1546
1547        let mut map = RangeMap2::default();
1548        map.insert(first.clone(), 1);
1549        map.insert(second.clone(), 2);
1550
1551        assert_eq!(map.remove(first.start..second.end), &[1, 2]);
1552    }
1553
1554    #[::fuchsia::test]
1555    fn test_large_insert_and_remove() {
1556        let mut map = RangeMap2::<u32, i32>::default();
1557        let num_entries = 1000;
1558
1559        // Insert a large number of entries
1560        for i in 0..num_entries {
1561            let start = i as u32 * 10;
1562            let end = start + 5;
1563            let value = i as i32;
1564            map.insert(start..end, value);
1565        }
1566
1567        // Verify that all inserted entries can be retrieved
1568        for i in 0..num_entries {
1569            let start = i as u32 * 10;
1570            let end = start + 5;
1571            let point = start + 2;
1572            if let Some((range, value)) = map.get(point) {
1573                assert!(range.start <= point && point < range.end);
1574                assert_eq!(*range, start..end);
1575                assert_eq!(*value, i as i32);
1576            } else {
1577                panic!("Expected to find a range for point {}", point);
1578            }
1579        }
1580
1581        // Remove a large number of entries
1582        for i in 0..num_entries {
1583            let start = i as u32 * 10;
1584            let end = start + 5;
1585            map.remove(start..end);
1586        }
1587
1588        // Verify that the map is empty after removing all entries
1589        assert!(map.is_empty());
1590    }
1591
1592    #[::fuchsia::test]
1593    fn test_large_insert_and_remove_overlapping() {
1594        let mut map = RangeMap2::<u32, i32>::default();
1595        let num_entries = 1000;
1596
1597        // Insert a large number of entries with overlapping ranges
1598        for i in 0..num_entries {
1599            let start = i as u32 * 5;
1600            let end = start + 20;
1601            let value = i as i32;
1602            map.insert(start..end, value);
1603        }
1604
1605        // Verify that all inserted entries can be retrieved
1606        for i in 0..num_entries {
1607            let point = i as u32 * 5 + 1;
1608            if let Some((range, value)) = map.get(point) {
1609                assert!(range.start <= point && point < range.end);
1610                assert_eq!(*value, i as i32);
1611            } else {
1612                panic!("Expected to find a range for point {}", point);
1613            }
1614        }
1615
1616        // Remove a large number of entries with overlapping ranges
1617        for i in 0..num_entries {
1618            let start = i as u32 * 5;
1619            let end = start + 20;
1620            map.remove(start..end);
1621        }
1622
1623        // Verify that the map is empty after removing all entries
1624        assert!(map.is_empty());
1625    }
1626
1627    #[::fuchsia::test]
1628    fn test_large_insert_and_get_specific_points() {
1629        let mut map = RangeMap2::<u32, i32>::default();
1630        let num_entries = 1000;
1631        let mut inserted_ranges = Vec::new();
1632
1633        // Insert a large number of entries
1634        for i in 0..num_entries {
1635            let start = i as u32 * 10;
1636            let end = start + 5;
1637            let value = i as i32;
1638            map.insert(start..end, value);
1639            inserted_ranges.push((start..end, value));
1640        }
1641
1642        // Verify that specific points can be retrieved correctly
1643        for (range, value) in &inserted_ranges {
1644            let point = range.start + 2;
1645            let (retrieved_range, retrieved_value) = map.get(point).unwrap();
1646            assert_eq!(retrieved_range, range);
1647            assert_eq!(retrieved_value, value);
1648        }
1649    }
1650
1651    #[::fuchsia::test]
1652    fn test_large_insert_and_iter() {
1653        let mut map = RangeMap2::<u32, i32>::default();
1654        let num_entries = 1000;
1655        let mut inserted_ranges = Vec::new();
1656
1657        // Insert a large number of entries
1658        for i in 0..num_entries {
1659            let start = i as u32 * 10;
1660            let end = start + 5;
1661            let value = i as i32;
1662            map.insert(start..end, value);
1663            inserted_ranges.push((start..end, value));
1664        }
1665
1666        // Verify that iter() returns all inserted entries
1667        let mut iter_ranges: Vec<(&Range<u32>, &i32)> = map.iter().collect();
1668        iter_ranges.sort_by_key(|(range, _)| range.start);
1669        let mut inserted_ranges_sorted: Vec<(Range<u32>, i32)> = inserted_ranges.clone();
1670        inserted_ranges_sorted.sort_by_key(|(range, _)| range.start);
1671
1672        assert_eq!(iter_ranges.len(), inserted_ranges_sorted.len());
1673        for (i, (range, value)) in iter_ranges.iter().enumerate() {
1674            assert_eq!(*range, &inserted_ranges_sorted[i].0);
1675            assert_eq!(*value, &inserted_ranges_sorted[i].1);
1676        }
1677    }
1678
1679    #[::fuchsia::test]
1680    fn test_large_insert_and_iter_starting_at() {
1681        let mut map = RangeMap2::<u32, i32>::default();
1682        let num_entries = 1000;
1683
1684        // Insert a large number of entries
1685        for i in 0..num_entries {
1686            let start = i as u32 * 10;
1687            let end = start + 5;
1688            let value = i as i32;
1689            map.insert(start..end, value);
1690        }
1691
1692        // Verify iter_starting_at()
1693        let start_point = 5000;
1694        let mut iter = map.iter_starting_at(start_point);
1695        while let Some((range, _)) = iter.next() {
1696            assert!(range.start >= start_point);
1697        }
1698    }
1699
1700    #[::fuchsia::test]
1701    fn test_large_insert_and_iter_ending_at() {
1702        let mut map = RangeMap2::<u32, i32>::default();
1703        let num_entries = 1000;
1704
1705        // Insert a large number of entries
1706        for i in 0..num_entries {
1707            let start = i as u32 * 10;
1708            let end = start + 5;
1709            let value = i as i32;
1710            map.insert(start..end, value);
1711        }
1712
1713        // Verify iter_ending_at()
1714        let end_point = 5000;
1715        let mut iter = map.iter_ending_at(end_point);
1716        while let Some((range, _)) = iter.next() {
1717            assert!(range.start < end_point);
1718        }
1719    }
1720
1721    #[::fuchsia::test]
1722    fn test_large_insert_and_intersection() {
1723        let mut map = RangeMap2::<u32, i32>::default();
1724        let num_entries = 1000;
1725
1726        // Insert a large number of entries
1727        for i in 0..num_entries {
1728            let start = i as u32 * 10;
1729            let end = start + 5;
1730            let value = i as i32;
1731            map.insert(start..end, value);
1732        }
1733
1734        // Verify intersection()
1735        let intersect_start = 4000;
1736        let intersect_end = 4050;
1737        let mut iter = map.intersection(intersect_start..intersect_end);
1738        while let Some((range, _)) = iter.next() {
1739            assert!((range.start < intersect_end && range.end > intersect_start));
1740        }
1741    }
1742
1743    #[::fuchsia::test]
1744    fn test_large_insert_and_last_range() {
1745        let mut map = RangeMap2::<u32, i32>::default();
1746        let num_entries = 1000;
1747        let mut last_range = None;
1748
1749        // Insert a large number of entries
1750        for i in 0..num_entries {
1751            let start = i as u32 * 10;
1752            let end = start + 5;
1753            let value = i as i32;
1754            map.insert(start..end, value);
1755            last_range = Some(start..end);
1756        }
1757
1758        // Verify last_range()
1759        if let Some(expected_last_range) = last_range {
1760            assert_eq!(map.last_range().unwrap(), &expected_last_range);
1761        }
1762    }
1763
1764    #[::fuchsia::test]
1765    fn test_large_insert_and_remove_alternating() {
1766        let mut map = RangeMap2::<u32, i32>::default();
1767        let num_entries = 1000;
1768
1769        // Insert and remove entries in an alternating pattern
1770        for i in 0..num_entries {
1771            let start = i as u32 * 10;
1772            let end = start + 5;
1773            let value = i as i32;
1774            map.insert(start..end, value);
1775
1776            if i % 2 == 0 {
1777                // Remove every other entry
1778                map.remove(start..end);
1779            }
1780        }
1781
1782        // Verify that only the non-removed entries remain
1783        for i in 0..num_entries {
1784            let start = i as u32 * 10;
1785            let end = start + 5;
1786            let point = start + 2;
1787            if i % 2 != 0 {
1788                if let Some((range, value)) = map.get(point) {
1789                    assert!(range.start <= point && point < range.end);
1790                    assert_eq!(*range, start..end);
1791                    assert_eq!(*value, i as i32);
1792                } else {
1793                    panic!("Expected to find a range for point {}", point);
1794                }
1795            } else {
1796                assert!(map.get(point).is_none());
1797            }
1798        }
1799    }
1800
1801    #[::fuchsia::test]
1802    fn test_large_insert_and_remove_multiple_times() {
1803        let mut map = RangeMap2::<u32, i32>::default();
1804        let num_entries = 200;
1805        let num_iterations = 5;
1806
1807        for _ in 0..num_iterations {
1808            // Insert a large number of entries
1809            for i in 0..num_entries {
1810                let start = i as u32 * 10;
1811                let end = start + 5;
1812                let value = i as i32;
1813                map.insert(start..end, value);
1814            }
1815
1816            // Remove a large number of entries
1817            for i in 0..num_entries {
1818                let start = i as u32 * 10;
1819                let end = start + 5;
1820                map.remove(start..end);
1821            }
1822        }
1823
1824        // Verify that the map is empty after multiple insert/remove cycles
1825        assert!(map.is_empty());
1826    }
1827
1828    #[::fuchsia::test]
1829    fn test_merging_ranges_with_equal_values() {
1830        let mut map = RangeMap2::<u32, i32>::default();
1831
1832        // Insert some initial ranges
1833        map.insert(10..20, 100);
1834        map.insert(30..40, 200);
1835        map.insert(50..60, 100);
1836
1837        // Merge adjacent ranges with equal values
1838        map.insert(20..30, 100);
1839        assert_eq!(map.get(15).unwrap(), (&(10..30), &100));
1840        assert_eq!(map.get(35).unwrap(), (&(30..40), &200));
1841        assert_eq!(map.get(55).unwrap(), (&(50..60), &100));
1842
1843        // Merge non-adjacent ranges with equal values
1844        map.insert(40..50, 100);
1845        assert_eq!(map.get(15).unwrap(), (&(10..30), &100));
1846        assert_eq!(map.get(35).unwrap(), (&(30..40), &200));
1847        assert_eq!(map.get(45).unwrap(), (&(40..60), &100));
1848
1849        // Insert a range with a different value
1850        map.insert(60..70, 300);
1851        assert_eq!(map.get(65).unwrap(), (&(60..70), &300));
1852
1853        // Merge a range with a different value
1854        map.insert(70..80, 300);
1855        assert_eq!(map.get(65).unwrap(), (&(60..80), &300));
1856
1857        // Merge a range with a different value at the beginning
1858        map.insert(0..10, 400);
1859        assert_eq!(map.get(5).unwrap(), (&(0..10), &400));
1860
1861        // Merge a range with a different value at the end
1862        map.insert(80..90, 400);
1863        assert_eq!(map.get(85).unwrap(), (&(80..90), &400));
1864
1865        // Merge a range with a different value in the middle
1866        map.insert(90..100, 400);
1867        assert_eq!(map.get(95).unwrap(), (&(80..100), &400));
1868        map.insert(100..110, 400);
1869        assert_eq!(map.get(95).unwrap(), (&(80..110), &400));
1870        map.insert(110..120, 400);
1871        assert_eq!(map.get(95).unwrap(), (&(80..120), &400));
1872
1873        // Merge a range with a different value in the middle
1874        map.insert(10..90, 400);
1875        assert_eq!(map.get(5).unwrap(), (&(0..120), &400));
1876    }
1877
1878    #[::fuchsia::test]
1879    fn test_large_insert_and_remove_with_gaps() {
1880        let mut map = RangeMap2::<u32, i32>::default();
1881        let num_entries = 500;
1882
1883        // Insert entries with gaps
1884        for i in 0..num_entries {
1885            let start = i as u32 * 20;
1886            let end = start + 5;
1887            let value = i as i32;
1888            map.insert(start..end, value);
1889        }
1890
1891        // Remove entries with gaps
1892        for i in 0..num_entries {
1893            if i % 2 == 0 {
1894                let start = i as u32 * 20;
1895                let end = start + 5;
1896                map.remove(start..end);
1897            }
1898        }
1899
1900        // Verify the state of the map
1901        for i in 0..num_entries {
1902            let start = i as u32 * 20;
1903            let end = start + 5;
1904            let point = start + 2;
1905
1906            if i % 2 != 0 {
1907                if let Some((range, value)) = map.get(point) {
1908                    assert!(range.start <= point && point < range.end);
1909                    assert_eq!(*range, start..end);
1910                    assert_eq!(*value, i as i32);
1911                } else {
1912                    panic!("Expected to find a range for point {}", point);
1913                }
1914            } else {
1915                assert!(map.get(point).is_none());
1916            }
1917        }
1918    }
1919
1920    #[::fuchsia::test]
1921    fn test_critical_removal() {
1922        fn range_for(i: u32) -> Range<u32> {
1923            let start = i * 20;
1924            let end = start + 5;
1925            start..end
1926        }
1927
1928        // Index 0 cannot be the critical index.
1929        for critial_index in 1..100 {
1930            let mut map = RangeMap2::<u32, i32>::default();
1931
1932            // Insert enough entries for force a split somewhere.
1933            for i in 0..100 {
1934                let value = i as i32;
1935                map.insert(range_for(i), value);
1936            }
1937
1938            // We don't know whether the split will occur at the critical index, but we try all the
1939            // possible indices to ensure that we test an index at which a split occurred.
1940            let critical_range = range_for(critial_index);
1941            map.remove(critical_range.clone());
1942
1943            // We now insert a range that spans the critical point. This range will be inserted to
1944            // left of the split.
1945            let value = -10 as i32;
1946            let spanning_range = (critical_range.start - 2)..(critical_range.start + 2);
1947            map.insert(spanning_range.clone(), value);
1948
1949            // Check to see if we can find the range by looking before the critical point.
1950            let key_before_split = critical_range.start - 1;
1951            assert_eq!(map.get(key_before_split), Some((&spanning_range, &value)));
1952
1953            // Check to see if we can find the range by looking after the critical point.
1954            let key_after_split = critical_range.start + 1;
1955            assert_eq!(map.get(key_after_split), Some((&spanning_range, &value)));
1956        }
1957    }
1958}