1use super::CompactArc;
36use std::marker::PhantomData;
37use std::mem;
38use std::ops::Bound;
39use std::ptr;
40use std::sync::atomic::{AtomicU32, Ordering};
41
42const MAX_KEYS: usize = 128;
48
49const _: () = assert!(
51 MAX_KEYS <= 255,
52 "MAX_KEYS must be <= 255 (NodePath uses u8 indices)"
53);
54
55const MIN_KEYS: usize = (MAX_KEYS - 1) / 2;
58
59#[repr(C)]
66struct NodeHeader {
67 len: u16,
68 is_leaf: u8,
69 _pad1: u8,
70 drop_count: AtomicU32,
74}
75
76pub struct NodePtr<V: Clone> {
80 ptr: CompactArc<[u8]>,
81 _marker: PhantomData<V>,
82}
83
84impl<V: Clone> Clone for NodePtr<V> {
85 fn clone(&self) -> Self {
86 let new_ptr = self.ptr.clone();
89
90 let header = unsafe { &*(self.ptr.data_ptr_mut() as *const NodeHeader) };
93 header.drop_count.fetch_add(1, Ordering::Relaxed);
94
95 Self {
96 ptr: new_ptr,
97 _marker: PhantomData,
98 }
99 }
100}
101
102impl<V: Clone> Drop for NodePtr<V> {
103 fn drop(&mut self) {
104 let header = unsafe { &*(self.ptr.data_ptr_mut() as *const NodeHeader) };
112 let old_count = header.drop_count.fetch_sub(1, Ordering::AcqRel);
113
114 if old_count != 1 {
115 return;
117 }
118
119 let len = self.len();
121 if self.is_leaf() {
122 unsafe {
126 let v_ptr = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
127 for i in 0..len {
128 ptr::drop_in_place(v_ptr.add(i));
129 }
130 }
131 } else {
132 unsafe {
136 let c_ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
137 for i in 0..=len {
138 ptr::drop_in_place(c_ptr.add(i));
139 }
140 }
141 }
142 }
143}
144
145impl<V: Clone> NodePtr<V> {
146 const fn keys_offset() -> usize {
148 mem::size_of::<NodeHeader>()
149 }
150
151 const fn values_offset() -> usize {
156 Self::keys_offset() + ((MAX_KEYS + 1) * 8)
157 }
158
159 const fn children_offset() -> usize {
162 Self::keys_offset() + ((MAX_KEYS + 1) * 8)
163 }
164
165 fn new_leaf() -> Self {
166 let size = Self::values_offset() + ((MAX_KEYS + 1) * mem::size_of::<V>());
167 let vec = vec![0u8; size];
168
169 let mut ptr = NodePtr {
170 ptr: CompactArc::from_vec(vec),
171 _marker: PhantomData,
172 };
173
174 let header = ptr.header_mut();
175 header.len = 0;
176 header.is_leaf = 1;
177 header.drop_count = AtomicU32::new(1);
178
179 ptr
180 }
181
182 fn new_internal() -> Self {
183 let size = Self::children_offset() + ((MAX_KEYS + 2) * mem::size_of::<NodePtr<V>>());
184 let vec = vec![0u8; size];
185
186 let mut ptr = NodePtr {
187 ptr: CompactArc::from_vec(vec),
188 _marker: PhantomData,
189 };
190
191 let header = ptr.header_mut();
192 header.len = 0;
193 header.is_leaf = 0;
194 header.drop_count = AtomicU32::new(1);
195
196 ptr
197 }
198
199 fn make_mut(&mut self) -> &mut Self {
200 let header = unsafe { &*(self.ptr.data_ptr_mut() as *const NodeHeader) };
204 if header.drop_count.load(Ordering::Acquire) != 1 {
205 *self = self.deep_clone();
208 }
209 self
210 }
211
212 fn deep_clone(&self) -> Self {
216 let len = self.len();
217
218 if self.is_leaf() {
219 let mut new_node = NodePtr::new_leaf();
220 unsafe {
227 let k_src = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *const i64;
228 let k_dst = new_node.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
229 ptr::copy_nonoverlapping(k_src, k_dst, len);
230 }
231
232 unsafe {
237 let v_src = self.ptr.data_ptr_mut().add(Self::values_offset()) as *const V;
238 let v_dst = new_node.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
239 for i in 0..len {
240 ptr::write(v_dst.add(i), (*v_src.add(i)).clone());
241 new_node.set_len(i + 1);
244 }
245 }
246
247 new_node
248 } else {
249 let mut new_node = NodePtr::new_internal();
250 new_node.set_len(len);
251
252 unsafe {
256 let k_src = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *const i64;
257 let k_dst = new_node.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
258 ptr::copy_nonoverlapping(k_src, k_dst, len);
259 }
260
261 unsafe {
267 let c_src =
268 self.ptr.data_ptr_mut().add(Self::children_offset()) as *const NodePtr<V>;
269 let c_dst =
270 new_node.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
271 for i in 0..=len {
272 ptr::write(c_dst.add(i), (*c_src.add(i)).clone());
274 }
275 }
276
277 new_node
278 }
279 }
280
281 fn header(&self) -> &NodeHeader {
282 unsafe { &*(self.ptr.data_ptr_mut() as *const NodeHeader) }
285 }
286
287 fn header_mut(&mut self) -> &mut NodeHeader {
288 unsafe { &mut *(self.ptr.data_ptr_mut() as *mut NodeHeader) }
294 }
295
296 fn is_leaf(&self) -> bool {
297 self.header().is_leaf == 1
298 }
299
300 fn len(&self) -> usize {
301 self.header().len as usize
302 }
303
304 fn set_len(&mut self, len: usize) {
305 self.header_mut().len = len as u16;
306 }
307
308 fn keys(&self) -> &[i64] {
309 let len = self.len();
310 unsafe {
313 let ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *const i64;
314 std::slice::from_raw_parts(ptr, len)
315 }
316 }
317
318 fn values(&self) -> &[V] {
319 assert!(self.is_leaf());
320 let len = self.len();
321 unsafe {
324 let ptr = self.ptr.data_ptr_mut().add(Self::values_offset()) as *const V;
325 std::slice::from_raw_parts(ptr, len)
326 }
327 }
328
329 fn values_mut_slice(&mut self) -> &mut [V] {
330 assert!(self.is_leaf());
331 let len = self.len();
332 unsafe {
335 let ptr = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
336 std::slice::from_raw_parts_mut(ptr, len)
337 }
338 }
339
340 fn children(&self) -> &[NodePtr<V>] {
341 assert!(!self.is_leaf());
342 let len = self.len() + 1; unsafe {
346 let ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *const NodePtr<V>;
347 std::slice::from_raw_parts(ptr, len)
348 }
349 }
350
351 fn children_mut_slice(&mut self) -> &mut [NodePtr<V>] {
352 assert!(!self.is_leaf());
353 let len = self.len() + 1;
354 unsafe {
357 let ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
358 std::slice::from_raw_parts_mut(ptr, len)
359 }
360 }
361
362 fn child(&self, index: usize) -> &NodePtr<V> {
363 &self.children()[index]
364 }
365
366 fn child_mut(&mut self, index: usize) -> &mut NodePtr<V> {
367 &mut self.children_mut_slice()[index]
368 }
369
370 fn search(&self, key: i64) -> Result<usize, usize> {
371 self.keys().binary_search(&key)
372 }
373
374 fn push_leaf(&mut self, key: i64, value: V) {
375 assert!(self.is_leaf());
376 let len = self.len();
377 assert!(len <= MAX_KEYS); unsafe {
382 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
383 ptr::write(k_ptr.add(len), key);
384
385 let v_ptr = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
386 ptr::write(v_ptr.add(len), value);
387 }
388 self.set_len(len + 1);
389 }
390
391 fn remove_leaf(&mut self, index: usize) -> V {
392 assert!(self.is_leaf());
393 let len = self.len();
394 assert!(index < len);
395
396 unsafe {
400 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
401 ptr::copy(k_ptr.add(index + 1), k_ptr.add(index), len - index - 1);
402
403 let v_ptr = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
404 let val = ptr::read(v_ptr.add(index));
405 ptr::copy(v_ptr.add(index + 1), v_ptr.add(index), len - index - 1);
406
407 self.set_len(len - 1);
408 val
409 }
410 }
411
412 fn insert_leaf(&mut self, index: usize, key: i64, value: V) {
413 assert!(self.is_leaf());
414 let len = self.len();
415 assert!(len <= MAX_KEYS);
416 assert!(
417 index <= len,
418 "insert_leaf: index {} > len {} for key {}",
419 index,
420 len,
421 key
422 );
423
424 unsafe {
427 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
428 let p_key = k_ptr.add(index);
429 ptr::copy(p_key, p_key.add(1), len - index);
430 ptr::write(p_key, key);
431
432 let v_ptr = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
433 let p_val = v_ptr.add(index);
434 ptr::copy(p_val, p_val.add(1), len - index);
435 ptr::write(p_val, value);
436 }
437 self.set_len(len + 1);
438 }
439
440 fn push_internal(&mut self, key: i64, child: NodePtr<V>) {
441 assert!(!self.is_leaf());
442 let len = self.len();
443 assert!(len <= MAX_KEYS);
444
445 unsafe {
448 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
449 ptr::write(k_ptr.add(len), key);
450
451 let c_ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
452 ptr::write(c_ptr.add(len + 1), child);
453 }
454 self.set_len(len + 1);
455 }
456
457 fn split_internal(&mut self) -> (i64, NodePtr<V>) {
458 let len = self.len();
459 let mid = len / 2;
460 let med_key = self.keys()[mid];
461 let mut right = NodePtr::new_internal();
462 let right_keys_count = len - mid - 1;
463 let right_children_count = right_keys_count + 1;
464
465 unsafe {
471 let k_src = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
472 let k_dst = right.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
473 ptr::copy_nonoverlapping(k_src.add(mid + 1), k_dst, right_keys_count);
474
475 let c_src = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
476 let c_dst = right.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
477 ptr::copy_nonoverlapping(c_src.add(mid + 1), c_dst, right_children_count);
478 }
479
480 right.set_len(right_keys_count);
481 self.set_len(mid);
482 (med_key, right)
483 }
484
485 fn split_internal_rightmost(&mut self) -> (i64, NodePtr<V>) {
489 let len = self.len();
490 debug_assert!(
491 len == MAX_KEYS + 1,
492 "split_internal_rightmost expects overflow node"
493 );
494
495 let med_key = self.keys()[len - 1];
497
498 let mut right = NodePtr::new_internal();
499
500 unsafe {
504 let c_src = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
505 let c_dst = right.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
506
507 ptr::write(c_dst, ptr::read(c_src.add(len)));
509 }
510
511 right.set_len(0); self.set_len(len - 1); (med_key, right)
515 }
516
517 fn borrow_from_left(&mut self, index: usize) {
518 assert!(!self.is_leaf());
519 let is_left_leaf = self.child(index - 1).is_leaf();
520
521 if is_left_leaf {
522 let (key, val) = {
523 let left = self.child_mut(index - 1).make_mut();
524 let left_len = left.len();
525 let key = left.keys()[left_len - 1];
526 let val = unsafe {
531 let val_ptr = left.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
532 ptr::read(val_ptr.add(left_len - 1))
533 };
534 left.set_len(left_len - 1);
535 (key, val)
536 };
537
538 let current = self.child_mut(index).make_mut();
539 current.insert_leaf(0, key, val);
540
541 unsafe {
545 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
546 ptr::write(k_ptr.add(index - 1), key);
547 }
548 } else {
549 let (child, key) = {
550 let left = self.child_mut(index - 1).make_mut();
551 let left_len = left.len();
552 let key = left.keys()[left_len - 1];
553 let child = unsafe {
558 let c_ptr =
559 left.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
560 ptr::read(c_ptr.add(left_len))
561 };
562 left.set_len(left_len - 1);
563 (child, key)
564 };
565
566 let separator_idx = index - 1;
567 let separator = self.keys()[separator_idx];
568
569 let current = self.child_mut(index).make_mut();
570 current.insert_internal_at_start(separator, child);
571
572 unsafe {
575 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
576 ptr::write(k_ptr.add(separator_idx), key);
577 }
578 }
579 }
580
581 fn insert_internal(&mut self, index: usize, key: i64, child: NodePtr<V>) {
582 assert!(!self.is_leaf());
583 let len = self.len();
584 assert!(len <= MAX_KEYS);
585
586 unsafe {
592 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
593 ptr::copy(k_ptr.add(index), k_ptr.add(index + 1), len - index);
594 ptr::write(k_ptr.add(index), key);
595
596 let c_ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
597 ptr::copy(c_ptr.add(index + 1), c_ptr.add(index + 2), len - index);
598 ptr::write(c_ptr.add(index + 1), child);
599 }
600 self.set_len(len + 1);
601 }
602
603 fn insert_internal_at_start(&mut self, key: i64, child: NodePtr<V>) {
604 let len = self.len();
605 unsafe {
611 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
612 ptr::copy(k_ptr, k_ptr.add(1), len);
613 ptr::write(k_ptr, key);
614
615 let c_ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
616 ptr::copy(c_ptr, c_ptr.add(1), len + 1);
617 ptr::write(c_ptr, child);
618 }
619 self.set_len(len + 1);
620 }
621
622 fn borrow_from_right(&mut self, index: usize) {
623 assert!(!self.is_leaf());
624 let is_right_leaf = self.child(index + 1).is_leaf();
625
626 if is_right_leaf {
627 let (key, val) = {
628 let right = self.child_mut(index + 1).make_mut();
629 let key = right.keys()[0];
630 let val = right.remove_leaf(0);
631 (key, val)
632 };
633
634 let current = self.child_mut(index).make_mut();
635 current.push_leaf(key, val);
636
637 let right = self.child(index + 1);
638 let new_sep = right.keys()[0];
639 unsafe {
643 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
644 ptr::write(k_ptr.add(index), new_sep);
645 }
646 } else {
647 let (child, key) = {
648 let right = self.child_mut(index + 1).make_mut();
649 let key = right.keys()[0];
650 let child = unsafe {
654 let c_ptr =
655 right.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
656 ptr::read(c_ptr)
657 };
658 right.remove_internal_at_start();
659 (child, key)
660 };
661
662 let separator = self.keys()[index];
663
664 let current = self.child_mut(index).make_mut();
665 current.push_internal(separator, child);
666
667 unsafe {
670 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
671 ptr::write(k_ptr.add(index), key);
672 }
673 }
674 }
675
676 fn remove_internal_at_start(&mut self) {
677 let len = self.len();
678 unsafe {
685 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
686 ptr::copy(k_ptr.add(1), k_ptr, len - 1);
687
688 let c_ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
689 ptr::copy(c_ptr.add(1), c_ptr, len);
690 }
691 self.set_len(len - 1);
692 }
693
694 fn merge_with_left(&mut self, index: usize) {
695 let separator = self.keys()[index - 1];
696
697 let right_raw = self.child_mut(index) as *mut NodePtr<V>;
699 let right = unsafe { (*right_raw).make_mut() };
705 let is_leaf = right.is_leaf();
706 let r_len = right.len();
707
708 if is_leaf {
709 let mut keys_vals: Vec<(i64, V)> = Vec::with_capacity(r_len);
711 unsafe {
716 let k_ptr = right.ptr.data_ptr_mut().add(Self::keys_offset()) as *const i64;
717 let v_ptr = right.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
718 for i in 0..r_len {
719 let key = *k_ptr.add(i);
720 let val = ptr::read(v_ptr.add(i)); keys_vals.push((key, val));
722 }
723 }
724 right.set_len(0);
726
727 let left = self.child_mut(index - 1).make_mut();
729 for (k, v) in keys_vals {
730 left.push_leaf(k, v);
731 }
732 } else {
733 let keys: Vec<i64> = right.keys().to_vec();
737 let children: Vec<NodePtr<V>> = unsafe {
740 let c_ptr =
741 right.ptr.data_ptr_mut().add(Self::children_offset()) as *const NodePtr<V>;
742 (0..=r_len).map(|i| ptr::read(c_ptr.add(i))).collect()
743 };
744
745 unsafe {
754 let header = &*(right.ptr.data_ptr_mut() as *const NodeHeader);
755 header.drop_count.store(2, Ordering::Release);
756 }
757 right.set_len(0);
758
759 let left = self.child_mut(index - 1).make_mut();
760
761 let mut children_iter = children.into_iter();
763 left.push_internal(separator, children_iter.next().unwrap());
764 for (key, child) in keys.into_iter().zip(children_iter) {
765 left.push_internal(key, child);
766 }
767 }
768
769 self.remove_key_and_child(index - 1, index);
770 }
771
772 fn merge_with_right(&mut self, index: usize) {
773 self.merge_with_left(index + 1)
774 }
775
776 fn remove_key_and_child(&mut self, key_idx: usize, child_idx: usize) {
777 let len = self.len();
778 unsafe {
785 let k_ptr = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
786 ptr::copy(
787 k_ptr.add(key_idx + 1),
788 k_ptr.add(key_idx),
789 len - key_idx - 1,
790 );
791
792 let c_ptr = self.ptr.data_ptr_mut().add(Self::children_offset()) as *mut NodePtr<V>;
793 ptr::drop_in_place(c_ptr.add(child_idx));
795 ptr::copy(
796 c_ptr.add(child_idx + 1),
797 c_ptr.add(child_idx),
798 len - child_idx,
799 );
800 }
801 self.set_len(len - 1);
802 }
803
804 fn split_leaf(&mut self) -> (i64, NodePtr<V>) {
805 let mid = self.len() / 2;
806 let right_count = self.len() - mid;
807
808 let mut right = NodePtr::new_leaf();
809
810 unsafe {
816 let k_src = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
817 let k_dst = right.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
818 ptr::copy_nonoverlapping(k_src.add(mid), k_dst, right_count);
819
820 let v_src = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
821 let v_dst = right.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
822 ptr::copy_nonoverlapping(v_src.add(mid), v_dst, right_count);
823 }
824
825 right.set_len(right_count);
826 self.set_len(mid);
827
828 let median = right.keys()[0];
829 (median, right)
830 }
831
832 fn split_leaf_rightmost(&mut self) -> (i64, NodePtr<V>) {
837 let len = self.len();
838 debug_assert!(
839 len == MAX_KEYS + 1,
840 "split_leaf_rightmost expects overflow node"
841 );
842
843 let mut right = NodePtr::new_leaf();
844
845 unsafe {
849 let k_src = self.ptr.data_ptr_mut().add(Self::keys_offset()) as *const i64;
850 let k_dst = right.ptr.data_ptr_mut().add(Self::keys_offset()) as *mut i64;
851
852 let v_src = self.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
853 let v_dst = right.ptr.data_ptr_mut().add(Self::values_offset()) as *mut V;
854
855 ptr::write(k_dst, *k_src.add(len - 1));
857
858 ptr::write(v_dst, ptr::read(v_src.add(len - 1)));
860 }
861
862 right.set_len(1);
863 self.set_len(len - 1); let median = right.keys()[0];
866 (median, right)
867 }
868}
869
870const MAX_TREE_DEPTH: usize = 16;
885
886#[derive(Clone, Copy)]
888pub struct NodePath {
889 indices: [u8; MAX_TREE_DEPTH],
890 len: u8,
891}
892
893impl Default for NodePath {
894 fn default() -> Self {
895 Self {
896 indices: [0; MAX_TREE_DEPTH],
897 len: 0,
898 }
899 }
900}
901
902impl NodePath {
903 pub fn new() -> Self {
904 Self::default()
905 }
906
907 #[cold]
908 #[inline(never)]
909 fn depth_overflow() -> ! {
910 panic!("B-Tree depth exceeded maximum")
911 }
912
913 pub fn push(&mut self, idx: usize) {
914 if (self.len as usize) < MAX_TREE_DEPTH {
915 self.indices[self.len as usize] = idx as u8;
916 self.len += 1;
917 } else {
918 Self::depth_overflow();
919 }
920 }
921
922 fn get(&self, depth: usize) -> usize {
923 self.indices[depth] as usize
924 }
925
926 fn iter(&self) -> impl Iterator<Item = usize> + '_ {
927 (0..self.len as usize).map(|i| self.indices[i] as usize)
928 }
929}
930
931pub struct CowBTree<V: Clone> {
932 root: Option<NodePtr<V>>,
933 max_key: i64,
935 len: usize,
937}
938
939impl<V: Clone> Default for CowBTree<V> {
940 fn default() -> Self {
941 Self::new()
942 }
943}
944
945impl<V: Clone> Clone for CowBTree<V> {
946 #[inline]
948 fn clone(&self) -> Self {
949 Self {
950 root: self.root.clone(),
951 max_key: self.max_key,
952 len: self.len,
953 }
954 }
955}
956
957impl<V: Clone> CowBTree<V> {
958 #[inline]
959 pub fn new() -> Self {
960 assert!(
961 mem::align_of::<V>() <= 8,
962 "CowBTree value type alignment must be <= 8"
963 );
964 Self {
965 root: None,
966 max_key: 0,
967 len: 0,
968 }
969 }
970
971 #[inline]
972 pub fn len(&self) -> usize {
973 self.len
974 }
975
976 #[inline]
977 pub fn is_empty(&self) -> bool {
978 self.root.is_none()
979 }
980
981 #[inline]
983 pub fn get(&self, key: i64) -> Option<&V> {
984 let mut node = self.root.as_ref()?;
985 loop {
986 match node.search(key) {
987 Ok(i) => {
988 if node.is_leaf() {
989 return Some(&node.values()[i]);
990 }
991 node = node.child(i + 1);
992 }
993 Err(i) => {
994 if node.is_leaf() {
995 return None;
996 }
997 node = node.child(i);
998 }
999 }
1000 }
1001 }
1002
1003 #[inline]
1005 pub fn get_mut(&mut self, key: i64) -> Option<&mut V> {
1006 let (path, leaf_idx) = self.search_path(key);
1007 let leaf_idx = leaf_idx.ok()?;
1008
1009 self.get_mut_with_path(key, &path, leaf_idx)
1010 }
1011
1012 fn get_mut_with_path(&mut self, _key: i64, path: &NodePath, leaf_idx: usize) -> Option<&mut V> {
1013 let root = self.root.as_mut()?;
1014 let mut node = root.make_mut();
1015
1016 for idx in path.iter() {
1017 let child = node.child_mut(idx);
1018 node = child.make_mut();
1019 }
1020
1021 if leaf_idx < node.len() {
1022 Some(&mut node.values_mut_slice()[leaf_idx])
1023 } else {
1024 None
1025 }
1026 }
1027
1028 #[inline]
1030 pub fn contains_key(&self, key: i64) -> bool {
1031 self.get(key).is_some()
1032 }
1033
1034 pub fn insert(&mut self, key: i64, value: V) -> Option<V> {
1036 if self.root.is_none() {
1037 let mut node = NodePtr::new_leaf();
1038 node.push_leaf(key, value);
1039 self.root = Some(node);
1040 self.max_key = key;
1041 self.len = 1;
1042 return None;
1043 }
1044
1045 if self.is_key_greater_than_max(key) {
1047 let root = self.root.as_mut().unwrap();
1048 let result = Self::insert_rightmost(root, key, value);
1049 self.max_key = key;
1050
1051 return match result {
1052 InsertResult::Done(old) => {
1053 if old.is_none() {
1054 self.len += 1;
1055 }
1056 old
1057 }
1058 InsertResult::Split(median, right) => {
1059 let old_root = self.root.take().unwrap();
1060 let mut new_root = NodePtr::new_internal();
1061
1062 unsafe {
1067 let c_ptr = new_root
1068 .ptr
1069 .data_ptr_mut()
1070 .add(NodePtr::<V>::children_offset())
1071 as *mut NodePtr<V>;
1072 ptr::write(c_ptr, old_root);
1073 }
1074
1075 new_root.push_internal(median, right);
1076 self.root = Some(new_root);
1077 self.len += 1;
1078 None
1079 }
1080 };
1081 }
1082
1083 let root = self.root.as_mut().unwrap();
1084 let result = Self::insert_recursive(root, key, value);
1085
1086 if key > self.max_key {
1087 self.max_key = key;
1088 }
1089
1090 match result {
1091 InsertResult::Done(old) => {
1092 if old.is_none() {
1093 self.len += 1;
1094 }
1095 old
1096 }
1097 InsertResult::Split(median, right) => {
1098 let old_root = self.root.take().unwrap();
1099 let mut new_root = NodePtr::new_internal();
1100 unsafe {
1105 let c_ptr = new_root
1106 .ptr
1107 .data_ptr_mut()
1108 .add(NodePtr::<V>::children_offset())
1109 as *mut NodePtr<V>;
1110 ptr::write(c_ptr, old_root);
1111 }
1112 new_root.push_internal(median, right);
1113 self.root = Some(new_root);
1114 self.len += 1;
1115 None
1116 }
1117 }
1118 }
1119
1120 #[inline]
1122 fn is_key_greater_than_max(&self, key: i64) -> bool {
1123 self.root.is_some() && key > self.max_key
1124 }
1125
1126 fn insert_rightmost(node: &mut NodePtr<V>, key: i64, value: V) -> InsertResult<V> {
1128 let (res, _) = Self::insert_rightmost_return_ptr(node, key, value);
1129 res
1130 }
1131
1132 fn insert_rightmost_return_ptr(
1133 node: &mut NodePtr<V>,
1134 key: i64,
1135 value: V,
1136 ) -> (InsertResult<V>, *mut V) {
1137 let node = node.make_mut();
1138
1139 if node.is_leaf() {
1140 node.push_leaf(key, value);
1141
1142 unsafe {
1148 let len = node.len();
1149 let v_ptr = node.ptr.data_ptr_mut().add(NodePtr::<V>::values_offset()) as *mut V;
1150 let ptr = v_ptr.add(len - 1);
1151
1152 if node.len() > MAX_KEYS {
1153 let (median, right) = node.split_leaf_rightmost();
1155 let v_ptr_new =
1157 right.ptr.data_ptr_mut().add(NodePtr::<V>::values_offset()) as *mut V;
1158 (InsertResult::Split(median, right), v_ptr_new)
1159 } else {
1160 (InsertResult::Done(None), ptr)
1161 }
1162 }
1163 } else {
1164 let last_idx = node.len();
1165 let child = node.child_mut(last_idx);
1166 let (result, ptr) = Self::insert_rightmost_return_ptr(child, key, value);
1167
1168 match result {
1169 InsertResult::Done(old) => (InsertResult::Done(old), ptr),
1170 InsertResult::Split(median, right) => {
1171 node.push_internal(median, right);
1172
1173 if node.len() > MAX_KEYS {
1174 let (m, r) = node.split_internal_rightmost();
1176 (InsertResult::Split(m, r), ptr)
1177 } else {
1178 (InsertResult::Done(None), ptr)
1179 }
1180 }
1181 }
1182 }
1183 }
1184
1185 fn insert_recursive(node: &mut NodePtr<V>, key: i64, value: V) -> InsertResult<V> {
1186 let node = node.make_mut();
1187
1188 if node.is_leaf() {
1189 match node.search(key) {
1190 Ok(i) => unsafe {
1195 let v_ptr =
1196 node.ptr.data_ptr_mut().add(NodePtr::<V>::values_offset()) as *mut V;
1197 let old = ptr::read(v_ptr.add(i));
1198 ptr::write(v_ptr.add(i), value);
1199 InsertResult::Done(Some(old))
1200 },
1201 Err(i) => {
1202 node.insert_leaf(i, key, value);
1203
1204 if node.len() > MAX_KEYS {
1205 let (median, right) = if i == node.len() - 1 {
1208 node.split_leaf_rightmost()
1209 } else {
1210 node.split_leaf()
1211 };
1212 InsertResult::Split(median, right)
1213 } else {
1214 InsertResult::Done(None)
1215 }
1216 }
1217 }
1218 } else {
1219 let i = match node.search(key) {
1220 Ok(i) => i + 1,
1221 Err(i) => i,
1222 };
1223
1224 let result = Self::insert_recursive(node.child_mut(i), key, value);
1225
1226 match result {
1227 InsertResult::Done(old) => InsertResult::Done(old),
1228 InsertResult::Split(median, right) => {
1229 node.insert_internal(i, median, right);
1230
1231 if node.len() > MAX_KEYS {
1232 let (m, r) = if i == node.len() - 1 {
1234 node.split_internal_rightmost()
1235 } else {
1236 node.split_internal()
1237 };
1238 InsertResult::Split(m, r)
1239 } else {
1240 InsertResult::Done(None)
1241 }
1242 }
1243 }
1244 }
1245 }
1246
1247 pub fn remove(&mut self, key: i64) -> Option<V> {
1249 let root = self.root.as_mut()?;
1250 let result = Self::remove_recursive(root, key);
1251
1252 if result.is_some() {
1253 self.len -= 1;
1254
1255 if self.len == 0 {
1256 self.root = None;
1257 self.max_key = 0;
1258 } else if self.max_key == key {
1259 self.refresh_max_key();
1260 }
1261
1262 if let Some(ref mut root) = self.root {
1263 let root = root.make_mut();
1264 if !root.is_leaf() && root.len() == 0 {
1265 let child_node = root.child_mut(0).clone();
1266 self.root = Some(child_node);
1267 } else if root.is_leaf() && root.len() == 0 {
1268 self.root = None;
1269 self.max_key = 0;
1270 }
1271 }
1272 }
1273
1274 result
1275 }
1276
1277 fn search_path(&self, key: i64) -> (NodePath, Result<usize, usize>) {
1281 let mut path = NodePath::new();
1282 if self.root.is_none() {
1283 return (path, Err(0));
1284 }
1285
1286 let mut node = self.root.as_ref().unwrap();
1287 loop {
1288 match node.search(key) {
1289 Ok(i) => {
1290 if node.is_leaf() {
1291 return (path, Ok(i));
1292 }
1293 path.push(i + 1);
1294 node = node.child(i + 1);
1295 }
1296 Err(i) => {
1297 if node.is_leaf() {
1298 return (path, Err(i));
1299 }
1300 path.push(i);
1301 node = node.child(i);
1302 }
1303 }
1304 }
1305 }
1306
1307 fn insert_with_path(
1310 node_ptr: &mut NodePtr<V>,
1311 key: i64,
1312 value: V,
1313 path: &NodePath,
1314 depth: usize,
1315 leaf_idx: usize,
1316 ) -> (InsertResult<V>, *mut V) {
1317 let node = node_ptr.make_mut();
1318
1319 if node.is_leaf() {
1320 let i = leaf_idx;
1322 node.insert_leaf(i, key, value);
1323
1324 if node.len() > MAX_KEYS {
1325 let (median, right_node) = node.split_leaf();
1326 let mid = MAX_KEYS.div_ceil(2);
1329
1330 let ptr = if i < mid {
1331 unsafe {
1335 let v_ptr =
1336 node.ptr.data_ptr_mut().add(NodePtr::<V>::values_offset()) as *mut V;
1337 v_ptr.add(i)
1338 }
1339 } else {
1340 let right_idx = i - mid;
1341 unsafe {
1345 let v_ptr = right_node
1346 .ptr
1347 .data_ptr_mut()
1348 .add(NodePtr::<V>::values_offset())
1349 as *mut V;
1350 v_ptr.add(right_idx)
1351 }
1352 };
1353
1354 (InsertResult::Split(median, right_node), ptr)
1355 } else {
1356 unsafe {
1359 let v_ptr =
1360 node.ptr.data_ptr_mut().add(NodePtr::<V>::values_offset()) as *mut V;
1361 let ptr = v_ptr.add(i);
1362 (InsertResult::Done(None), ptr)
1363 }
1364 }
1365 } else {
1366 let i = path.get(depth);
1367
1368 let (result, ptr) =
1369 Self::insert_with_path(node.child_mut(i), key, value, path, depth + 1, leaf_idx);
1370
1371 match result {
1372 InsertResult::Done(old) => (InsertResult::Done(old), ptr),
1373 InsertResult::Split(median, right) => {
1374 node.insert_internal(i, median, right);
1375
1376 if node.len() > MAX_KEYS {
1377 let (m, r) = if i == node.len() - 1 {
1379 node.split_internal_rightmost()
1380 } else {
1381 node.split_internal()
1382 };
1383 (InsertResult::Split(m, r), ptr)
1384 } else {
1385 (InsertResult::Done(None), ptr)
1386 }
1387 }
1388 }
1389 }
1390 }
1391
1392 fn refresh_max_key(&mut self) {
1393 if let Some(root) = &self.root {
1394 let mut node = root;
1395 loop {
1396 if node.is_leaf() {
1397 self.max_key = node.keys().last().copied().unwrap_or(0);
1398 break;
1399 }
1400 node = node.children().last().unwrap();
1401 }
1402 } else {
1403 self.max_key = 0;
1404 }
1405 }
1406
1407 fn remove_recursive(node: &mut NodePtr<V>, key: i64) -> Option<V> {
1408 let node = node.make_mut();
1409
1410 if node.is_leaf() {
1411 match node.search(key) {
1412 Ok(i) => Some(node.remove_leaf(i)),
1413 Err(_) => None,
1414 }
1415 } else {
1416 let i = match node.search(key) {
1417 Ok(i) => i + 1,
1418 Err(i) => i,
1419 };
1420
1421 if node.child(i).len() <= MIN_KEYS {
1422 Self::ensure_child_can_lose_key(node, i);
1423 }
1424
1425 let new_i = match node.search(key) {
1426 Ok(i) => i + 1,
1427 Err(i) => i,
1428 };
1429
1430 let i = new_i.min(node.len()); Self::remove_recursive(node.child_mut(i), key)
1432 }
1433 }
1434
1435 fn ensure_child_can_lose_key(node: &mut NodePtr<V>, i: usize) {
1436 let can_borrow_left = i > 0 && node.child(i - 1).len() > MIN_KEYS;
1437 let can_borrow_right = i < node.len() && node.child(i + 1).len() > MIN_KEYS;
1438
1439 if can_borrow_left {
1440 node.borrow_from_left(i);
1441 } else if can_borrow_right {
1442 node.borrow_from_right(i);
1443 } else if i > 0 {
1444 node.merge_with_left(i);
1445 } else if i < node.len() {
1446 node.merge_with_right(i);
1447 }
1448 }
1449
1450 pub fn iter_chunks(&self) -> impl Iterator<Item = (&[i64], &[V])> {
1453 CowBTreeChunkIter::new(self.root.as_ref())
1454 }
1455
1456 pub fn iter(&self) -> impl Iterator<Item = (&i64, &V)> {
1458 self.iter_chunks()
1459 .flat_map(|(keys, values)| keys.iter().zip(values.iter()))
1460 }
1461
1462 pub fn keys(&self) -> impl Iterator<Item = i64> + '_ {
1464 self.iter().map(|(k, _)| *k)
1465 }
1466
1467 pub fn values(&self) -> impl Iterator<Item = &V> {
1469 self.iter().map(|(_, v)| v)
1470 }
1471
1472 pub fn range_chunks<R>(&self, range: R) -> impl Iterator<Item = (&[i64], &[V])>
1475 where
1476 R: std::ops::RangeBounds<i64>,
1477 {
1478 CowBTreeRangeChunkIter::new(self.root.as_ref(), range)
1479 }
1480
1481 pub fn range<R>(&self, range: R) -> impl Iterator<Item = (&i64, &V)>
1483 where
1484 R: std::ops::RangeBounds<i64>,
1485 {
1486 self.range_chunks(range)
1487 .flat_map(|(keys, values)| keys.iter().zip(values.iter()))
1488 }
1489
1490 pub fn iter_rev_chunks(&self) -> impl Iterator<Item = (&[i64], &[V])> {
1492 CowBTreeRevChunkIter::new(self.root.as_ref())
1493 }
1494
1495 pub fn iter_rev(&self) -> impl Iterator<Item = (&i64, &V)> {
1497 self.iter_rev_chunks()
1498 .flat_map(|(keys, values)| keys.iter().zip(values.iter()).rev())
1499 }
1500
1501 pub fn range_rev_chunks<R>(&self, range: R) -> impl Iterator<Item = (&[i64], &[V])>
1504 where
1505 R: std::ops::RangeBounds<i64>,
1506 {
1507 CowBTreeRevRangeChunkIter::new(self.root.as_ref(), range)
1508 }
1509
1510 pub fn range_rev<R>(&self, range: R) -> impl Iterator<Item = (&i64, &V)>
1512 where
1513 R: std::ops::RangeBounds<i64>,
1514 {
1515 self.range_rev_chunks(range)
1516 .flat_map(|(keys, values)| keys.iter().zip(values.iter()).rev())
1517 }
1518
1519 pub fn clear(&mut self) {
1521 self.root = None;
1522 self.max_key = 0;
1523 self.len = 0;
1524 }
1525
1526 #[inline]
1528 pub fn max_key(&self) -> Option<i64> {
1529 if self.root.is_some() {
1530 Some(self.max_key)
1531 } else {
1532 None
1533 }
1534 }
1535
1536 fn insert_rightmost_entry(&mut self, key: i64, value: V) -> *mut V {
1537 let root = self.root.as_mut().unwrap();
1538 let (result, ptr) = Self::insert_rightmost_return_ptr(root, key, value);
1539 self.max_key = key;
1540
1541 match result {
1542 InsertResult::Done(old) => {
1543 if old.is_none() {
1544 self.len += 1;
1545 }
1546 }
1547 InsertResult::Split(median, right) => {
1548 let old_root = self.root.take().unwrap();
1549 let mut new_root = NodePtr::new_internal();
1550 unsafe {
1555 let c_ptr = new_root
1556 .ptr
1557 .data_ptr_mut()
1558 .add(NodePtr::<V>::children_offset())
1559 as *mut NodePtr<V>;
1560 ptr::write(c_ptr, old_root);
1561 }
1562 new_root.push_internal(median, right);
1563 self.root = Some(new_root);
1564 self.len += 1;
1565 }
1566 }
1567
1568 ptr
1569 }
1570
1571 pub fn entry(&mut self, key: i64) -> Entry<'_, V> {
1577 if self.is_key_greater_than_max(key) {
1579 return Entry::Vacant(VacantEntry {
1580 tree: self,
1581 key,
1582 path: NodePath::new(), leaf_idx: 0, is_rightmost: true,
1585 });
1586 }
1587
1588 let (path, leaf_result) = self.search_path(key);
1589 match leaf_result {
1590 Ok(idx) => Entry::Occupied(OccupiedEntry {
1591 tree: self,
1592 key,
1593 path,
1594 leaf_idx: idx,
1595 }),
1596 Err(idx) => Entry::Vacant(VacantEntry {
1597 tree: self,
1598 key,
1599 path,
1600 leaf_idx: idx,
1601 is_rightmost: false,
1602 }),
1603 }
1604 }
1605
1606 fn insert_using_path(
1607 &mut self,
1608 key: i64,
1609 value: V,
1610 path: &NodePath,
1611 leaf_idx: usize,
1612 ) -> *mut V {
1613 if self.root.is_none() {
1614 self.insert(key, value);
1615 return self.get_mut(key).unwrap() as *mut V;
1616 }
1617
1618 let root = self.root.as_mut().unwrap();
1619 let (result, ptr) = Self::insert_with_path(root, key, value, path, 0, leaf_idx);
1620
1621 if key > self.max_key {
1622 self.max_key = key;
1623 }
1624
1625 match result {
1626 InsertResult::Done(old) => {
1627 if old.is_none() {
1628 self.len += 1;
1629 }
1630 }
1631 InsertResult::Split(median, right) => {
1632 let old_root = self.root.take().unwrap();
1633 let mut new_root = NodePtr::new_internal();
1634 unsafe {
1639 let c_ptr = new_root
1640 .ptr
1641 .data_ptr_mut()
1642 .add(NodePtr::<V>::children_offset())
1643 as *mut NodePtr<V>;
1644 ptr::write(c_ptr, old_root);
1645 }
1646 new_root.push_internal(median, right);
1647 self.root = Some(new_root);
1648 self.len += 1;
1649 }
1650 }
1651
1652 ptr
1653 }
1654}
1655
1656pub enum Entry<'a, V: Clone> {
1658 Occupied(OccupiedEntry<'a, V>),
1659 Vacant(VacantEntry<'a, V>),
1660}
1661
1662impl<'a, V: Clone> Entry<'a, V> {
1663 pub fn or_insert(self, default: V) -> &'a mut V {
1664 match self {
1665 Entry::Occupied(entry) => entry.into_mut(),
1666 Entry::Vacant(entry) => entry.insert(default),
1667 }
1668 }
1669
1670 pub fn and_modify<F>(self, f: F) -> Self
1671 where
1672 F: FnOnce(&mut V),
1673 {
1674 match self {
1675 Entry::Occupied(mut entry) => {
1676 f(entry.get_mut());
1677 Entry::Occupied(entry)
1678 }
1679 Entry::Vacant(entry) => Entry::Vacant(entry),
1680 }
1681 }
1682
1683 pub fn key(&self) -> i64 {
1684 match self {
1685 Entry::Occupied(entry) => entry.key(),
1686 Entry::Vacant(entry) => entry.key(),
1687 }
1688 }
1689}
1690
1691pub struct OccupiedEntry<'a, V: Clone> {
1693 tree: &'a mut CowBTree<V>,
1694 key: i64,
1695 path: NodePath,
1696 leaf_idx: usize,
1697}
1698
1699impl<'a, V: Clone> OccupiedEntry<'a, V> {
1700 #[inline]
1701 pub fn key(&self) -> i64 {
1702 self.key
1703 }
1704
1705 #[inline]
1707 pub fn get(&self) -> &V {
1708 let mut node = self.tree.root.as_ref().unwrap();
1709 for idx in self.path.iter() {
1710 node = node.child(idx);
1711 }
1712 &node.values()[self.leaf_idx]
1713 }
1714
1715 pub fn get_mut(&mut self) -> &mut V {
1716 self.tree
1717 .get_mut_with_path(self.key, &self.path, self.leaf_idx)
1718 .unwrap()
1719 }
1720
1721 pub fn into_mut(self) -> &'a mut V {
1722 self.tree
1723 .get_mut_with_path(self.key, &self.path, self.leaf_idx)
1724 .unwrap()
1725 }
1726
1727 pub fn insert(&mut self, value: V) -> V {
1728 let node = self.tree.root.as_mut().unwrap();
1729 let mut node = node.make_mut();
1730
1731 for idx in self.path.iter() {
1732 let child = node.child_mut(idx);
1733 node = child.make_mut();
1734 }
1735
1736 unsafe {
1741 let v_ptr = node.ptr.data_ptr_mut().add(NodePtr::<V>::values_offset()) as *mut V;
1742 let ptr = v_ptr.add(self.leaf_idx);
1743 let old = ptr::read(ptr);
1744 ptr::write(ptr, value);
1745 old
1746 }
1747 }
1748}
1749
1750pub struct VacantEntry<'a, V: Clone> {
1752 tree: &'a mut CowBTree<V>,
1753 key: i64,
1754 path: NodePath,
1755 leaf_idx: usize,
1757 is_rightmost: bool,
1759}
1760
1761impl<'a, V: Clone> VacantEntry<'a, V> {
1762 #[inline]
1763 pub fn key(&self) -> i64 {
1764 self.key
1765 }
1766
1767 #[inline]
1768 pub fn insert(self, value: V) -> &'a mut V {
1769 if self.is_rightmost {
1770 let ptr = self.tree.insert_rightmost_entry(self.key, value);
1771 unsafe { &mut *ptr }
1776 } else {
1777 let ptr = self
1778 .tree
1779 .insert_using_path(self.key, value, &self.path, self.leaf_idx);
1780 unsafe { &mut *ptr }
1785 }
1786 }
1787}
1788
1789enum InsertResult<V: Clone> {
1790 Done(Option<V>),
1791 Split(i64, NodePtr<V>),
1792}
1793
1794struct CowBTreeChunkIter<'a, V: Clone> {
1796 stack: Vec<(&'a NodePtr<V>, usize)>,
1799 current_leaf: Option<&'a NodePtr<V>>,
1801}
1802
1803impl<'a, V: Clone> CowBTreeChunkIter<'a, V> {
1804 fn new(root: Option<&'a NodePtr<V>>) -> Self {
1805 let mut iter = Self {
1806 stack: Vec::new(),
1807 current_leaf: None,
1808 };
1809 if let Some(root) = root {
1810 iter.descend_to_leftmost(root);
1811 }
1812 iter
1813 }
1814
1815 fn descend_to_leftmost(&mut self, mut node: &'a NodePtr<V>) {
1817 while !node.is_leaf() {
1818 self.stack.push((node, 1));
1819 node = node.child(0);
1820 }
1821 self.current_leaf = Some(node);
1822 }
1823}
1824
1825impl<'a, V: Clone> Iterator for CowBTreeChunkIter<'a, V> {
1826 type Item = (&'a [i64], &'a [V]);
1827
1828 fn next(&mut self) -> Option<Self::Item> {
1829 if let Some(leaf) = self.current_leaf.take() {
1830 return Some((leaf.keys(), leaf.values()));
1831 }
1832
1833 loop {
1834 let (node, idx) = self.stack.last_mut()?;
1835
1836 if *idx < node.len() + 1 {
1837 let child_idx = *idx;
1838 *idx += 1;
1839 let child = node.child(child_idx);
1840 self.descend_to_leftmost(child);
1841
1842 if let Some(leaf) = self.current_leaf.take() {
1843 return Some((leaf.keys(), leaf.values()));
1844 }
1845 } else {
1846 self.stack.pop();
1847 }
1848 }
1849 }
1850}
1851
1852struct CowBTreeRangeChunkIter<'a, V: Clone, R> {
1855 stack: Vec<(&'a NodePtr<V>, usize)>,
1856 range: R,
1857 current_leaf: Option<&'a NodePtr<V>>,
1858 current_idx: usize,
1859 finished: bool,
1860}
1861
1862impl<'a, V: Clone, R: std::ops::RangeBounds<i64>> CowBTreeRangeChunkIter<'a, V, R> {
1863 fn new(root: Option<&'a NodePtr<V>>, range: R) -> Self {
1864 let mut iter = Self {
1865 stack: Vec::new(),
1866 range,
1867 current_leaf: None,
1868 current_idx: 0,
1869 finished: false,
1870 };
1871 if let Some(root) = root {
1872 iter.seek_to_start(root);
1873 } else {
1874 iter.finished = true;
1875 }
1876 iter
1877 }
1878
1879 fn seek_to_start(&mut self, mut node: &'a NodePtr<V>) {
1880 let start_key = match self.range.start_bound() {
1881 Bound::Included(&k) => Some(k),
1882 Bound::Excluded(&k) => Some(k),
1883 Bound::Unbounded => None,
1884 };
1885
1886 loop {
1887 if node.is_leaf() {
1888 let keys = node.keys();
1889 let mut idx = if let Some(k) = start_key {
1890 match keys.binary_search(&k) {
1891 Ok(i) => i,
1892 Err(i) => i,
1893 }
1894 } else {
1895 0
1896 };
1897
1898 if let Bound::Excluded(&k) = self.range.start_bound() {
1899 if idx < keys.len() && keys[idx] == k {
1900 idx += 1;
1901 }
1902 }
1903
1904 self.current_leaf = Some(node);
1905 self.current_idx = idx;
1906 break;
1907 } else {
1908 let idx = if let Some(k) = start_key {
1909 match node.search(k) {
1910 Ok(i) => i + 1,
1911 Err(i) => i,
1912 }
1913 } else {
1914 0
1915 };
1916
1917 self.stack.push((node, idx + 1));
1918 node = node.child(idx);
1919 }
1920 }
1921 }
1922}
1923
1924impl<'a, V: Clone, R: std::ops::RangeBounds<i64>> Iterator for CowBTreeRangeChunkIter<'a, V, R> {
1925 type Item = (&'a [i64], &'a [V]);
1926
1927 fn next(&mut self) -> Option<Self::Item> {
1928 if self.finished {
1929 return None;
1930 }
1931
1932 loop {
1933 if let Some(leaf) = self.current_leaf {
1934 let keys = leaf.keys();
1935 let values = leaf.values();
1936 if self.current_idx < keys.len() {
1937 let start = self.current_idx;
1938 let end = match self.range.end_bound() {
1939 Bound::Unbounded => keys.len(),
1940 Bound::Included(&k) => {
1941 if keys.last().unwrap() <= &k {
1942 keys.len()
1943 } else {
1944 let pos = keys[start..].partition_point(|&x| x <= k);
1945 self.finished = true;
1946 start + pos
1947 }
1948 }
1949 Bound::Excluded(&k) => {
1950 if keys.last().unwrap() < &k {
1951 keys.len()
1952 } else {
1953 let pos = keys[start..].partition_point(|&x| x < k);
1954 self.finished = true;
1955 start + pos
1956 }
1957 }
1958 };
1959
1960 if start >= end {
1961 self.finished = true;
1962 self.current_leaf = None;
1963 return None;
1964 }
1965
1966 self.current_idx = end;
1967 let result = (&keys[start..end], &values[start..end]);
1968
1969 if end == keys.len() && !self.finished {
1970 self.current_leaf = None;
1971 } else {
1972 self.finished = true;
1973 self.current_leaf = None;
1974 }
1975
1976 return Some(result);
1977 } else {
1978 self.current_leaf = None;
1979 }
1980 }
1981
1982 if self.finished {
1983 return None;
1984 }
1985
1986 if let Some((node, idx)) = self.stack.last_mut() {
1987 if *idx < node.len() + 1 {
1988 let child_idx = *idx;
1989 *idx += 1;
1990 let mut child = node.child(child_idx);
1991
1992 loop {
1993 if child.is_leaf() {
1994 self.current_leaf = Some(child);
1995 self.current_idx = 0;
1996 break;
1997 } else {
1998 self.stack.push((child, 1));
1999 child = child.child(0);
2000 }
2001 }
2002 } else {
2003 self.stack.pop();
2004 if self.stack.is_empty() {
2005 self.finished = true;
2006 return None;
2007 }
2008 }
2009 } else {
2010 self.finished = true;
2011 return None;
2012 }
2013 }
2014 }
2015}
2016
2017impl<V: Clone + std::fmt::Debug> std::fmt::Debug for CowBTree<V> {
2018 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2019 f.debug_map().entries(self.iter()).finish()
2020 }
2021}
2022
2023struct CowBTreeRevChunkIter<'a, V: Clone> {
2025 stack: Vec<(&'a NodePtr<V>, usize)>,
2028 current_leaf: Option<&'a NodePtr<V>>,
2030}
2031
2032impl<'a, V: Clone> CowBTreeRevChunkIter<'a, V> {
2033 fn new(root: Option<&'a NodePtr<V>>) -> Self {
2034 let mut iter = Self {
2035 stack: Vec::new(),
2036 current_leaf: None,
2037 };
2038 if let Some(root) = root {
2039 iter.descend_to_rightmost(root);
2040 }
2041 iter
2042 }
2043
2044 fn descend_to_rightmost(&mut self, mut node: &'a NodePtr<V>) {
2046 while !node.is_leaf() {
2047 let last_child = node.len(); self.stack.push((node, last_child));
2051 node = node.child(last_child);
2052 }
2053 self.current_leaf = Some(node);
2054 }
2055}
2056
2057impl<'a, V: Clone> Iterator for CowBTreeRevChunkIter<'a, V> {
2058 type Item = (&'a [i64], &'a [V]);
2059
2060 fn next(&mut self) -> Option<Self::Item> {
2061 if let Some(leaf) = self.current_leaf.take() {
2062 return Some((leaf.keys(), leaf.values()));
2063 }
2064
2065 loop {
2066 let (node, next_plus_one) = self.stack.last_mut()?;
2067
2068 if *next_plus_one > 0 {
2069 let child_idx = *next_plus_one - 1;
2070 *next_plus_one = child_idx;
2071 let child = node.child(child_idx);
2072 self.descend_to_rightmost(child);
2073
2074 if let Some(leaf) = self.current_leaf.take() {
2075 return Some((leaf.keys(), leaf.values()));
2076 }
2077 } else {
2078 self.stack.pop();
2079 }
2080 }
2081 }
2082}
2083
2084struct CowBTreeRevRangeChunkIter<'a, V: Clone, R> {
2087 stack: Vec<(&'a NodePtr<V>, usize)>,
2088 range: R,
2089 current_leaf: Option<&'a NodePtr<V>>,
2090 current_end_idx: usize,
2092 finished: bool,
2093}
2094
2095impl<'a, V: Clone, R: std::ops::RangeBounds<i64>> CowBTreeRevRangeChunkIter<'a, V, R> {
2096 fn new(root: Option<&'a NodePtr<V>>, range: R) -> Self {
2097 let mut iter = Self {
2098 stack: Vec::new(),
2099 range,
2100 current_leaf: None,
2101 current_end_idx: 0,
2102 finished: false,
2103 };
2104 if let Some(root) = root {
2105 iter.seek_to_end(root);
2106 } else {
2107 iter.finished = true;
2108 }
2109 iter
2110 }
2111
2112 fn descend_to_rightmost(&mut self, mut node: &'a NodePtr<V>) {
2114 while !node.is_leaf() {
2115 let last_child = node.len();
2116 self.stack.push((node, last_child));
2117 node = node.child(last_child);
2118 }
2119 self.current_leaf = Some(node);
2120 self.current_end_idx = node.len();
2121 }
2122
2123 fn seek_to_end(&mut self, mut node: &'a NodePtr<V>) {
2125 let end_key = match self.range.end_bound() {
2126 Bound::Included(&k) | Bound::Excluded(&k) => Some(k),
2127 Bound::Unbounded => None,
2128 };
2129
2130 if end_key.is_none() {
2131 self.descend_to_rightmost(node);
2133 return;
2134 }
2135
2136 let k = end_key.unwrap();
2137
2138 loop {
2139 if node.is_leaf() {
2140 let keys = node.keys();
2141 let idx = match keys.binary_search(&k) {
2142 Ok(i) => match self.range.end_bound() {
2143 Bound::Included(_) => i + 1,
2144 _ => i,
2145 },
2146 Err(i) => i,
2147 };
2148
2149 if idx > 0 {
2150 self.current_leaf = Some(node);
2151 self.current_end_idx = idx;
2152 }
2153 break;
2156 } else {
2157 let child_idx = match node.search(k) {
2158 Ok(i) => i + 1,
2159 Err(i) => i,
2160 };
2161 self.stack.push((node, child_idx));
2163 node = node.child(child_idx);
2164 }
2165 }
2166 }
2167}
2168
2169impl<'a, V: Clone, R: std::ops::RangeBounds<i64>> Iterator for CowBTreeRevRangeChunkIter<'a, V, R> {
2170 type Item = (&'a [i64], &'a [V]);
2171
2172 fn next(&mut self) -> Option<Self::Item> {
2173 if self.finished {
2174 return None;
2175 }
2176
2177 loop {
2178 if let Some(leaf) = self.current_leaf {
2179 let keys = leaf.keys();
2180 let values = leaf.values();
2181 let end = self.current_end_idx;
2182
2183 let start = if end == 0 {
2185 end } else {
2187 match self.range.start_bound() {
2188 Bound::Unbounded => 0,
2189 Bound::Included(&k) => {
2190 if keys[0] >= k {
2191 0 } else {
2193 self.finished = true;
2194 keys[..end].partition_point(|&x| x < k)
2195 }
2196 }
2197 Bound::Excluded(&k) => {
2198 if keys[0] > k {
2199 0
2200 } else {
2201 self.finished = true;
2202 keys[..end].partition_point(|&x| x <= k)
2203 }
2204 }
2205 }
2206 };
2207
2208 self.current_leaf = None;
2209
2210 if start < end {
2211 return Some((&keys[start..end], &values[start..end]));
2212 }
2213
2214 if self.finished {
2215 return None;
2216 }
2217 }
2218
2219 let (node, next_plus_one) = self.stack.last_mut()?;
2221
2222 if *next_plus_one > 0 {
2223 let child_idx = *next_plus_one - 1;
2224 *next_plus_one = child_idx;
2225 let child = node.child(child_idx);
2226 self.descend_to_rightmost(child);
2227 } else {
2228 self.stack.pop();
2229 }
2230 }
2231 }
2232}
2233
2234include!("cow_btree/tests.rs");