1use 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
12const B: usize = 6;
16
17const NODE_CAPACITY: usize = 2 * B;
19
20pub trait Gap {
22 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#[derive(Debug, Default, Clone, Copy)]
34struct Cursor {
35 depth: u8,
37
38 indices: [u8; 7],
40}
41
42impl Cursor {
43 fn with_index(index: usize) -> Self {
45 let mut cursor = Self::default();
46 cursor.push(index);
47 cursor
48 }
49
50 fn is_empty(&self) -> bool {
55 self.depth == 0
56 }
57
58 fn push(&mut self, index: usize) {
62 self.indices[self.depth as usize] = index as u8;
63 self.depth += 1;
64 }
65
66 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 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 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 fn back(&self) -> usize {
107 self.indices[0] as usize
108 }
109
110 fn increment_back(&mut self) {
116 self.indices[0] += 1;
117 }
118
119 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
145enum CursorPosition {
147 Left,
151
152 Right,
156}
157
158fn 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 let range = &keys[mid];
170 if key < &range.start {
171 right = mid;
173 } else if key < &range.end {
174 return mid;
176 } else {
177 left = mid + 1;
179 }
180 }
181 left
184}
185
186#[derive(Clone)]
192struct NodeLeaf<K: Ord + Copy + Gap, V: Clone> {
193 max_gap: u64,
195
196 keys: ArrayVec<Range<K>, NODE_CAPACITY>,
202
203 values: ArrayVec<V, NODE_CAPACITY>,
205}
206
207impl<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
218enum InsertResult<K: Ord + Copy + Gap, V: Clone> {
220 Inserted,
222
223 SplitLeaf(K, Arc<NodeLeaf<K, V>>),
228
229 SplitInternal(K, Arc<NodeInternal<K, V>>),
234}
235
236struct RemoveState<K: Ord + Copy + Gap, V: Clone> {
238 maybe_split_key: Option<K>,
243
244 removed_value: V,
246}
247
248enum RemoveResult<K: Ord + Copy + Gap, V: Clone> {
253 NotFound,
255
256 Removed(RemoveState<K, V>),
259
260 Underflow(RemoveState<K, V>),
269}
270
271impl<K, V> NodeLeaf<K, V>
272where
273 K: Ord + Copy + Gap,
274 V: Clone,
275{
276 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 fn empty() -> Self {
287 Self { max_gap: 0, keys: ArrayVec::new(), values: ArrayVec::new() }
288 }
289
290 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 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 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 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 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 #[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 fn update_max_gap(&mut self) {
363 self.max_gap = self.compute_max_gap();
364 }
365
366 fn measure_gap(&self, other: &Self) -> u64 {
368 self.keys[self.keys.len() - 1].end.measure_gap(&other.keys[0].start)
372 }
373
374 fn key_range(&self) -> Range<K> {
376 self.keys[0].start..self.keys[self.keys.len() - 1].end
377 }
378
379 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 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 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#[derive(Clone, Debug)]
469enum ChildList<K: Ord + Copy + Gap, V: Clone> {
470 Leaf(ArrayVec<Arc<NodeLeaf<K, V>>, NODE_CAPACITY>),
472
473 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 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 fn len(&self) -> usize {
492 match self {
493 ChildList::Leaf(children) => children.len(),
494 ChildList::Internal(children) => children.len(),
495 }
496 }
497
498 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 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 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 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 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 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 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 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 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#[derive(Clone, Debug)]
605struct NodeInternal<K: Ord + Copy + Gap, V: Clone> {
606 max_gap: u64,
608
609 subtree_key_range: Range<K>,
611
612 keys: ArrayVec<K, NODE_CAPACITY>,
618
619 children: ChildList<K, V>,
621}
622
623fn 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 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 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 p + 1
665 } else {
666 p
668 }
669 }
670
671 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 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 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 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 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 #[cfg(test)]
746 fn validate_max_gap(&self) {
747 let computed = self.compute_max_gap();
748 assert_eq!(computed, self.max_gap);
749
750 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 #[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 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 fn measure_gap(&self, other: &Self) -> u64 {
807 self.subtree_key_range.end.measure_gap(&other.subtree_key_range.start)
808 }
809
810 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 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 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 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 fn rebalance_child(&mut self, index: usize) {
899 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 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 let split = n / 2;
924 let right_node = Arc::make_mut(right_shared_node);
925 if left_node.values.len() < split {
926 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 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 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 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 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 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 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 self.keys[left] = split_key;
997 }
998 }
999 }
1000 }
1001
1002 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 self.children.size_at(index) == 0 {
1031 debug_assert!(maybe_split_key.is_none());
1034 self.children.remove(index);
1035 if index == 0 {
1036 maybe_split_key = Some(self.keys.remove(0));
1039 } else {
1040 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 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#[derive(Clone, Debug)]
1116enum Node<K: Ord + Copy + Gap, V: Clone> {
1117 Internal(Arc<NodeInternal<K, V>>),
1119
1120 Leaf(Arc<NodeLeaf<K, V>>),
1122}
1123
1124impl<K, V> Node<K, V>
1125where
1126 K: Ord + Copy + Gap,
1127 V: Clone,
1128{
1129 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 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 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 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 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 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 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 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 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 #[cfg(test)]
1226 fn validate_keys(&self) {
1227 if let Node::Internal(node) = self {
1228 node.validate_keys();
1229 }
1230 }
1231}
1232
1233#[derive(Clone, Debug)]
1235enum NodeRef<'a, K: Ord + Copy + Gap, V: Clone> {
1236 Internal(&'a Arc<NodeInternal<K, V>>),
1238
1239 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 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#[derive(Debug)]
1259pub struct Iter<'a, K: Ord + Copy + Gap, V: Clone> {
1260 forward: Cursor,
1264
1265 backward: Cursor,
1270
1271 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 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#[derive(Clone, Debug)]
1365pub struct RangeMap<K: Ord + Copy + Gap, V: Clone + Eq> {
1366 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 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 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 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 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 pub fn get_keys(&self, range: Range<K>) -> impl Iterator<Item = &Range<K>> {
1427 self.range(range).map(|(key, _)| key)
1428 }
1429
1430 pub fn last_range(&self) -> Option<&Range<K>> {
1432 self.node.last_key_value().map(|(key, _)| key)
1433 }
1434
1435 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 pub fn find_gap_end(&self, gap: u64, upper_bound: &K) -> K {
1448 self.node.find_gap_end(gap, upper_bound)
1449 }
1450
1451 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 removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1464
1465 if old_range.end > range.end {
1469 self.insert_range_internal(range.end..old_range.end, v.clone());
1470 }
1471
1472 if old_range.start < range.start {
1476 self.insert_range_internal(old_range.start..range.start, v);
1477 }
1478
1479 }
1483
1484 if let Some((cursor, old_range, v)) = self.get_cursor_key_value(&range.end) {
1485 if old_range.start < range.end {
1488 removed_values.push(self.remove_at(cursor).expect("entry should exist"));
1490
1491 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 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 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 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 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 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 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 }
1616 Node::Internal(node) => {
1617 if node.children.len() == 1 {
1619 self.node = node.children.get(0);
1620 }
1621 }
1622 }
1623 Some(removed_value)
1624 }
1625 }
1626 }
1627
1628 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 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 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 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 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 #[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 assert!(map.is_empty());
1696
1697 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 map.remove(start..end);
2145 }
2146 }
2147
2148 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 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 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 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 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 assert_eq!(map.node.len(), NODE_CAPACITY);
2207
2208 {
2209 let mut map = map.clone();
2210
2211 let split_start = 15;
2213 let split_end = 20;
2214 map.insert(split_start..split_end, -1);
2215
2216 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 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 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 map.insert(10..20, 100);
2247 map.insert(30..40, 200);
2248 map.insert(50..60, 100);
2249
2250 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 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 map.insert(60..70, 300);
2264 assert_eq!(map.get(65).unwrap(), (&(60..70), &300));
2265
2266 map.insert(70..80, 300);
2268 assert_eq!(map.get(65).unwrap(), (&(60..80), &300));
2269
2270 map.insert(0..10, 400);
2272 assert_eq!(map.get(5).unwrap(), (&(0..10), &400));
2273
2274 map.insert(80..90, 400);
2276 assert_eq!(map.get(85).unwrap(), (&(80..90), &400));
2277
2278 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 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 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 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 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 for critial_index in 1..100 {
2343 let mut map = RangeMap::<u32, i32>::default();
2344
2345 for i in 0..100 {
2347 let value = i as i32;
2348 map.insert(range_for(i), value);
2349 }
2350
2351 let critical_range = range_for(critial_index);
2354 map.remove(critical_range.clone());
2355
2356 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 let key_before_split = critical_range.start - 1;
2364 assert_eq!(map.get(key_before_split), Some((&spanning_range, &value)));
2365
2366 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 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 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 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 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 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); assert_eq!(map.find_gap_end(10, &100), 100);
2492
2493 map.insert(40..50, 1);
2495
2496 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 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 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 let mut map = map.clone();
2523 map.remove(critical_range.clone());
2524 }
2525
2526 {
2527 let mut map = map.clone();
2528
2529 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 map.remove(critical_range);
2540 }
2541 }
2542}