1use crate::chunk::{Chunk, Loc, MutUpdate, Update, UpdateChunk};
2use alloc::{
3 sync::{Arc, Weak},
4 vec::Vec,
5};
6use arrayvec::ArrayVec;
7#[cfg(feature = "pool")]
8use core::mem::MaybeUninit;
9use core::{
10 borrow::Borrow,
11 cmp::{max, min, Eq, Ord, Ordering, PartialEq, PartialOrd},
12 default::Default,
13 fmt::{self, Debug, Formatter},
14 hash::{Hash, Hasher},
15 iter,
16 marker::PhantomData,
17 ops::{Bound, Deref, Index, RangeBounds, RangeFull},
18 slice,
19};
20
21#[cfg(feature = "pool")]
22use core::{
23 cell::Cell,
24 mem::{self, ManuallyDrop},
25 ptr,
26 sync::atomic::{AtomicUsize, Ordering as AOrdering},
27};
28#[cfg(feature = "pool")]
29use poolshark::{
30 local::{insert_raw, take},
31 location_id, Discriminant, IsoPoolable, Poolable,
32};
33#[cfg(feature = "pool")]
34use std::{
35 boxed::Box,
36 collections::BTreeMap,
37 sync::{Mutex, MutexGuard},
38};
39
40const MAX_DEPTH: usize = 64;
42
43#[cfg(feature = "pool")]
55const MAX_DROP_DEPTH: usize = 256;
56
57#[cfg(feature = "pool")]
58std::thread_local! {
59 static DROP_DEPTH: Cell<usize> = const { Cell::new(0) };
64}
65
66#[cfg(feature = "pool")]
74struct Deferred(unsafe fn(*mut ()), *mut ());
75
76#[cfg(feature = "pool")]
77unsafe impl Send for Deferred {}
78
79#[cfg(feature = "pool")]
90static DROP_DEFERRED: Mutex<BTreeMap<usize, Vec<Deferred>>> = Mutex::new(BTreeMap::new());
91
92#[cfg(feature = "pool")]
97static DROP_DEFERRED_LEN: AtomicUsize = AtomicUsize::new(0);
98
99#[cfg(feature = "pool")]
100fn deferred_lock() -> MutexGuard<'static, BTreeMap<usize, Vec<Deferred>>> {
101 match DROP_DEFERRED.lock() {
104 Ok(g) => g,
105 Err(e) => e.into_inner(),
106 }
107}
108
109#[cfg(feature = "pool")]
118unsafe fn drop_deferred_node<K: Ord + Clone, V: Clone, const SIZE: usize>(p: *mut ()) {
119 drop(Box::from_raw(p as *mut Node<K, V, SIZE>))
120}
121
122#[cfg(feature = "pool")]
130struct DepthGuard {
131 depth: usize,
132 tag: usize,
133}
134
135#[cfg(feature = "pool")]
136impl Drop for DepthGuard {
137 fn drop(&mut self) {
138 let _ = DROP_DEPTH.try_with(|d| d.set(self.depth));
139 if self.depth == 0 && DROP_DEFERRED_LEN.load(AOrdering::Relaxed) > 0 {
140 drain_deferred(self.tag)
141 }
142 }
143}
144
145#[cfg(feature = "pool")]
171fn drain_deferred(tag: usize) {
172 let _ = DROP_DEPTH.try_with(|d| d.set(1));
173 let mut panic = None;
174 loop {
175 let Some(batch) = deferred_lock().remove(&tag) else {
181 break;
182 };
183 DROP_DEFERRED_LEN.fetch_sub(batch.len(), AOrdering::Relaxed);
184 for Deferred(f, p) in batch {
185 let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| unsafe {
186 f(p)
187 }));
188 if let Err(e) = r {
189 if panic.is_none() {
190 panic = Some(e)
191 }
192 }
193 }
194 }
195 let _ = DROP_DEPTH.try_with(|d| d.set(0));
196 if let Some(e) = panic {
197 if !std::thread::panicking() {
198 std::panic::resume_unwind(e)
199 }
200 }
201}
202
203fn pack_height_and_size(height: u8, size: usize) -> u64 {
204 assert!((size & 0x00ffffff_ffffffff) == size);
205 ((height as u64) << 56) | (size as u64)
206}
207
208#[cfg(feature = "pool")]
209#[derive(Debug)]
210#[repr(C)]
211pub(crate) struct NodeInner<K: Ord + Clone, V: Clone, const SIZE: usize> {
212 elts: MaybeUninit<Chunk<K, V, SIZE>>,
213 min_key: MaybeUninit<K>,
214 max_key: MaybeUninit<K>,
215 left: Tree<K, V, SIZE>,
216 right: Tree<K, V, SIZE>,
217 height_and_size: u64,
218}
219
220#[cfg(feature = "pool")]
221impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for NodeInner<K, V, SIZE> {
222 fn clone(&self) -> Self {
223 unsafe {
224 Self {
225 elts: MaybeUninit::new(self.elts.assume_init_ref().clone()),
226 min_key: MaybeUninit::new(self.min_key.assume_init_ref().clone()),
227 max_key: MaybeUninit::new(self.max_key.assume_init_ref().clone()),
228 left: self.left.clone(),
229 right: self.right.clone(),
230 height_and_size: self.height_and_size,
231 }
232 }
233 }
234}
235
236#[cfg(not(feature = "pool"))]
237#[derive(Clone, Debug)]
238pub(crate) struct NodeInner<K: Ord + Clone, V: Clone, const SIZE: usize> {
239 elts: Chunk<K, V, SIZE>,
240 min_key: K,
241 max_key: K,
242 left: Tree<K, V, SIZE>,
243 right: Tree<K, V, SIZE>,
244 height_and_size: u64,
245}
246
247#[cfg(feature = "pool")]
248pub(crate) struct Node<K: Ord + Clone, V: Clone, const SIZE: usize>(
249 ManuallyDrop<Arc<NodeInner<K, V, SIZE>>>,
250);
251
252#[cfg(not(feature = "pool"))]
253pub(crate) struct Node<K: Ord + Clone, V: Clone, const SIZE: usize>(
254 Arc<NodeInner<K, V, SIZE>>,
255);
256
257#[cfg(feature = "pool")]
258impl<K: Ord + Clone, V: Clone, const SIZE: usize> Poolable for Node<K, V, SIZE> {
259 fn capacity(&self) -> usize {
260 1
261 }
262
263 fn empty() -> Self {
264 let n = NodeInner {
265 elts: MaybeUninit::zeroed(),
266 min_key: MaybeUninit::zeroed(),
267 max_key: MaybeUninit::zeroed(),
268 left: Tree::Empty,
269 right: Tree::Empty,
270 height_and_size: 0,
271 };
272 Node(ManuallyDrop::new(Arc::new(n)))
273 }
274
275 fn really_dropped(&mut self) -> bool {
276 unreachable!()
277 }
278
279 fn reset(&mut self) {
280 unreachable!()
281 }
282}
283
284#[cfg(feature = "pool")]
285unsafe impl<K: Ord + Clone, V: Clone, const SIZE: usize> IsoPoolable
286 for Node<K, V, SIZE>
287{
288 const DISCRIMINANT: Option<Discriminant> =
289 Discriminant::new_p2_size::<K, V, SIZE>(location_id!());
290}
291
292#[cfg(feature = "pool")]
293impl<K: Ord + Clone, V: Clone, const SIZE: usize> Drop for Node<K, V, SIZE> {
294 fn drop(&mut self) {
295 let cell = DROP_DEPTH
305 .try_with(|d| (d.get(), d as *const Cell<usize> as usize))
306 .ok();
307 let Some((depth, tag)) = cell else {
308 return self.really_drop();
309 };
310 if depth >= MAX_DROP_DEPTH {
311 let b = Box::new(unsafe { ptr::read(self) });
312 let raw = Box::into_raw(b) as *mut ();
313 deferred_lock().entry(tag).or_default().push(Deferred(
314 drop_deferred_node::<K, V, SIZE> as unsafe fn(*mut ()),
315 raw,
316 ));
317 DROP_DEFERRED_LEN.fetch_add(1, AOrdering::Relaxed);
318 return;
319 }
320 let _guard = DepthGuard { depth, tag };
321 let _ = DROP_DEPTH.try_with(|d| d.set(depth + 1));
322 self.really_drop()
323 }
324}
325
326#[cfg(feature = "pool")]
327impl<K: Ord + Clone, V: Clone, const SIZE: usize> Node<K, V, SIZE> {
328 fn really_drop(&mut self) {
331 match Arc::get_mut(&mut self.0) {
332 None => unsafe { ManuallyDrop::drop(&mut self.0) },
333 Some(inner) => {
334 unsafe { inner.reset() }
335 if let Some(mut n) = unsafe { insert_raw(ptr::read(self)) } {
336 unsafe { ManuallyDrop::drop(&mut n.0) };
337 mem::forget(n); }
339 }
340 }
341 }
342}
343
344#[cfg(feature = "pool")]
345impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for Node<K, V, SIZE> {
346 fn clone(&self) -> Self {
347 Self(ManuallyDrop::new(Arc::clone(&*self.0)))
348 }
349}
350
351#[cfg(not(feature = "pool"))]
352impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for Node<K, V, SIZE> {
353 fn clone(&self) -> Self {
354 Self(Arc::clone(&self.0))
355 }
356}
357
358impl<K: Ord + Clone, V: Clone, const SIZE: usize> Deref for Node<K, V, SIZE> {
359 type Target = NodeInner<K, V, SIZE>;
360
361 fn deref(&self) -> &Self::Target {
362 &*self.0
363 }
364}
365
366impl<K: Ord + Clone, V: Clone, const SIZE: usize> Node<K, V, SIZE> {
367 fn downgrade(&self) -> WeakNode<K, V, SIZE> {
368 WeakNode(Arc::downgrade(&self.0))
369 }
370
371 #[cfg(feature = "pool")]
372 fn make_mut<'a>(&'a mut self) -> &'a mut NodeInner<K, V, SIZE> {
373 match Arc::get_mut(&mut *self.0).map(|n| n as *mut _) {
374 Some(t) => unsafe { &mut *t },
375 None => {
376 let mut n = take::<Node<K, V, SIZE>>();
377 *Arc::get_mut(&mut *n.0).unwrap() = (**self.0).clone();
378 *self = n;
379 Arc::get_mut(&mut *self.0).unwrap()
380 }
381 }
382 }
383
384 #[cfg(not(feature = "pool"))]
385 fn make_mut(&mut self) -> &mut NodeInner<K, V, SIZE> {
386 Arc::make_mut(&mut self.0)
387 }
388
389 #[cfg(feature = "pool")]
390 fn arc(&self) -> &Arc<NodeInner<K, V, SIZE>> {
391 &self.0
392 }
393
394 #[cfg(not(feature = "pool"))]
395 fn arc(&self) -> &Arc<NodeInner<K, V, SIZE>> {
396 &self.0
397 }
398}
399
400#[derive(Clone)]
401pub(crate) struct WeakNode<K: Ord + Clone, V: Clone, const SIZE: usize>(
402 Weak<NodeInner<K, V, SIZE>>,
403);
404
405impl<K: Ord + Clone, V: Clone, const SIZE: usize> WeakNode<K, V, SIZE> {
406 fn upgrade(&self) -> Option<Node<K, V, SIZE>> {
407 #[cfg(feature = "pool")]
408 {
409 Weak::upgrade(&self.0).map(|n| Node(ManuallyDrop::new(n)))
410 }
411 #[cfg(not(feature = "pool"))]
412 {
413 Weak::upgrade(&self.0).map(Node)
414 }
415 }
416}
417
418impl<K, V, const SIZE: usize> NodeInner<K, V, SIZE>
419where
420 K: Ord + Clone,
421 V: Clone,
422{
423 #[cfg(feature = "pool")]
424 unsafe fn reset(&mut self) {
425 let Self {
426 elts,
427 min_key,
428 max_key,
429 left,
430 right,
431 height_and_size,
432 } = self;
433 if *height_and_size > 0 {
434 unsafe {
435 elts.assume_init_drop();
436 min_key.assume_init_drop();
437 max_key.assume_init_drop();
438 }
439 *left = Tree::Empty;
440 *right = Tree::Empty;
441 *height_and_size = 0
442 }
443 }
444
445 #[cfg(feature = "pool")]
447 fn elts(&self) -> &Chunk<K, V, SIZE> {
448 unsafe { self.elts.assume_init_ref() }
449 }
450
451 #[cfg(not(feature = "pool"))]
452 fn elts(&self) -> &Chunk<K, V, SIZE> {
453 &self.elts
454 }
455
456 #[cfg(feature = "pool")]
458 fn elts_mut(&mut self) -> &mut Chunk<K, V, SIZE> {
459 unsafe { self.elts.assume_init_mut() }
460 }
461
462 #[cfg(not(feature = "pool"))]
463 fn elts_mut(&mut self) -> &mut Chunk<K, V, SIZE> {
464 &mut self.elts
465 }
466
467 #[cfg(feature = "pool")]
469 fn min_key(&self) -> &K {
470 unsafe { self.min_key.assume_init_ref() }
471 }
472
473 #[cfg(not(feature = "pool"))]
474 fn min_key(&self) -> &K {
475 &self.min_key
476 }
477
478 #[cfg(feature = "pool")]
480 fn max_key(&self) -> &K {
481 unsafe { self.max_key.assume_init_ref() }
482 }
483
484 #[cfg(not(feature = "pool"))]
485 fn max_key(&self) -> &K {
486 &self.max_key
487 }
488
489 fn height(&self) -> u8 {
490 (self.height_and_size >> 56) as u8
491 }
492
493 #[cfg(feature = "pool")]
494 fn mutated(&mut self) {
495 unsafe {
496 if let Some((min, max)) = self.elts().min_max_key() {
497 *self.min_key.assume_init_mut() = min;
498 *self.max_key.assume_init_mut() = max;
499 }
500 self.height_and_size = pack_height_and_size(
501 1 + max(self.left.height(), self.right.height()),
502 self.left.len() + self.right.len(),
503 );
504 }
505 }
506
507 #[cfg(not(feature = "pool"))]
508 fn mutated(&mut self) {
509 if let Some((min, max)) = self.elts().min_max_key() {
510 self.min_key = min;
511 self.max_key = max;
512 }
513 self.height_and_size = pack_height_and_size(
514 1 + max(self.left.height(), self.right.height()),
515 self.left.len() + self.right.len(),
516 );
517 }
518}
519
520#[derive(Clone)]
521pub(crate) enum WeakTree<K: Ord + Clone, V: Clone, const SIZE: usize> {
522 Empty,
523 Node(WeakNode<K, V, SIZE>),
524}
525
526impl<K: Ord + Clone, V: Clone, const SIZE: usize> WeakTree<K, V, SIZE> {
527 pub(crate) fn upgrade(&self) -> Option<Tree<K, V, SIZE>> {
528 match self {
529 WeakTree::Empty => Some(Tree::Empty),
530 WeakTree::Node(n) => n.upgrade().map(Tree::Node),
531 }
532 }
533}
534
535#[derive(Clone)]
536pub(crate) enum Tree<K: Ord + Clone, V: Clone, const SIZE: usize> {
537 Empty,
538 Node(Node<K, V, SIZE>),
539}
540
541impl<K, V, const SIZE: usize> Hash for Tree<K, V, SIZE>
542where
543 K: Hash + Ord + Clone,
544 V: Hash + Clone,
545{
546 fn hash<H: Hasher>(&self, state: &mut H) {
547 for elt in self {
548 elt.hash(state)
549 }
550 }
551}
552
553impl<K, V, const SIZE: usize> Default for Tree<K, V, SIZE>
554where
555 K: Ord + Clone,
556 V: Clone,
557{
558 fn default() -> Tree<K, V, SIZE> {
559 Tree::Empty
560 }
561}
562
563impl<K, V, const SIZE: usize> PartialEq for Tree<K, V, SIZE>
564where
565 K: PartialEq + Ord + Clone,
566 V: PartialEq + Clone,
567{
568 fn eq(&self, other: &Tree<K, V, SIZE>) -> bool {
569 self.len() == other.len() && self.into_iter().zip(other).all(|(e0, e1)| e0 == e1)
570 }
571}
572
573impl<K, V, const SIZE: usize> Eq for Tree<K, V, SIZE>
574where
575 K: Eq + Ord + Clone,
576 V: Eq + Clone,
577{
578}
579
580impl<K, V, const SIZE: usize> PartialOrd for Tree<K, V, SIZE>
581where
582 K: Ord + Clone,
583 V: PartialOrd + Clone,
584{
585 fn partial_cmp(&self, other: &Tree<K, V, SIZE>) -> Option<Ordering> {
586 self.into_iter().partial_cmp(other.into_iter())
587 }
588}
589
590impl<K, V, const SIZE: usize> Ord for Tree<K, V, SIZE>
591where
592 K: Ord + Clone,
593 V: Ord + Clone,
594{
595 fn cmp(&self, other: &Tree<K, V, SIZE>) -> Ordering {
596 self.into_iter().cmp(other.into_iter())
597 }
598}
599
600impl<K, V, const SIZE: usize> Debug for Tree<K, V, SIZE>
601where
602 K: Debug + Ord + Clone,
603 V: Debug + Clone,
604{
605 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
606 f.debug_map().entries(self.into_iter()).finish()
607 }
608}
609
610impl<'a, Q, K, V, const SIZE: usize> Index<&'a Q> for Tree<K, V, SIZE>
611where
612 Q: Ord,
613 K: Ord + Clone + Borrow<Q>,
614 V: Clone,
615{
616 type Output = V;
617 fn index(&self, k: &Q) -> &V {
618 self.get(k).expect("element not found for key")
619 }
620}
621
622pub struct Iter<'a, R, Q, K, V, const SIZE: usize>
623where
624 Q: Ord + ?Sized,
625 R: RangeBounds<Q> + 'a,
626 K: 'a + Borrow<Q> + Ord + Clone,
627 V: 'a + Clone,
628{
629 q: PhantomData<Q>,
630 stack: ArrayVec<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>,
631 elts: Option<iter::Zip<slice::Iter<'a, K>, slice::Iter<'a, V>>>,
632 current: Option<&'a K>,
633 stack_rev: ArrayVec<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>,
634 elts_rev: Option<iter::Zip<slice::Iter<'a, K>, slice::Iter<'a, V>>>,
635 current_rev: Option<&'a K>,
636 bounds: R,
637}
638
639impl<'a, R, Q, K, V, const SIZE: usize> Iter<'a, R, Q, K, V, SIZE>
640where
641 Q: Ord + ?Sized,
642 R: RangeBounds<Q> + 'a,
643 K: 'a + Borrow<Q> + Ord + Clone,
644 V: 'a + Clone,
645{
646 fn any_elts_above_lbound(&self, n: &'a Node<K, V, SIZE>) -> bool {
648 let l = n.elts().len();
649 match self.bounds.start_bound() {
650 Bound::Unbounded => true,
651 Bound::Included(bound) => l == 0 || n.elts().key(l - 1).borrow() >= bound,
652 Bound::Excluded(bound) => l == 0 || n.elts().key(l - 1).borrow() > bound,
653 }
654 }
655
656 fn any_elts_below_ubound(&self, n: &'a Node<K, V, SIZE>) -> bool {
657 let l = n.elts().len();
658 match self.bounds.end_bound() {
659 Bound::Unbounded => true,
660 Bound::Included(bound) => l == 0 || n.elts().key(0).borrow() <= bound,
661 Bound::Excluded(bound) => l == 0 || n.elts().key(0).borrow() < bound,
662 }
663 }
664
665 fn any_elts_in_bounds(&self, n: &'a Node<K, V, SIZE>) -> bool {
666 self.any_elts_above_lbound(n) && self.any_elts_below_ubound(n)
667 }
668
669 fn above_lbound(&self, k: &'a K) -> bool {
670 match self.bounds.start_bound() {
671 Bound::Unbounded => true,
672 Bound::Included(bound) => k.borrow() >= bound,
673 Bound::Excluded(bound) => k.borrow() > bound,
674 }
675 }
676
677 fn below_ubound(&self, k: &'a K) -> bool {
678 match self.bounds.end_bound() {
679 Bound::Unbounded => true,
680 Bound::Included(bound) => k.borrow() <= bound,
681 Bound::Excluded(bound) => k.borrow() < bound,
682 }
683 }
684}
685
686impl<'a, R, Q, K, V, const SIZE: usize> Iterator for Iter<'a, R, Q, K, V, SIZE>
687where
688 Q: Ord + ?Sized,
689 R: RangeBounds<Q> + 'a,
690 K: 'a + Borrow<Q> + Ord + Clone,
691 V: 'a + Clone,
692{
693 type Item = (&'a K, &'a V);
694 fn next(&mut self) -> Option<Self::Item> {
695 loop {
696 loop {
697 let (k, v) = match &mut self.elts {
698 None => break,
699 Some(s) => match s.next() {
700 Some((k, v)) => (k, v),
701 None => break,
702 },
703 };
704 if let Some(back) = self.current_rev {
705 if k >= back {
706 return None;
707 }
708 }
709 if !self.below_ubound(k) {
710 return None;
711 }
712 self.current = Some(k);
713 if self.above_lbound(k) {
714 return Some((k, v));
715 }
716 }
717 if self.stack.is_empty() {
718 return None;
719 }
720 self.elts = None;
721 let top = self.stack.len() - 1;
722 let (visited, current) = self.stack[top];
723 if visited {
724 if self.any_elts_in_bounds(current) {
725 self.elts = Some(current.elts().into_iter());
726 }
727 self.stack.pop();
728 match current.right {
729 Tree::Empty => (),
730 Tree::Node(ref n) => {
731 if self.any_elts_below_ubound(n) || !n.left.is_empty() {
732 self.stack.push((false, n))
733 }
734 }
735 };
736 } else {
737 self.stack[top].0 = true;
738 match current.left {
739 Tree::Empty => (),
740 Tree::Node(ref n) => {
741 if self.any_elts_above_lbound(n) || !n.right.is_empty() {
742 self.stack.push((false, n))
743 }
744 }
745 }
746 }
747 }
748 }
749}
750
751impl<'a, R, Q, K, V, const SIZE: usize> DoubleEndedIterator for Iter<'a, R, Q, K, V, SIZE>
752where
753 Q: Ord + ?Sized,
754 R: RangeBounds<Q> + 'a,
755 K: 'a + Borrow<Q> + Ord + Clone,
756 V: 'a + Clone,
757{
758 fn next_back(&mut self) -> Option<Self::Item> {
759 loop {
760 loop {
761 let (k, v) = match &mut self.elts_rev {
762 &mut None => break,
763 &mut Some(ref mut s) => match s.next_back() {
764 None => break,
765 Some((k, v)) => (k, v),
766 },
767 };
768 if let Some(front) = self.current {
769 if k <= front {
770 return None;
771 }
772 }
773 if !self.above_lbound(k) {
774 return None;
775 }
776 self.current_rev = Some(k);
777 if self.below_ubound(k) {
778 return Some((k, v));
779 }
780 }
781 if self.stack_rev.is_empty() {
782 return None;
783 }
784 self.elts_rev = None;
785 let top = self.stack_rev.len() - 1;
786 let (visited, current) = self.stack_rev[top];
787 if visited {
788 if self.any_elts_in_bounds(current) {
789 self.elts_rev = Some(current.elts().into_iter());
790 }
791 self.stack_rev.pop();
792 match current.left {
793 Tree::Empty => (),
794 Tree::Node(ref n) => {
795 if self.any_elts_above_lbound(n) || !n.right.is_empty() {
796 self.stack_rev.push((false, n))
797 }
798 }
799 };
800 } else {
801 self.stack_rev[top].0 = true;
802 match current.right {
803 Tree::Empty => (),
804 Tree::Node(ref n) => {
805 if self.any_elts_below_ubound(n) || !n.left.is_empty() {
806 self.stack_rev.push((false, n))
807 }
808 }
809 }
810 }
811 }
812 }
813}
814
815pub struct IterMut<'a, R, Q, K, V, const SIZE: usize>
816where
817 Q: Ord + ?Sized,
818 R: RangeBounds<Q> + 'a,
819 K: 'a + Borrow<Q> + Ord + Clone,
820 V: 'a + Clone,
821{
822 q: PhantomData<Q>,
823 stack: ArrayVec<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>,
824 elts: Option<iter::Zip<slice::Iter<'a, K>, slice::IterMut<'a, V>>>,
825 current: Option<&'a K>,
826 stack_rev: ArrayVec<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>,
827 elts_rev: Option<iter::Zip<slice::Iter<'a, K>, slice::IterMut<'a, V>>>,
828 current_rev: Option<&'a K>,
829 bounds: R,
830}
831
832impl<'a, R, Q, K, V, const SIZE: usize> IterMut<'a, R, Q, K, V, SIZE>
833where
834 Q: Ord + ?Sized,
835 R: RangeBounds<Q> + 'a,
836 K: 'a + Borrow<Q> + Ord + Clone,
837 V: 'a + Clone,
838{
839 fn any_elts_above_lbound(&self, n: &'a NodeInner<K, V, SIZE>) -> bool {
841 let l = n.elts().len();
842 match self.bounds.start_bound() {
843 Bound::Unbounded => true,
844 Bound::Included(bound) => l == 0 || n.elts().key(l - 1).borrow() >= bound,
845 Bound::Excluded(bound) => l == 0 || n.elts().key(l - 1).borrow() > bound,
846 }
847 }
848
849 fn any_elts_below_ubound(&self, n: &'a NodeInner<K, V, SIZE>) -> bool {
850 let l = n.elts().len();
851 match self.bounds.end_bound() {
852 Bound::Unbounded => true,
853 Bound::Included(bound) => l == 0 || n.elts().key(0).borrow() <= bound,
854 Bound::Excluded(bound) => l == 0 || n.elts().key(0).borrow() < bound,
855 }
856 }
857
858 fn any_elts_in_bounds(&self, n: &'a NodeInner<K, V, SIZE>) -> bool {
859 self.any_elts_above_lbound(n) && self.any_elts_below_ubound(n)
860 }
861
862 fn above_lbound(&self, k: &'a K) -> bool {
863 match self.bounds.start_bound() {
864 Bound::Unbounded => true,
865 Bound::Included(bound) => k.borrow() >= bound,
866 Bound::Excluded(bound) => k.borrow() > bound,
867 }
868 }
869
870 fn below_ubound(&self, k: &'a K) -> bool {
871 match self.bounds.end_bound() {
872 Bound::Unbounded => true,
873 Bound::Included(bound) => k.borrow() <= bound,
874 Bound::Excluded(bound) => k.borrow() < bound,
875 }
876 }
877}
878
879impl<'a, R, Q, K, V, const SIZE: usize> Iterator for IterMut<'a, R, Q, K, V, SIZE>
880where
881 Q: Ord + ?Sized,
882 R: RangeBounds<Q> + 'a,
883 K: 'a + Borrow<Q> + Ord + Clone,
884 V: 'a + Clone,
885{
886 type Item = (&'a K, &'a mut V);
887 fn next(&mut self) -> Option<Self::Item> {
888 loop {
889 loop {
890 let (k, v) = match &mut self.elts {
891 &mut None => break,
892 &mut Some(ref mut s) => match s.next() {
893 Some((k, v)) => (k, v),
894 None => break,
895 },
896 };
897 if let Some(back) = self.current_rev {
898 if k >= back {
899 return None;
900 }
901 }
902 if !self.below_ubound(k) {
903 return None;
904 }
905 self.current = Some(k);
906 if self.above_lbound(k) {
907 return Some((k, v));
908 }
909 }
910 if self.stack.is_empty() {
911 return None;
912 }
913 self.elts = None;
914 let top = self.stack.len() - 1;
915 let (visited, current) = self.stack[top];
916 if visited {
917 if self.any_elts_in_bounds(unsafe { &*current }) {
918 self.elts =
919 Some((unsafe { (*current).make_mut().elts_mut() }).into_iter());
920 }
921 self.stack.pop();
922 match unsafe { &mut (*current).make_mut().right } {
923 Tree::Empty => (),
924 Tree::Node(ref mut n) => {
925 if self.any_elts_below_ubound(n) || !n.left.is_empty() {
926 self.stack.push((false, n))
927 }
928 }
929 };
930 } else {
931 self.stack[top].0 = true;
932 match unsafe { &mut (*current).make_mut().left } {
933 Tree::Empty => (),
934 Tree::Node(n) => {
935 if self.any_elts_above_lbound(n) || !n.right.is_empty() {
936 self.stack.push((false, n))
937 }
938 }
939 }
940 }
941 }
942 }
943}
944
945impl<'a, R, Q, K, V, const SIZE: usize> DoubleEndedIterator
946 for IterMut<'a, R, Q, K, V, SIZE>
947where
948 Q: Ord + ?Sized,
949 R: RangeBounds<Q> + 'a,
950 K: 'a + Borrow<Q> + Ord + Clone,
951 V: 'a + Clone,
952{
953 fn next_back(&mut self) -> Option<Self::Item> {
954 loop {
955 loop {
956 let (k, v) = match &mut self.elts_rev {
957 &mut None => break,
958 &mut Some(ref mut s) => match s.next_back() {
959 None => break,
960 Some((k, v)) => (k, v),
961 },
962 };
963 if let Some(front) = self.current {
964 if k <= front {
965 return None;
966 }
967 }
968 if !self.above_lbound(k) {
969 return None;
970 }
971 self.current_rev = Some(k);
972 if self.below_ubound(k) {
973 return Some((k, v));
974 }
975 }
976 if self.stack_rev.is_empty() {
977 return None;
978 }
979 self.elts_rev = None;
980 let top = self.stack_rev.len() - 1;
981 let (visited, current) = self.stack_rev[top];
982 if visited {
983 if self.any_elts_in_bounds(unsafe { &*current }) {
984 self.elts_rev =
985 Some((unsafe { (*current).make_mut().elts_mut() }).into_iter());
986 }
987 self.stack_rev.pop();
988 match unsafe { &mut (*current).make_mut().left } {
989 Tree::Empty => (),
990 Tree::Node(ref mut n) => {
991 if self.any_elts_above_lbound(n) || !n.right.is_empty() {
992 self.stack_rev.push((false, n))
993 }
994 }
995 };
996 } else {
997 self.stack_rev[top].0 = true;
998 match unsafe { &mut (*current).make_mut().right } {
999 Tree::Empty => (),
1000 Tree::Node(ref mut n) => {
1001 if self.any_elts_below_ubound(n) || !n.left.is_empty() {
1002 self.stack_rev.push((false, n))
1003 }
1004 }
1005 }
1006 }
1007 }
1008 }
1009}
1010
1011impl<'a, K, V, const SIZE: usize> IntoIterator for &'a Tree<K, V, SIZE>
1012where
1013 K: 'a + Ord + Clone,
1014 V: 'a + Clone,
1015{
1016 type Item = (&'a K, &'a V);
1017 type IntoIter = Iter<'a, RangeFull, K, K, V, SIZE>;
1018 fn into_iter(self) -> Self::IntoIter {
1019 self.range(..)
1020 }
1021}
1022
1023impl<K, V, const SIZE: usize> Tree<K, V, SIZE>
1024where
1025 K: Ord + Clone,
1026 V: Clone,
1027{
1028 pub(crate) fn new() -> Self {
1029 Tree::Empty
1030 }
1031
1032 pub(crate) fn downgrade(&self) -> WeakTree<K, V, SIZE> {
1033 match self {
1034 Tree::Empty => WeakTree::Empty,
1035 Tree::Node(n) => WeakTree::Node(n.downgrade()),
1036 }
1037 }
1038
1039 pub(crate) fn strong_count(&self) -> usize {
1040 match self {
1041 Tree::Empty => 0,
1042 Tree::Node(n) => Arc::strong_count(&n.0),
1043 }
1044 }
1045
1046 pub(crate) fn weak_count(&self) -> usize {
1047 match self {
1048 Tree::Empty => 0,
1049 Tree::Node(n) => Arc::weak_count(&n.0),
1050 }
1051 }
1052
1053 pub(crate) fn range<'a, Q, R>(&'a self, r: R) -> Iter<'a, R, Q, K, V, SIZE>
1054 where
1055 Q: Ord + ?Sized + 'a,
1056 K: Borrow<Q>,
1057 R: RangeBounds<Q> + 'a,
1058 {
1059 match self {
1060 &Tree::Empty => Iter {
1061 q: PhantomData,
1062 bounds: r,
1063 stack: ArrayVec::<_, MAX_DEPTH>::new(),
1064 elts: None,
1065 current: None,
1066 stack_rev: ArrayVec::<_, MAX_DEPTH>::new(),
1067 elts_rev: None,
1068 current_rev: None,
1069 },
1070 &Tree::Node(ref n) => {
1071 let mut stack =
1072 ArrayVec::<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>::new();
1073 let mut stack_rev =
1074 ArrayVec::<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>::new();
1075 stack.push((false, n));
1076 stack_rev.push((false, n));
1077 Iter {
1078 q: PhantomData,
1079 bounds: r,
1080 stack,
1081 elts: None,
1082 current: None,
1083 stack_rev,
1084 elts_rev: None,
1085 current_rev: None,
1086 }
1087 }
1088 }
1089 }
1090
1091 pub(crate) fn range_mut_cow<'a, Q, R>(
1092 &'a mut self,
1093 r: R,
1094 ) -> IterMut<'a, R, Q, K, V, SIZE>
1095 where
1096 Q: Ord + ?Sized + 'a,
1097 K: Borrow<Q>,
1098 R: RangeBounds<Q> + 'a,
1099 {
1100 match self {
1101 Tree::Empty => IterMut {
1102 q: PhantomData,
1103 bounds: r,
1104 stack: ArrayVec::<_, MAX_DEPTH>::new(),
1105 elts: None,
1106 current: None,
1107 stack_rev: ArrayVec::<_, MAX_DEPTH>::new(),
1108 elts_rev: None,
1109 current_rev: None,
1110 },
1111 Tree::Node(ref mut n) => {
1112 let mut stack =
1113 ArrayVec::<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>::new();
1114 let mut stack_rev =
1115 ArrayVec::<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>::new();
1116 stack.push((false, n));
1117 stack_rev.push((false, n));
1118 IterMut {
1119 q: PhantomData,
1120 bounds: r,
1121 stack,
1122 elts: None,
1123 current: None,
1124 stack_rev,
1125 elts_rev: None,
1126 current_rev: None,
1127 }
1128 }
1129 }
1130 }
1131
1132 pub(crate) fn iter_mut_cow<'a, Q>(
1133 &'a mut self,
1134 ) -> IterMut<'a, RangeFull, Q, K, V, SIZE>
1135 where
1136 Q: Ord + ?Sized + 'a,
1137 K: Borrow<Q>,
1138 {
1139 self.range_mut_cow(..)
1140 }
1141
1142 fn add_min_elts(&self, elts: &Chunk<K, V, SIZE>) -> Self {
1143 match self {
1144 Tree::Empty => Tree::create(&Tree::Empty, elts.clone(), &Tree::Empty),
1145 Tree::Node(ref n) => {
1146 Tree::bal(&n.left.add_min_elts(elts), n.elts().clone(), &n.right)
1147 }
1148 }
1149 }
1150
1151 fn add_max_elts(&self, elts: &Chunk<K, V, SIZE>) -> Self {
1152 match self {
1153 Tree::Empty => Tree::create(&Tree::Empty, elts.clone(), &Tree::Empty),
1154 Tree::Node(ref n) => {
1155 Tree::bal(&n.left, n.elts().clone(), &n.right.add_max_elts(elts))
1156 }
1157 }
1158 }
1159
1160 fn join(
1164 l: &Tree<K, V, SIZE>,
1165 elts: &Chunk<K, V, SIZE>,
1166 r: &Tree<K, V, SIZE>,
1167 ) -> Self {
1168 match (l, r) {
1169 (Tree::Empty, _) => r.add_min_elts(elts),
1170 (_, Tree::Empty) => l.add_max_elts(elts),
1171 (Tree::Node(ref ln), Tree::Node(ref rn)) => {
1172 let (ln_height, rn_height) = (ln.height(), rn.height());
1173 if ln_height > rn_height + 2 {
1174 Tree::bal(
1175 &ln.left,
1176 ln.elts().clone(),
1177 &Tree::join(&ln.right, elts, r),
1178 )
1179 } else if rn_height > ln_height + 2 {
1180 Tree::bal(
1181 &Tree::join(l, elts, &rn.left),
1182 rn.elts().clone(),
1183 &rn.right,
1184 )
1185 } else {
1186 Tree::create(l, elts.clone(), r)
1187 }
1188 }
1189 }
1190 }
1191
1192 fn split(&self, vmin: &K, vmax: &K) -> (Self, Option<Chunk<K, V, SIZE>>, Self) {
1198 match self {
1199 Tree::Empty => (Tree::Empty, None, Tree::Empty),
1200 Tree::Node(ref n) => {
1201 if vmax < n.min_key() {
1202 let (ll, inter, rl) = n.left.split(vmin, vmax);
1203 (ll, inter, Tree::join(&rl, n.elts(), &n.right))
1204 } else if vmin > n.max_key() {
1205 let (lr, inter, rr) = n.right.split(vmin, vmax);
1206 (Tree::join(&n.left, n.elts(), &lr), inter, rr)
1207 } else {
1208 (n.left.clone(), Some(n.elts().clone()), n.right.clone())
1209 }
1210 }
1211 }
1212 }
1213
1214 fn merge_root_to<F>(
1218 from: &Tree<K, V, SIZE>,
1219 to: &Tree<K, V, SIZE>,
1220 f: &mut F,
1221 ) -> (Self, Self)
1222 where
1223 F: FnMut(&K, &V, &V) -> Option<V>,
1224 {
1225 match (from, to) {
1226 (Tree::Empty, to) => (Tree::Empty, to.clone()),
1227 (Tree::Node(ref n), to) => {
1228 let to =
1229 to.update_chunk(n.elts().to_vec(), &mut |k0, v0, cur| match cur {
1230 None => Some((k0, v0)),
1231 Some((_, v1)) => f(&k0, &v0, v1).map(|v| (k0, v)),
1232 });
1233 if n.height() == 1 {
1234 (Tree::Empty, to)
1235 } else {
1236 match n.right {
1237 Tree::Empty => (n.left.clone(), to),
1238 Tree::Node(_) => {
1239 let elts = n.right.min_elts().unwrap();
1240 let right = n.right.remove_min_elts();
1241 (Tree::join(&n.left, elts, &right), to)
1242 }
1243 }
1244 }
1245 }
1246 }
1247 }
1248
1249 pub(crate) fn union<F>(
1253 t0: &Tree<K, V, SIZE>,
1254 t1: &Tree<K, V, SIZE>,
1255 f: &mut F,
1256 ) -> Self
1257 where
1258 F: FnMut(&K, &V, &V) -> Option<V>,
1259 {
1260 match (t0, t1) {
1261 (Tree::Empty, Tree::Empty) => Tree::Empty,
1262 (Tree::Empty, t1) => t1.clone(),
1263 (t0, Tree::Empty) => t0.clone(),
1264 (Tree::Node(ref n0), Tree::Node(ref n1)) => {
1265 if n0.height() > n1.height() {
1266 match t1.split(n0.min_key(), n0.max_key()) {
1267 (_, Some(_), _) => {
1268 let (t0, t1) = Tree::merge_root_to(&t0, &t1, f);
1269 Tree::union(&t0, &t1, f)
1270 }
1271 (l1, None, r1) => Tree::join(
1272 &Tree::union(&n0.left, &l1, f),
1273 n0.elts(),
1274 &Tree::union(&n0.right, &r1, f),
1275 ),
1276 }
1277 } else {
1278 match t0.split(n1.min_key(), n1.max_key()) {
1279 (_, Some(_), _) => {
1280 let (t1, t0) = Tree::merge_root_to(&t1, &t0, f);
1281 Tree::union(&t0, &t1, f)
1282 }
1283 (l0, None, r0) => Tree::join(
1284 &Tree::union(&l0, &n1.left, f),
1285 n1.elts(),
1286 &Tree::union(&r0, &n1.right, f),
1287 ),
1288 }
1289 }
1290 }
1291 }
1292 }
1293
1294 fn intersect_int<F>(
1295 t0: &Tree<K, V, SIZE>,
1296 t1: &Tree<K, V, SIZE>,
1297 r: &mut Vec<(K, V)>,
1298 f: &mut F,
1299 ) where
1300 F: FnMut(&K, &V, &V) -> Option<V>,
1301 {
1302 match (t0, t1) {
1303 (Tree::Empty, _) => (),
1304 (_, Tree::Empty) => (),
1305 (Tree::Node(ref n0), t1) => match t1.split(n0.min_key(), n0.max_key()) {
1306 (l1, None, r1) => {
1307 Tree::intersect_int(&n0.left, &l1, r, f);
1308 Tree::intersect_int(&n0.right, &r1, r, f);
1309 }
1310 (l1, Some(elts), r1) if elts.len() == 0 => {
1311 Tree::intersect_int(&n0.left, &l1, r, f);
1312 Tree::intersect_int(&n0.right, &r1, r, f);
1313 }
1314 (l1, Some(elts), r1) => {
1315 let (min_k, max_k) = elts.min_max_key().unwrap();
1316 Chunk::intersect(n0.elts(), &elts, r, f);
1317 if n0.min_key() < &min_k && n0.max_key() > &max_k {
1318 Tree::intersect_int(t0, &Tree::concat(&l1, &r1), r, f)
1319 } else if n0.min_key() >= &min_k && n0.max_key() <= &max_k {
1320 let t0 = Tree::concat(&n0.left, &n0.right);
1321 let t1 = Tree::join(&l1, &elts, &r1);
1322 Tree::intersect_int(&t0, &t1, r, f);
1323 } else if n0.min_key() < &min_k {
1324 let tl = Tree::join(&n0.left, n0.elts(), &Tree::Empty);
1325 Tree::intersect_int(&tl, &l1, r, f);
1326 let tr = Tree::join(&Tree::Empty, &elts, &r1);
1327 Tree::intersect_int(&n0.right, &tr, r, f);
1328 } else {
1329 let tl = Tree::join(&l1, &elts, &Tree::Empty);
1330 Tree::intersect_int(&tl, &n0.left, r, f);
1331 let tr = Tree::join(&Tree::Empty, n0.elts(), &n0.right);
1332 Tree::intersect_int(&r1, &tr, r, f);
1333 }
1334 }
1335 },
1336 }
1337 }
1338
1339 pub(crate) fn intersect<F>(
1340 t0: &Tree<K, V, SIZE>,
1341 t1: &Tree<K, V, SIZE>,
1342 f: &mut F,
1343 ) -> Self
1344 where
1345 F: FnMut(&K, &V, &V) -> Option<V>,
1346 {
1347 let mut r = Vec::new();
1348 Tree::intersect_int(t0, t1, &mut r, f);
1349 Tree::Empty.insert_many(r.into_iter())
1350 }
1351
1352 pub(crate) fn diff<F>(t0: &Tree<K, V, SIZE>, t1: &Tree<K, V, SIZE>, f: &mut F) -> Self
1353 where
1354 F: FnMut(&K, &V, &V) -> Option<V>,
1355 {
1356 let mut actions = Vec::new();
1357 Tree::intersect_int(t0, t1, &mut Vec::new(), &mut |k, v0, v1| {
1358 actions.push((k.clone(), f(k, v0, v1)));
1359 None
1360 });
1361 t0.update_many(actions, &mut |k, v, _| v.map(|v| (k, v)))
1362 }
1363
1364 fn is_empty(&self) -> bool {
1365 match self {
1366 Tree::Empty => true,
1367 Tree::Node(..) => false,
1368 }
1369 }
1370
1371 pub(crate) fn len(&self) -> usize {
1372 match self {
1373 Tree::Empty => 0,
1374 Tree::Node(n) => {
1375 let size_of_children = (n.height_and_size & 0x00ffffff_ffffffff) as usize;
1379 n.elts().len() + size_of_children
1380 }
1381 }
1382 }
1383
1384 fn height(&self) -> u8 {
1385 match self {
1386 Tree::Empty => 0,
1387 Tree::Node(ref n) => n.height(),
1388 }
1389 }
1390
1391 #[cfg(feature = "pool")]
1392 fn create(
1393 l: &Tree<K, V, SIZE>,
1394 elts: Chunk<K, V, SIZE>,
1395 r: &Tree<K, V, SIZE>,
1396 ) -> Self {
1397 let (min_key, max_key) = elts.min_max_key().unwrap();
1398 let height_and_size =
1399 pack_height_and_size(1 + max(l.height(), r.height()), l.len() + r.len());
1400 let mut t = take::<Node<K, V, SIZE>>();
1401 let inner = Arc::get_mut(&mut t.0).unwrap();
1402 inner.elts = MaybeUninit::new(elts);
1403 inner.min_key = MaybeUninit::new(min_key);
1404 inner.max_key = MaybeUninit::new(max_key);
1405 inner.left = l.clone();
1406 inner.right = r.clone();
1407 inner.height_and_size = height_and_size;
1408 Tree::Node(t)
1409 }
1410
1411 #[cfg(not(feature = "pool"))]
1412 fn create(
1413 l: &Tree<K, V, SIZE>,
1414 elts: Chunk<K, V, SIZE>,
1415 r: &Tree<K, V, SIZE>,
1416 ) -> Self {
1417 let (min_key, max_key) = elts.min_max_key().unwrap();
1418 let height_and_size =
1419 pack_height_and_size(1 + max(l.height(), r.height()), l.len() + r.len());
1420 let n = NodeInner {
1421 elts,
1422 min_key,
1423 max_key,
1424 left: l.clone(),
1425 right: r.clone(),
1426 height_and_size,
1427 };
1428 Tree::Node(Node(Arc::new(n)))
1429 }
1430
1431 fn in_bal(l: &Tree<K, V, SIZE>, r: &Tree<K, V, SIZE>) -> bool {
1432 let (hl, hr) = (l.height(), r.height());
1433 (hl <= hr.saturating_add(2)) && (hr <= hl.saturating_add(2))
1434 }
1435
1436 fn compact(self) -> Self {
1437 match self {
1438 Tree::Empty => self,
1439 Tree::Node(ref tn) => {
1440 let len = tn.elts().len();
1441 if len > SIZE >> 1 {
1442 self
1443 } else {
1444 match tn.right.min_elts() {
1445 None => self,
1446 Some(chunk) => {
1447 let n = SIZE - len;
1448 let to_add =
1449 chunk.into_iter().map(|(k, v)| (k.clone(), v.clone()));
1450 let overflow = chunk
1451 .into_iter()
1452 .skip(n)
1453 .map(|(k, v)| (k.clone(), v.clone()));
1454 let elts = tn.elts().append(to_add);
1455 let t =
1456 Tree::bal(&tn.left, elts, &tn.right.remove_min_elts());
1457 if n >= chunk.len() {
1458 t
1459 } else {
1460 t.insert_many(overflow)
1461 }
1462 }
1463 }
1464 }
1465 }
1466 }
1467 }
1468
1469 fn bal(l: &Tree<K, V, SIZE>, elts: Chunk<K, V, SIZE>, r: &Tree<K, V, SIZE>) -> Self {
1470 let (hl, hr) = (l.height(), r.height());
1471 if hl > hr.saturating_add(2) {
1472 match *l {
1473 Tree::Empty => panic!("tree heights wrong"),
1474 Tree::Node(ref ln) => {
1475 if ln.left.height() >= ln.right.height() {
1476 Tree::create(
1477 &ln.left,
1478 ln.elts().clone(),
1479 &Tree::create(&ln.right, elts, r),
1480 )
1481 .compact()
1482 } else {
1483 match ln.right {
1484 Tree::Empty => panic!("tree heights wrong"),
1485 Tree::Node(ref lrn) => Tree::create(
1486 &Tree::create(&ln.left, ln.elts().clone(), &lrn.left),
1487 lrn.elts().clone(),
1488 &Tree::create(&lrn.right, elts, r),
1489 )
1490 .compact(),
1491 }
1492 }
1493 }
1494 }
1495 } else if hr > hl.saturating_add(2) {
1496 match *r {
1497 Tree::Empty => panic!("tree heights are wrong"),
1498 Tree::Node(ref rn) => {
1499 if rn.right.height() >= rn.left.height() {
1500 Tree::create(
1501 &Tree::create(l, elts, &rn.left),
1502 rn.elts().clone(),
1503 &rn.right,
1504 )
1505 .compact()
1506 } else {
1507 match rn.left {
1508 Tree::Empty => panic!("tree heights are wrong"),
1509 Tree::Node(ref rln) => Tree::create(
1510 &Tree::create(l, elts, &rln.left),
1511 rln.elts().clone(),
1512 &Tree::create(&rln.right, rn.elts().clone(), &rn.right),
1513 )
1514 .compact(),
1515 }
1516 }
1517 }
1518 }
1519 } else {
1520 Tree::create(l, elts, r).compact()
1521 }
1522 }
1523
1524 fn update_chunk<Q, D, F>(&self, chunk: Vec<(Q, D)>, f: &mut F) -> Self
1525 where
1526 Q: Ord,
1527 K: Borrow<Q>,
1528 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1529 {
1530 if chunk.len() == 0 {
1531 return self.clone();
1532 }
1533 match self {
1534 &Tree::Empty => {
1535 let chunk = Chunk::create_with(chunk, f);
1536 if chunk.len() == 0 {
1537 Tree::Empty
1538 } else {
1539 Tree::create(&Tree::Empty, chunk, &Tree::Empty)
1540 }
1541 }
1542 &Tree::Node(ref tn) => {
1543 let leaf = match (&tn.left, &tn.right) {
1544 (&Tree::Empty, &Tree::Empty) => true,
1545 (_, _) => false,
1546 };
1547 match tn.elts().update_chunk(chunk, leaf, f) {
1548 UpdateChunk::Updated {
1549 elts,
1550 update_left,
1551 update_right,
1552 overflow_right,
1553 } => {
1554 let l = tn.left.update_chunk(update_left, f);
1555 let r = tn.right.insert_chunk(overflow_right);
1556 let r = r.update_chunk(update_right, f);
1557 Tree::bal(&l, elts, &r)
1558 }
1559 UpdateChunk::Removed {
1560 not_done,
1561 update_left,
1562 update_right,
1563 } => {
1564 let l = tn.left.update_chunk(update_left, f);
1565 let r = tn.right.update_chunk(update_right, f);
1566 let t = Tree::concat(&l, &r);
1567 t.update_chunk(not_done, f)
1568 }
1569 UpdateChunk::UpdateLeft(chunk) => {
1570 let l = tn.left.update_chunk(chunk, f);
1571 Tree::bal(&l, tn.elts().clone(), &tn.right)
1572 }
1573 UpdateChunk::UpdateRight(chunk) => {
1574 let r = tn.right.update_chunk(chunk, f);
1575 Tree::bal(&tn.left, tn.elts().clone(), &r)
1576 }
1577 }
1578 }
1579 }
1580 }
1581
1582 fn insert_chunk(&self, chunk: Vec<(K, V)>) -> Self {
1583 self.update_chunk(chunk, &mut |k, v, _| Some((k, v)))
1584 }
1585
1586 pub(crate) fn update_many<Q, D, E, F>(&self, elts: E, f: &mut F) -> Self
1587 where
1588 E: IntoIterator<Item = (Q, D)>,
1589 Q: Ord,
1590 K: Borrow<Q>,
1591 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1592 {
1593 let mut elts = {
1594 let mut v = elts.into_iter().collect::<Vec<(Q, D)>>();
1595 v.sort_by(|(ref k0, _), (ref k1, _)| k0.cmp(k1));
1596 v.dedup_by(|t0, t1| t0.0 == t1.0);
1597 v
1598 };
1599 let mut t = self.clone();
1600 while elts.len() > 0 {
1601 let chunk = elts.drain(0..min(SIZE, elts.len())).collect::<Vec<_>>();
1602 t = t.update_chunk(chunk, f)
1603 }
1604 t
1605 }
1606
1607 pub(crate) fn insert_many<E: IntoIterator<Item = (K, V)>>(&self, elts: E) -> Self {
1608 self.update_many(elts, &mut |k, v, _| Some((k, v)))
1609 }
1610
1611 pub(crate) fn update_cow<Q, D, F>(&mut self, q: Q, d: D, f: &mut F) -> Option<V>
1612 where
1613 Q: Ord,
1614 K: Borrow<Q>,
1615 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1616 {
1617 match self {
1618 Tree::Empty => match f(q, d, None) {
1619 None => None,
1620 Some((k, v)) => {
1621 *self =
1622 Tree::create(&Tree::Empty, Chunk::singleton(k, v), &Tree::Empty);
1623 None
1624 }
1625 },
1626 Tree::Node(ref mut tn) => {
1627 let tn = tn.make_mut();
1629 let leaf = match (&tn.left, &tn.right) {
1630 (&Tree::Empty, &Tree::Empty) => true,
1631 (_, _) => false,
1632 };
1633 match tn.elts_mut().update_mut(q, d, leaf, f) {
1634 MutUpdate::UpdateLeft(k, d) => {
1635 let prev = tn.left.update_cow(k, d, f);
1636 if !Tree::in_bal(&tn.left, &tn.right) {
1637 *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right)
1638 } else {
1639 tn.mutated();
1640 }
1641 prev
1642 }
1643 MutUpdate::UpdateRight(k, d) => {
1644 let prev = tn.right.update_cow(k, d, f);
1645 if !Tree::in_bal(&tn.left, &tn.right) {
1646 *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right)
1647 } else {
1648 tn.mutated();
1649 }
1650 prev
1651 }
1652 MutUpdate::Updated { overflow, previous } => match overflow {
1653 None => {
1654 if tn.elts().len() > 0 {
1655 tn.mutated();
1656 previous
1657 } else {
1658 *self = Tree::concat(&tn.left, &tn.right);
1659 previous
1660 }
1661 }
1662 Some((ovk, ovv)) => {
1663 let _ = tn.right.insert_cow(ovk, ovv);
1664 if tn.elts().len() > 0 {
1665 if !Tree::in_bal(&tn.left, &tn.right) {
1666 *self =
1667 Tree::bal(&tn.left, tn.elts().clone(), &tn.right);
1668 previous
1669 } else {
1670 tn.mutated();
1671 previous
1672 }
1673 } else {
1674 *self = Tree::concat(&tn.left, &tn.right);
1676 previous
1677 }
1678 }
1679 },
1680 }
1681 }
1682 }
1683 }
1684
1685 pub(crate) fn update<Q, D, F>(&self, q: Q, d: D, f: &mut F) -> (Self, Option<V>)
1686 where
1687 Q: Ord,
1688 K: Borrow<Q>,
1689 F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1690 {
1691 match self {
1692 Tree::Empty => match f(q, d, None) {
1693 None => (self.clone(), None),
1694 Some((k, v)) => (
1695 Tree::create(&Tree::Empty, Chunk::singleton(k, v), &Tree::Empty),
1696 None,
1697 ),
1698 },
1699 Tree::Node(ref tn) => {
1700 let leaf = match (&tn.left, &tn.right) {
1701 (&Tree::Empty, &Tree::Empty) => true,
1702 (_, _) => false,
1703 };
1704 match tn.elts().update(q, d, leaf, f) {
1705 Update::UpdateLeft(k, d) => {
1706 let (l, prev) = tn.left.update(k, d, f);
1707 (Tree::bal(&l, tn.elts().clone(), &tn.right), prev)
1708 }
1709 Update::UpdateRight(k, d) => {
1710 let (r, prev) = tn.right.update(k, d, f);
1711 (Tree::bal(&tn.left, tn.elts().clone(), &r), prev)
1712 }
1713 Update::Updated {
1714 elts,
1715 overflow,
1716 previous,
1717 } => match overflow {
1718 None => {
1719 if elts.len() == 0 {
1720 (Tree::concat(&tn.left, &tn.right), previous)
1721 } else {
1722 (Tree::create(&tn.left, elts, &tn.right), previous)
1723 }
1724 }
1725 Some((ovk, ovv)) => {
1726 let (r, _) = tn.right.insert(ovk, ovv);
1727 if elts.len() == 0 {
1728 (Tree::concat(&tn.left, &r), previous)
1729 } else {
1730 (Tree::bal(&tn.left, elts, &r), previous)
1731 }
1732 }
1733 },
1734 }
1735 }
1736 }
1737 }
1738
1739 pub(crate) fn insert(&self, k: K, v: V) -> (Self, Option<V>) {
1740 self.update(k, v, &mut |k, v, _| Some((k, v)))
1741 }
1742
1743 pub(crate) fn insert_cow(&mut self, k: K, v: V) -> Option<V> {
1744 self.update_cow(k, v, &mut |k, v, _| Some((k, v)))
1745 }
1746
1747 fn min_elts<'a>(&'a self) -> Option<&'a Chunk<K, V, SIZE>> {
1748 match self {
1749 Tree::Empty => None,
1750 Tree::Node(ref tn) => match tn.left {
1751 Tree::Empty => Some(tn.elts()),
1752 Tree::Node(_) => tn.left.min_elts(),
1753 },
1754 }
1755 }
1756
1757 fn remove_min_elts(&self) -> Self {
1758 match self {
1759 Tree::Empty => panic!("remove min elt"),
1760 Tree::Node(ref tn) => match tn.left {
1761 Tree::Empty => tn.right.clone(),
1762 Tree::Node(_) => {
1763 Tree::bal(&tn.left.remove_min_elts(), tn.elts().clone(), &tn.right)
1764 }
1765 },
1766 }
1767 }
1768
1769 fn concat(l: &Tree<K, V, SIZE>, r: &Tree<K, V, SIZE>) -> Tree<K, V, SIZE> {
1770 match (l, r) {
1771 (Tree::Empty, _) => r.clone(),
1772 (_, Tree::Empty) => l.clone(),
1773 (_, _) => {
1774 let elts = match r.min_elts() {
1775 Some(e) => e,
1776 None => &Chunk::empty(), };
1778 Tree::bal(l, elts.clone(), &r.remove_min_elts())
1779 }
1780 }
1781 }
1782
1783 pub(crate) fn remove<Q: ?Sized + Ord>(&self, k: &Q) -> (Self, Option<V>)
1784 where
1785 K: Borrow<Q>,
1786 {
1787 match self {
1788 &Tree::Empty => (Tree::Empty, None),
1789 &Tree::Node(ref tn) => match tn.elts().get(k) {
1790 Loc::NotPresent(_) => (self.clone(), None),
1791 Loc::Here(i) => {
1792 let p = tn.elts().val(i).clone();
1793 let elts = tn.elts().remove_elt_at(i);
1794 if elts.len() == 0 {
1795 (Tree::concat(&tn.left, &tn.right), Some(p))
1796 } else {
1797 (Tree::create(&tn.left, elts, &tn.right), Some(p))
1798 }
1799 }
1800 Loc::InLeft => {
1801 let (l, p) = tn.left.remove(k);
1802 (Tree::bal(&l, tn.elts().clone(), &tn.right), p)
1803 }
1804 Loc::InRight => {
1805 let (r, p) = tn.right.remove(k);
1806 (Tree::bal(&tn.left, tn.elts().clone(), &r), p)
1807 }
1808 },
1809 }
1810 }
1811
1812 pub(crate) fn remove_cow<Q: ?Sized + Ord>(&mut self, k: &Q) -> Option<V>
1813 where
1814 K: Borrow<Q>,
1815 {
1816 match self {
1817 Tree::Empty => None,
1818 Tree::Node(ref mut tn) => {
1819 let tn = tn.make_mut();
1821 match tn.elts().get(k) {
1822 Loc::NotPresent(_) => None,
1823 Loc::Here(i) => {
1824 let (_, p) = tn.elts_mut().remove_elt_at_mut(i);
1825 if tn.elts().len() == 0 {
1826 *self = Tree::concat(&tn.left, &tn.right);
1827 Some(p)
1828 } else {
1829 tn.mutated();
1830 Some(p)
1831 }
1832 }
1833 Loc::InLeft => {
1834 let p = tn.left.remove_cow(k);
1835 if !Tree::in_bal(&tn.left, &tn.right) {
1836 *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right);
1837 } else {
1838 tn.mutated()
1839 }
1840 p
1841 }
1842 Loc::InRight => {
1843 let p = tn.right.remove_cow(k);
1844 if !Tree::in_bal(&tn.left, &tn.right) {
1845 *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right);
1846 } else {
1847 tn.mutated()
1848 }
1849 p
1850 }
1851 }
1852 }
1853 }
1854 }
1855
1856 fn get_gen<'a, Q, F, R>(&'a self, k: &Q, f: F) -> Option<R>
1861 where
1862 Q: ?Sized + Ord,
1863 K: Borrow<Q>,
1864 F: FnOnce(&'a Chunk<K, V, SIZE>, usize) -> R,
1865 R: 'a,
1866 {
1867 match self {
1868 Tree::Empty => None,
1869 Tree::Node(n) => {
1870 let mut tn = n;
1871 loop {
1872 match (k.cmp(tn.min_key().borrow()), k.cmp(tn.max_key().borrow())) {
1873 (Ordering::Less, _) => match tn.left {
1874 Tree::Empty => break None,
1875 Tree::Node(ref n) => tn = n,
1876 },
1877 (_, Ordering::Greater) => match tn.right {
1878 Tree::Empty => break None,
1879 Tree::Node(ref n) => tn = n,
1880 },
1881 (_, _) => {
1882 let e = tn.elts();
1883 break e.get_local(k).map(|i| f(e, i));
1884 }
1885 }
1886 }
1887 }
1888 }
1889 }
1890
1891 pub(crate) fn get<'a, Q>(&'a self, k: &Q) -> Option<&'a V>
1892 where
1893 Q: ?Sized + Ord,
1894 K: Borrow<Q>,
1895 {
1896 self.get_gen(k, |e, i| e.val(i))
1897 }
1898
1899 pub(crate) fn get_key<'a, Q>(&'a self, k: &Q) -> Option<&'a K>
1900 where
1901 Q: ?Sized + Ord,
1902 K: Borrow<Q>,
1903 {
1904 self.get_gen(k, |e, i| e.key(i))
1905 }
1906
1907 pub(crate) fn get_full<'a, Q>(&'a self, k: &Q) -> Option<(&'a K, &'a V)>
1908 where
1909 Q: ?Sized + Ord,
1910 K: Borrow<Q>,
1911 {
1912 self.get_gen(k, |e, i| e.kv(i))
1913 }
1914
1915 pub(crate) fn get_mut_cow<'a, Q>(&'a mut self, k: &Q) -> Option<&'a mut V>
1916 where
1917 Q: ?Sized + Ord,
1918 K: Borrow<Q>,
1919 {
1920 match self {
1921 Tree::Empty => None,
1922 Tree::Node(tn) => {
1923 let tn = tn.make_mut();
1924 match (k.cmp(tn.min_key().borrow()), k.cmp(tn.max_key().borrow())) {
1925 (Ordering::Less, _) => tn.left.get_mut_cow(k),
1926 (_, Ordering::Greater) => tn.right.get_mut_cow(k),
1927 (_, _) => match tn.elts().get_local(k) {
1928 Some(i) => Some(tn.elts_mut().val_mut(i)),
1929 None => None,
1930 },
1931 }
1932 }
1933 }
1934 }
1935
1936 pub(crate) fn get_or_insert_cow<'a, F>(&'a mut self, k: K, f: F) -> &'a mut V
1937 where
1938 F: FnOnce() -> V,
1939 {
1940 match self.get_mut_cow(&k).map(|v| v as *mut V) {
1941 Some(v) => unsafe { &mut *v },
1942 None => {
1943 self.insert_cow(k.clone(), f());
1944 self.get_mut_cow(&k).unwrap()
1945 }
1946 }
1947 }
1948
1949 pub(crate) fn root(&self) -> Option<NodeRef<'_, K, V, SIZE>> {
1950 match self {
1951 Tree::Empty => None,
1952 Tree::Node(n) => Some(NodeRef(n)),
1953 }
1954 }
1955
1956 pub(crate) fn from_root(root: Option<NodeHandle<K, V, SIZE>>) -> Self {
1957 match root {
1958 None => Tree::Empty,
1959 Some(h) => Tree::Node(h.0),
1960 }
1961 }
1962}
1963
1964impl<K, V, const SIZE: usize> Tree<K, V, SIZE>
1965where
1966 K: Ord + Clone + Debug,
1967 V: Clone + Debug,
1968{
1969 #[allow(dead_code)]
1970 pub(crate) fn invariant(&self) -> () {
1971 fn in_range<K, V, const SIZE: usize>(
1972 lower: Option<&K>,
1973 upper: Option<&K>,
1974 elts: &Chunk<K, V, SIZE>,
1975 ) -> bool
1976 where
1977 K: Ord + Clone + Debug,
1978 V: Clone + Debug,
1979 {
1980 (match lower {
1981 None => true,
1982 Some(lower) => elts
1983 .into_iter()
1984 .all(|(k, _)| lower.cmp(k) == Ordering::Less),
1985 }) && (match upper {
1986 None => true,
1987 Some(upper) => elts
1988 .into_iter()
1989 .all(|(k, _)| upper.cmp(k) == Ordering::Greater),
1990 })
1991 }
1992
1993 fn sorted<K, V, const SIZE: usize>(elts: &Chunk<K, V, SIZE>) -> bool
1994 where
1995 K: Ord + Clone + Debug,
1996 V: Clone + Debug,
1997 {
1998 if elts.len() == 1 {
1999 true
2000 } else {
2001 for i in 0..(elts.len() - 1) {
2002 match elts.key(i).cmp(&elts.key(i + 1)) {
2003 Ordering::Greater => return false,
2004 Ordering::Less => (),
2005 Ordering::Equal => panic!("duplicates found: {:#?}", elts),
2006 }
2007 }
2008 true
2009 }
2010 }
2011
2012 fn check<K, V, const SIZE: usize>(
2013 t: &Tree<K, V, SIZE>,
2014 lower: Option<&K>,
2015 upper: Option<&K>,
2016 len: usize,
2017 ) -> (u8, usize)
2018 where
2019 K: Ord + Clone + Debug,
2020 V: Clone + Debug,
2021 {
2022 match *t {
2023 Tree::Empty => (0, len),
2024 Tree::Node(ref tn) => {
2025 if !in_range(lower, upper, tn.elts()) {
2026 panic!("tree invariant violated lower\n{:#?}\n\nupper\n{:#?}\n\nelts\n{:#?}\n\ntree\n{:#?}",
2027 lower, upper, &tn.elts, t)
2028 };
2029 if !sorted(tn.elts()) {
2030 panic!("elements isn't sorted")
2031 };
2032 let (thl, len) =
2033 check(&tn.left, lower, tn.elts().min_elt().map(|(k, _)| k), len);
2034 let (thr, len) =
2035 check(&tn.right, tn.elts().max_elt().map(|(k, _)| k), upper, len);
2036 let th = max(thl, thr).saturating_add(1);
2037 let (hl, hr) = (tn.left.height(), tn.right.height());
2038 let ub = max(hl, hr) - min(hl, hr);
2039 if thl != hl {
2040 panic!("left node height is wrong")
2041 };
2042 if thr != hr {
2043 panic!("right node height is wrong")
2044 };
2045 let h = t.height();
2046 if th != h {
2047 panic!("node height is wrong {} vs {}", th, h)
2048 };
2049 if ub > 2 {
2050 panic!("tree is unbalanced {:#?} tree: {:#?}", ub, t)
2051 };
2052 (th, len + tn.elts().len())
2053 }
2054 }
2055 }
2056
2057 let (_height, tlen) = check(self, None, None, 0);
2059 let len = self.len();
2060 if len != tlen {
2061 panic!("len is wrong {} vs {}", len, tlen)
2062 }
2063 }
2064}
2065
2066pub struct NodeRef<'a, K: Ord + Clone, V: Clone, const SIZE: usize>(&'a Node<K, V, SIZE>);
2074
2075impl<'a, K: Ord + Clone, V: Clone, const SIZE: usize> NodeRef<'a, K, V, SIZE> {
2076 pub fn identity(&self) -> usize {
2079 Arc::as_ptr(self.0.arc()) as *const () as usize
2080 }
2081
2082 pub fn keep(&self) -> NodeHandle<K, V, SIZE> {
2084 NodeHandle(self.0.clone())
2085 }
2086
2087 pub fn len(&self) -> usize {
2089 self.0.elts().len()
2090 }
2091
2092 pub fn pairs(&self) -> impl Iterator<Item = (&'a K, &'a V)> + 'a {
2094 let node: &'a Node<K, V, SIZE> = self.0;
2095 let chunk: &'a Chunk<K, V, SIZE> = node.elts();
2096 (0..chunk.len()).map(move |i| chunk.kv(i))
2097 }
2098
2099 pub fn left(&self) -> Option<NodeRef<'a, K, V, SIZE>> {
2100 let node: &'a Node<K, V, SIZE> = self.0;
2101 node.left.root()
2102 }
2103
2104 pub fn right(&self) -> Option<NodeRef<'a, K, V, SIZE>> {
2105 let node: &'a Node<K, V, SIZE> = self.0;
2106 node.right.root()
2107 }
2108}
2109
2110pub struct NodeHandle<K: Ord + Clone, V: Clone, const SIZE: usize>(Node<K, V, SIZE>);
2115
2116impl<K: Ord + Clone, V: Clone, const SIZE: usize> Debug for NodeHandle<K, V, SIZE> {
2117 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
2118 let v = self.view();
2119 write!(f, "NodeHandle({:#x}, {} pairs)", v.identity(), v.len())
2120 }
2121}
2122
2123impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for NodeHandle<K, V, SIZE> {
2124 fn clone(&self) -> Self {
2125 Self(self.0.clone())
2126 }
2127}
2128
2129impl<K: Ord + Clone, V: Clone, const SIZE: usize> NodeHandle<K, V, SIZE> {
2130 pub fn view(&self) -> NodeRef<'_, K, V, SIZE> {
2131 NodeRef(&self.0)
2132 }
2133
2134 pub unsafe fn create<I: IntoIterator<Item = (K, V)>>(
2145 left: Option<Self>,
2146 pairs: I,
2147 right: Option<Self>,
2148 ) -> Self {
2149 let chunk = Chunk::empty().append(pairs);
2150 let l = Tree::from_root(left);
2151 let r = Tree::from_root(right);
2152 match Tree::create(&l, chunk, &r) {
2153 Tree::Node(node) => NodeHandle(node),
2154 Tree::Empty => unreachable!("create of a non-empty chunk"),
2155 }
2156 }
2157}