immutable_chunkmap/
avl.rs

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    mem::{self, ManuallyDrop},
24    ptr,
25};
26#[cfg(feature = "pool")]
27use poolshark::{
28    local::{insert_raw, take},
29    location_id, Discriminant, IsoPoolable, Poolable,
30};
31
32// until we get 128 bit machines with exabytes of memory
33const MAX_DEPTH: usize = 64;
34
35fn pack_height_and_size(height: u8, size: usize) -> u64 {
36    assert!((size & 0x00ffffff_ffffffff) == size);
37    ((height as u64) << 56) | (size as u64)
38}
39
40#[cfg(feature = "pool")]
41#[derive(Debug)]
42#[repr(C)]
43pub(crate) struct NodeInner<K: Ord + Clone, V: Clone, const SIZE: usize> {
44    elts: MaybeUninit<Chunk<K, V, SIZE>>,
45    min_key: MaybeUninit<K>,
46    max_key: MaybeUninit<K>,
47    left: Tree<K, V, SIZE>,
48    right: Tree<K, V, SIZE>,
49    height_and_size: u64,
50}
51
52#[cfg(feature = "pool")]
53impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for NodeInner<K, V, SIZE> {
54    fn clone(&self) -> Self {
55        unsafe {
56            Self {
57                elts: MaybeUninit::new(self.elts.assume_init_ref().clone()),
58                min_key: MaybeUninit::new(self.min_key.assume_init_ref().clone()),
59                max_key: MaybeUninit::new(self.max_key.assume_init_ref().clone()),
60                left: self.left.clone(),
61                right: self.right.clone(),
62                height_and_size: self.height_and_size,
63            }
64        }
65    }
66}
67
68#[cfg(not(feature = "pool"))]
69#[derive(Clone, Debug)]
70pub(crate) struct NodeInner<K: Ord + Clone, V: Clone, const SIZE: usize> {
71    elts: Chunk<K, V, SIZE>,
72    min_key: K,
73    max_key: K,
74    left: Tree<K, V, SIZE>,
75    right: Tree<K, V, SIZE>,
76    height_and_size: u64,
77}
78
79#[cfg(feature = "pool")]
80pub(crate) struct Node<K: Ord + Clone, V: Clone, const SIZE: usize>(
81    ManuallyDrop<Arc<NodeInner<K, V, SIZE>>>,
82);
83
84#[cfg(not(feature = "pool"))]
85pub(crate) struct Node<K: Ord + Clone, V: Clone, const SIZE: usize>(
86    Arc<NodeInner<K, V, SIZE>>,
87);
88
89#[cfg(feature = "pool")]
90impl<K: Ord + Clone, V: Clone, const SIZE: usize> Poolable for Node<K, V, SIZE> {
91    fn capacity(&self) -> usize {
92        1
93    }
94
95    fn empty() -> Self {
96        let n = NodeInner {
97            elts: MaybeUninit::zeroed(),
98            min_key: MaybeUninit::zeroed(),
99            max_key: MaybeUninit::zeroed(),
100            left: Tree::Empty,
101            right: Tree::Empty,
102            height_and_size: 0,
103        };
104        Node(ManuallyDrop::new(Arc::new(n)))
105    }
106
107    fn really_dropped(&mut self) -> bool {
108        unreachable!()
109    }
110
111    fn reset(&mut self) {
112        unreachable!()
113    }
114}
115
116#[cfg(feature = "pool")]
117unsafe impl<K: Ord + Clone, V: Clone, const SIZE: usize> IsoPoolable
118    for Node<K, V, SIZE>
119{
120    const DISCRIMINANT: Option<Discriminant> =
121        Discriminant::new_p2_size::<K, V, SIZE>(location_id!());
122}
123
124#[cfg(feature = "pool")]
125impl<K: Ord + Clone, V: Clone, const SIZE: usize> Drop for Node<K, V, SIZE> {
126    fn drop(&mut self) {
127        match Arc::get_mut(&mut self.0) {
128            None => unsafe { ManuallyDrop::drop(&mut self.0) },
129            Some(inner) => {
130                unsafe { inner.reset() }
131                if let Some(mut n) = unsafe { insert_raw(ptr::read(self)) } {
132                    unsafe { ManuallyDrop::drop(&mut n.0) };
133                    mem::forget(n); // don't call ourselves recursively
134                }
135            }
136        }
137    }
138}
139
140#[cfg(feature = "pool")]
141impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for Node<K, V, SIZE> {
142    fn clone(&self) -> Self {
143        Self(ManuallyDrop::new(Arc::clone(&*self.0)))
144    }
145}
146
147#[cfg(not(feature = "pool"))]
148impl<K: Ord + Clone, V: Clone, const SIZE: usize> Clone for Node<K, V, SIZE> {
149    fn clone(&self) -> Self {
150        Self(Arc::clone(&self.0))
151    }
152}
153
154impl<K: Ord + Clone, V: Clone, const SIZE: usize> Deref for Node<K, V, SIZE> {
155    type Target = NodeInner<K, V, SIZE>;
156
157    fn deref(&self) -> &Self::Target {
158        &*self.0
159    }
160}
161
162impl<K: Ord + Clone, V: Clone, const SIZE: usize> Node<K, V, SIZE> {
163    fn downgrade(&self) -> WeakNode<K, V, SIZE> {
164        WeakNode(Arc::downgrade(&self.0))
165    }
166
167    #[cfg(feature = "pool")]
168    fn make_mut<'a>(&'a mut self) -> &'a mut NodeInner<K, V, SIZE> {
169        match Arc::get_mut(&mut *self.0).map(|n| n as *mut _) {
170            Some(t) => unsafe { &mut *t },
171            None => {
172                let mut n = take::<Node<K, V, SIZE>>();
173                *Arc::get_mut(&mut *n.0).unwrap() = (**self.0).clone();
174                *self = n;
175                Arc::get_mut(&mut *self.0).unwrap()
176            }
177        }
178    }
179
180    #[cfg(not(feature = "pool"))]
181    fn make_mut(&mut self) -> &mut NodeInner<K, V, SIZE> {
182        Arc::make_mut(&mut self.0)
183    }
184}
185
186#[derive(Clone)]
187pub(crate) struct WeakNode<K: Ord + Clone, V: Clone, const SIZE: usize>(
188    Weak<NodeInner<K, V, SIZE>>,
189);
190
191impl<K: Ord + Clone, V: Clone, const SIZE: usize> WeakNode<K, V, SIZE> {
192    fn upgrade(&self) -> Option<Node<K, V, SIZE>> {
193        #[cfg(feature = "pool")]
194        {
195            Weak::upgrade(&self.0).map(|n| Node(ManuallyDrop::new(n)))
196        }
197        #[cfg(not(feature = "pool"))]
198        {
199            Weak::upgrade(&self.0).map(Node)
200        }
201    }
202}
203
204impl<K, V, const SIZE: usize> NodeInner<K, V, SIZE>
205where
206    K: Ord + Clone,
207    V: Clone,
208{
209    #[cfg(feature = "pool")]
210    unsafe fn reset(&mut self) {
211        let Self {
212            elts,
213            min_key,
214            max_key,
215            left,
216            right,
217            height_and_size,
218        } = self;
219        if *height_and_size > 0 {
220            unsafe {
221                elts.assume_init_drop();
222                min_key.assume_init_drop();
223                max_key.assume_init_drop();
224            }
225            *left = Tree::Empty;
226            *right = Tree::Empty;
227            *height_and_size = 0
228        }
229    }
230
231    // a node that is not in the pool will never have elts set to None
232    #[cfg(feature = "pool")]
233    fn elts(&self) -> &Chunk<K, V, SIZE> {
234        unsafe { self.elts.assume_init_ref() }
235    }
236
237    #[cfg(not(feature = "pool"))]
238    fn elts(&self) -> &Chunk<K, V, SIZE> {
239        &self.elts
240    }
241
242    // a node that is not in the pool will never have elts set to None
243    #[cfg(feature = "pool")]
244    fn elts_mut(&mut self) -> &mut Chunk<K, V, SIZE> {
245        unsafe { self.elts.assume_init_mut() }
246    }
247
248    #[cfg(not(feature = "pool"))]
249    fn elts_mut(&mut self) -> &mut Chunk<K, V, SIZE> {
250        &mut self.elts
251    }
252
253    // a node that is not in the pool will never have min_key set to None
254    #[cfg(feature = "pool")]
255    fn min_key(&self) -> &K {
256        unsafe { self.min_key.assume_init_ref() }
257    }
258
259    #[cfg(not(feature = "pool"))]
260    fn min_key(&self) -> &K {
261        &self.min_key
262    }
263
264    // a node that is not in the pool will never have max_key set to None
265    #[cfg(feature = "pool")]
266    fn max_key(&self) -> &K {
267        unsafe { self.max_key.assume_init_ref() }
268    }
269
270    #[cfg(not(feature = "pool"))]
271    fn max_key(&self) -> &K {
272        &self.max_key
273    }
274
275    fn height(&self) -> u8 {
276        (self.height_and_size >> 56) as u8
277    }
278
279    #[cfg(feature = "pool")]
280    fn mutated(&mut self) {
281        unsafe {
282            if let Some((min, max)) = self.elts().min_max_key() {
283                *self.min_key.assume_init_mut() = min;
284                *self.max_key.assume_init_mut() = max;
285            }
286            self.height_and_size = pack_height_and_size(
287                1 + max(self.left.height(), self.right.height()),
288                self.left.len() + self.right.len(),
289            );
290        }
291    }
292
293    #[cfg(not(feature = "pool"))]
294    fn mutated(&mut self) {
295        if let Some((min, max)) = self.elts().min_max_key() {
296            self.min_key = min;
297            self.max_key = max;
298        }
299        self.height_and_size = pack_height_and_size(
300            1 + max(self.left.height(), self.right.height()),
301            self.left.len() + self.right.len(),
302        );
303    }
304}
305
306#[derive(Clone)]
307pub(crate) enum WeakTree<K: Ord + Clone, V: Clone, const SIZE: usize> {
308    Empty,
309    Node(WeakNode<K, V, SIZE>),
310}
311
312impl<K: Ord + Clone, V: Clone, const SIZE: usize> WeakTree<K, V, SIZE> {
313    pub(crate) fn upgrade(&self) -> Option<Tree<K, V, SIZE>> {
314        match self {
315            WeakTree::Empty => Some(Tree::Empty),
316            WeakTree::Node(n) => n.upgrade().map(Tree::Node),
317        }
318    }
319}
320
321#[derive(Clone)]
322pub(crate) enum Tree<K: Ord + Clone, V: Clone, const SIZE: usize> {
323    Empty,
324    Node(Node<K, V, SIZE>),
325}
326
327impl<K, V, const SIZE: usize> Hash for Tree<K, V, SIZE>
328where
329    K: Hash + Ord + Clone,
330    V: Hash + Clone,
331{
332    fn hash<H: Hasher>(&self, state: &mut H) {
333        for elt in self {
334            elt.hash(state)
335        }
336    }
337}
338
339impl<K, V, const SIZE: usize> Default for Tree<K, V, SIZE>
340where
341    K: Ord + Clone,
342    V: Clone,
343{
344    fn default() -> Tree<K, V, SIZE> {
345        Tree::Empty
346    }
347}
348
349impl<K, V, const SIZE: usize> PartialEq for Tree<K, V, SIZE>
350where
351    K: PartialEq + Ord + Clone,
352    V: PartialEq + Clone,
353{
354    fn eq(&self, other: &Tree<K, V, SIZE>) -> bool {
355        self.len() == other.len() && self.into_iter().zip(other).all(|(e0, e1)| e0 == e1)
356    }
357}
358
359impl<K, V, const SIZE: usize> Eq for Tree<K, V, SIZE>
360where
361    K: Eq + Ord + Clone,
362    V: Eq + Clone,
363{
364}
365
366impl<K, V, const SIZE: usize> PartialOrd for Tree<K, V, SIZE>
367where
368    K: Ord + Clone,
369    V: PartialOrd + Clone,
370{
371    fn partial_cmp(&self, other: &Tree<K, V, SIZE>) -> Option<Ordering> {
372        self.into_iter().partial_cmp(other.into_iter())
373    }
374}
375
376impl<K, V, const SIZE: usize> Ord for Tree<K, V, SIZE>
377where
378    K: Ord + Clone,
379    V: Ord + Clone,
380{
381    fn cmp(&self, other: &Tree<K, V, SIZE>) -> Ordering {
382        self.into_iter().cmp(other.into_iter())
383    }
384}
385
386impl<K, V, const SIZE: usize> Debug for Tree<K, V, SIZE>
387where
388    K: Debug + Ord + Clone,
389    V: Debug + Clone,
390{
391    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
392        f.debug_map().entries(self.into_iter()).finish()
393    }
394}
395
396impl<'a, Q, K, V, const SIZE: usize> Index<&'a Q> for Tree<K, V, SIZE>
397where
398    Q: Ord,
399    K: Ord + Clone + Borrow<Q>,
400    V: Clone,
401{
402    type Output = V;
403    fn index(&self, k: &Q) -> &V {
404        self.get(k).expect("element not found for key")
405    }
406}
407
408pub struct Iter<'a, R, Q, K, V, const SIZE: usize>
409where
410    Q: Ord + ?Sized,
411    R: RangeBounds<Q> + 'a,
412    K: 'a + Borrow<Q> + Ord + Clone,
413    V: 'a + Clone,
414{
415    q: PhantomData<Q>,
416    stack: ArrayVec<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>,
417    elts: Option<iter::Zip<slice::Iter<'a, K>, slice::Iter<'a, V>>>,
418    current: Option<&'a K>,
419    stack_rev: ArrayVec<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>,
420    elts_rev: Option<iter::Zip<slice::Iter<'a, K>, slice::Iter<'a, V>>>,
421    current_rev: Option<&'a K>,
422    bounds: R,
423}
424
425impl<'a, R, Q, K, V, const SIZE: usize> Iter<'a, R, Q, K, V, SIZE>
426where
427    Q: Ord + ?Sized,
428    R: RangeBounds<Q> + 'a,
429    K: 'a + Borrow<Q> + Ord + Clone,
430    V: 'a + Clone,
431{
432    // is at least one element of the chunk in bounds
433    fn any_elts_above_lbound(&self, n: &'a Node<K, V, SIZE>) -> bool {
434        let l = n.elts().len();
435        match self.bounds.start_bound() {
436            Bound::Unbounded => true,
437            Bound::Included(bound) => l == 0 || n.elts().key(l - 1).borrow() >= bound,
438            Bound::Excluded(bound) => l == 0 || n.elts().key(l - 1).borrow() > bound,
439        }
440    }
441
442    fn any_elts_below_ubound(&self, n: &'a Node<K, V, SIZE>) -> bool {
443        let l = n.elts().len();
444        match self.bounds.end_bound() {
445            Bound::Unbounded => true,
446            Bound::Included(bound) => l == 0 || n.elts().key(0).borrow() <= bound,
447            Bound::Excluded(bound) => l == 0 || n.elts().key(0).borrow() < bound,
448        }
449    }
450
451    fn any_elts_in_bounds(&self, n: &'a Node<K, V, SIZE>) -> bool {
452        self.any_elts_above_lbound(n) && self.any_elts_below_ubound(n)
453    }
454
455    fn above_lbound(&self, k: &'a K) -> bool {
456        match self.bounds.start_bound() {
457            Bound::Unbounded => true,
458            Bound::Included(bound) => k.borrow() >= bound,
459            Bound::Excluded(bound) => k.borrow() > bound,
460        }
461    }
462
463    fn below_ubound(&self, k: &'a K) -> bool {
464        match self.bounds.end_bound() {
465            Bound::Unbounded => true,
466            Bound::Included(bound) => k.borrow() <= bound,
467            Bound::Excluded(bound) => k.borrow() < bound,
468        }
469    }
470}
471
472impl<'a, R, Q, K, V, const SIZE: usize> Iterator for Iter<'a, R, Q, K, V, SIZE>
473where
474    Q: Ord + ?Sized,
475    R: RangeBounds<Q> + 'a,
476    K: 'a + Borrow<Q> + Ord + Clone,
477    V: 'a + Clone,
478{
479    type Item = (&'a K, &'a V);
480    fn next(&mut self) -> Option<Self::Item> {
481        loop {
482            loop {
483                let (k, v) = match &mut self.elts {
484                    None => break,
485                    Some(s) => match s.next() {
486                        Some((k, v)) => (k, v),
487                        None => break,
488                    },
489                };
490                if let Some(back) = self.current_rev {
491                    if k >= back {
492                        return None;
493                    }
494                }
495                if !self.below_ubound(k) {
496                    return None;
497                }
498                self.current = Some(k);
499                if self.above_lbound(k) {
500                    return Some((k, v));
501                }
502            }
503            if self.stack.is_empty() {
504                return None;
505            }
506            self.elts = None;
507            let top = self.stack.len() - 1;
508            let (visited, current) = self.stack[top];
509            if visited {
510                if self.any_elts_in_bounds(current) {
511                    self.elts = Some(current.elts().into_iter());
512                }
513                self.stack.pop();
514                match current.right {
515                    Tree::Empty => (),
516                    Tree::Node(ref n) => {
517                        if self.any_elts_below_ubound(n) || !n.left.is_empty() {
518                            self.stack.push((false, n))
519                        }
520                    }
521                };
522            } else {
523                self.stack[top].0 = true;
524                match current.left {
525                    Tree::Empty => (),
526                    Tree::Node(ref n) => {
527                        if self.any_elts_above_lbound(n) || !n.right.is_empty() {
528                            self.stack.push((false, n))
529                        }
530                    }
531                }
532            }
533        }
534    }
535}
536
537impl<'a, R, Q, K, V, const SIZE: usize> DoubleEndedIterator for Iter<'a, R, Q, K, V, SIZE>
538where
539    Q: Ord + ?Sized,
540    R: RangeBounds<Q> + 'a,
541    K: 'a + Borrow<Q> + Ord + Clone,
542    V: 'a + Clone,
543{
544    fn next_back(&mut self) -> Option<Self::Item> {
545        loop {
546            loop {
547                let (k, v) = match &mut self.elts_rev {
548                    &mut None => break,
549                    &mut Some(ref mut s) => match s.next_back() {
550                        None => break,
551                        Some((k, v)) => (k, v),
552                    },
553                };
554                if let Some(front) = self.current {
555                    if k <= front {
556                        return None;
557                    }
558                }
559                if !self.above_lbound(k) {
560                    return None;
561                }
562                self.current_rev = Some(k);
563                if self.below_ubound(k) {
564                    return Some((k, v));
565                }
566            }
567            if self.stack_rev.is_empty() {
568                return None;
569            }
570            self.elts_rev = None;
571            let top = self.stack_rev.len() - 1;
572            let (visited, current) = self.stack_rev[top];
573            if visited {
574                if self.any_elts_in_bounds(current) {
575                    self.elts_rev = Some(current.elts().into_iter());
576                }
577                self.stack_rev.pop();
578                match current.left {
579                    Tree::Empty => (),
580                    Tree::Node(ref n) => {
581                        if self.any_elts_above_lbound(n) || !n.right.is_empty() {
582                            self.stack_rev.push((false, n))
583                        }
584                    }
585                };
586            } else {
587                self.stack_rev[top].0 = true;
588                match current.right {
589                    Tree::Empty => (),
590                    Tree::Node(ref n) => {
591                        if self.any_elts_below_ubound(n) || !n.left.is_empty() {
592                            self.stack_rev.push((false, n))
593                        }
594                    }
595                }
596            }
597        }
598    }
599}
600
601pub struct IterMut<'a, R, Q, K, V, const SIZE: usize>
602where
603    Q: Ord + ?Sized,
604    R: RangeBounds<Q> + 'a,
605    K: 'a + Borrow<Q> + Ord + Clone,
606    V: 'a + Clone,
607{
608    q: PhantomData<Q>,
609    stack: ArrayVec<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>,
610    elts: Option<iter::Zip<slice::Iter<'a, K>, slice::IterMut<'a, V>>>,
611    current: Option<&'a K>,
612    stack_rev: ArrayVec<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>,
613    elts_rev: Option<iter::Zip<slice::Iter<'a, K>, slice::IterMut<'a, V>>>,
614    current_rev: Option<&'a K>,
615    bounds: R,
616}
617
618impl<'a, R, Q, K, V, const SIZE: usize> IterMut<'a, R, Q, K, V, SIZE>
619where
620    Q: Ord + ?Sized,
621    R: RangeBounds<Q> + 'a,
622    K: 'a + Borrow<Q> + Ord + Clone,
623    V: 'a + Clone,
624{
625    // is at least one element of the chunk in bounds
626    fn any_elts_above_lbound(&self, n: &'a NodeInner<K, V, SIZE>) -> bool {
627        let l = n.elts().len();
628        match self.bounds.start_bound() {
629            Bound::Unbounded => true,
630            Bound::Included(bound) => l == 0 || n.elts().key(l - 1).borrow() >= bound,
631            Bound::Excluded(bound) => l == 0 || n.elts().key(l - 1).borrow() > bound,
632        }
633    }
634
635    fn any_elts_below_ubound(&self, n: &'a NodeInner<K, V, SIZE>) -> bool {
636        let l = n.elts().len();
637        match self.bounds.end_bound() {
638            Bound::Unbounded => true,
639            Bound::Included(bound) => l == 0 || n.elts().key(0).borrow() <= bound,
640            Bound::Excluded(bound) => l == 0 || n.elts().key(0).borrow() < bound,
641        }
642    }
643
644    fn any_elts_in_bounds(&self, n: &'a NodeInner<K, V, SIZE>) -> bool {
645        self.any_elts_above_lbound(n) && self.any_elts_below_ubound(n)
646    }
647
648    fn above_lbound(&self, k: &'a K) -> bool {
649        match self.bounds.start_bound() {
650            Bound::Unbounded => true,
651            Bound::Included(bound) => k.borrow() >= bound,
652            Bound::Excluded(bound) => k.borrow() > bound,
653        }
654    }
655
656    fn below_ubound(&self, k: &'a K) -> bool {
657        match self.bounds.end_bound() {
658            Bound::Unbounded => true,
659            Bound::Included(bound) => k.borrow() <= bound,
660            Bound::Excluded(bound) => k.borrow() < bound,
661        }
662    }
663}
664
665impl<'a, R, Q, K, V, const SIZE: usize> Iterator for IterMut<'a, R, Q, K, V, SIZE>
666where
667    Q: Ord + ?Sized,
668    R: RangeBounds<Q> + 'a,
669    K: 'a + Borrow<Q> + Ord + Clone,
670    V: 'a + Clone,
671{
672    type Item = (&'a K, &'a mut V);
673    fn next(&mut self) -> Option<Self::Item> {
674        loop {
675            loop {
676                let (k, v) = match &mut self.elts {
677                    &mut None => break,
678                    &mut Some(ref mut s) => match s.next() {
679                        Some((k, v)) => (k, v),
680                        None => break,
681                    },
682                };
683                if let Some(back) = self.current_rev {
684                    if k >= back {
685                        return None;
686                    }
687                }
688                if !self.below_ubound(k) {
689                    return None;
690                }
691                self.current = Some(k);
692                if self.above_lbound(k) {
693                    return Some((k, v));
694                }
695            }
696            if self.stack.is_empty() {
697                return None;
698            }
699            self.elts = None;
700            let top = self.stack.len() - 1;
701            let (visited, current) = self.stack[top];
702            if visited {
703                if self.any_elts_in_bounds(unsafe { &*current }) {
704                    self.elts =
705                        Some((unsafe { (*current).make_mut().elts_mut() }).into_iter());
706                }
707                self.stack.pop();
708                match unsafe { &mut (*current).make_mut().right } {
709                    Tree::Empty => (),
710                    Tree::Node(ref mut n) => {
711                        if self.any_elts_below_ubound(n) || !n.left.is_empty() {
712                            self.stack.push((false, n))
713                        }
714                    }
715                };
716            } else {
717                self.stack[top].0 = true;
718                match unsafe { &mut (*current).make_mut().left } {
719                    Tree::Empty => (),
720                    Tree::Node(n) => {
721                        if self.any_elts_above_lbound(n) || !n.right.is_empty() {
722                            self.stack.push((false, n))
723                        }
724                    }
725                }
726            }
727        }
728    }
729}
730
731impl<'a, R, Q, K, V, const SIZE: usize> DoubleEndedIterator
732    for IterMut<'a, R, Q, K, V, SIZE>
733where
734    Q: Ord + ?Sized,
735    R: RangeBounds<Q> + 'a,
736    K: 'a + Borrow<Q> + Ord + Clone,
737    V: 'a + Clone,
738{
739    fn next_back(&mut self) -> Option<Self::Item> {
740        loop {
741            loop {
742                let (k, v) = match &mut self.elts_rev {
743                    &mut None => break,
744                    &mut Some(ref mut s) => match s.next_back() {
745                        None => break,
746                        Some((k, v)) => (k, v),
747                    },
748                };
749                if let Some(front) = self.current {
750                    if k <= front {
751                        return None;
752                    }
753                }
754                if !self.above_lbound(k) {
755                    return None;
756                }
757                self.current_rev = Some(k);
758                if self.below_ubound(k) {
759                    return Some((k, v));
760                }
761            }
762            if self.stack_rev.is_empty() {
763                return None;
764            }
765            self.elts_rev = None;
766            let top = self.stack_rev.len() - 1;
767            let (visited, current) = self.stack_rev[top];
768            if visited {
769                if self.any_elts_in_bounds(unsafe { &*current }) {
770                    self.elts_rev =
771                        Some((unsafe { (*current).make_mut().elts_mut() }).into_iter());
772                }
773                self.stack_rev.pop();
774                match unsafe { &mut (*current).make_mut().left } {
775                    Tree::Empty => (),
776                    Tree::Node(ref mut n) => {
777                        if self.any_elts_above_lbound(n) || !n.right.is_empty() {
778                            self.stack_rev.push((false, n))
779                        }
780                    }
781                };
782            } else {
783                self.stack_rev[top].0 = true;
784                match unsafe { &mut (*current).make_mut().right } {
785                    Tree::Empty => (),
786                    Tree::Node(ref mut n) => {
787                        if self.any_elts_below_ubound(n) || !n.left.is_empty() {
788                            self.stack_rev.push((false, n))
789                        }
790                    }
791                }
792            }
793        }
794    }
795}
796
797impl<'a, K, V, const SIZE: usize> IntoIterator for &'a Tree<K, V, SIZE>
798where
799    K: 'a + Ord + Clone,
800    V: 'a + Clone,
801{
802    type Item = (&'a K, &'a V);
803    type IntoIter = Iter<'a, RangeFull, K, K, V, SIZE>;
804    fn into_iter(self) -> Self::IntoIter {
805        self.range(..)
806    }
807}
808
809impl<K, V, const SIZE: usize> Tree<K, V, SIZE>
810where
811    K: Ord + Clone,
812    V: Clone,
813{
814    pub(crate) fn new() -> Self {
815        Tree::Empty
816    }
817
818    pub(crate) fn downgrade(&self) -> WeakTree<K, V, SIZE> {
819        match self {
820            Tree::Empty => WeakTree::Empty,
821            Tree::Node(n) => WeakTree::Node(n.downgrade()),
822        }
823    }
824
825    pub(crate) fn strong_count(&self) -> usize {
826        match self {
827            Tree::Empty => 0,
828            Tree::Node(n) => Arc::strong_count(&n.0),
829        }
830    }
831
832    pub(crate) fn weak_count(&self) -> usize {
833        match self {
834            Tree::Empty => 0,
835            Tree::Node(n) => Arc::weak_count(&n.0),
836        }
837    }
838
839    pub(crate) fn range<'a, Q, R>(&'a self, r: R) -> Iter<'a, R, Q, K, V, SIZE>
840    where
841        Q: Ord + ?Sized + 'a,
842        K: Borrow<Q>,
843        R: RangeBounds<Q> + 'a,
844    {
845        match self {
846            &Tree::Empty => Iter {
847                q: PhantomData,
848                bounds: r,
849                stack: ArrayVec::<_, MAX_DEPTH>::new(),
850                elts: None,
851                current: None,
852                stack_rev: ArrayVec::<_, MAX_DEPTH>::new(),
853                elts_rev: None,
854                current_rev: None,
855            },
856            &Tree::Node(ref n) => {
857                let mut stack =
858                    ArrayVec::<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>::new();
859                let mut stack_rev =
860                    ArrayVec::<(bool, &'a Node<K, V, SIZE>), MAX_DEPTH>::new();
861                stack.push((false, n));
862                stack_rev.push((false, n));
863                Iter {
864                    q: PhantomData,
865                    bounds: r,
866                    stack,
867                    elts: None,
868                    current: None,
869                    stack_rev,
870                    elts_rev: None,
871                    current_rev: None,
872                }
873            }
874        }
875    }
876
877    pub(crate) fn range_mut_cow<'a, Q, R>(
878        &'a mut self,
879        r: R,
880    ) -> IterMut<'a, R, Q, K, V, SIZE>
881    where
882        Q: Ord + ?Sized + 'a,
883        K: Borrow<Q>,
884        R: RangeBounds<Q> + 'a,
885    {
886        match self {
887            Tree::Empty => IterMut {
888                q: PhantomData,
889                bounds: r,
890                stack: ArrayVec::<_, MAX_DEPTH>::new(),
891                elts: None,
892                current: None,
893                stack_rev: ArrayVec::<_, MAX_DEPTH>::new(),
894                elts_rev: None,
895                current_rev: None,
896            },
897            Tree::Node(ref mut n) => {
898                let mut stack =
899                    ArrayVec::<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>::new();
900                let mut stack_rev =
901                    ArrayVec::<(bool, *mut Node<K, V, SIZE>), MAX_DEPTH>::new();
902                stack.push((false, n));
903                stack_rev.push((false, n));
904                IterMut {
905                    q: PhantomData,
906                    bounds: r,
907                    stack,
908                    elts: None,
909                    current: None,
910                    stack_rev,
911                    elts_rev: None,
912                    current_rev: None,
913                }
914            }
915        }
916    }
917
918    pub(crate) fn iter_mut_cow<'a, Q>(
919        &'a mut self,
920    ) -> IterMut<'a, RangeFull, Q, K, V, SIZE>
921    where
922        Q: Ord + ?Sized + 'a,
923        K: Borrow<Q>,
924    {
925        self.range_mut_cow(..)
926    }
927
928    fn add_min_elts(&self, elts: &Chunk<K, V, SIZE>) -> Self {
929        match self {
930            Tree::Empty => Tree::create(&Tree::Empty, elts.clone(), &Tree::Empty),
931            Tree::Node(ref n) => {
932                Tree::bal(&n.left.add_min_elts(elts), n.elts().clone(), &n.right)
933            }
934        }
935    }
936
937    fn add_max_elts(&self, elts: &Chunk<K, V, SIZE>) -> Self {
938        match self {
939            Tree::Empty => Tree::create(&Tree::Empty, elts.clone(), &Tree::Empty),
940            Tree::Node(ref n) => {
941                Tree::bal(&n.left, n.elts().clone(), &n.right.add_max_elts(elts))
942            }
943        }
944    }
945
946    // This is the same as create except it makes no assumption about the tree
947    // heights or tree balance, so you can pass it anything, and it will return
948    // a balanced tree.
949    fn join(
950        l: &Tree<K, V, SIZE>,
951        elts: &Chunk<K, V, SIZE>,
952        r: &Tree<K, V, SIZE>,
953    ) -> Self {
954        match (l, r) {
955            (Tree::Empty, _) => r.add_min_elts(elts),
956            (_, Tree::Empty) => l.add_max_elts(elts),
957            (Tree::Node(ref ln), Tree::Node(ref rn)) => {
958                let (ln_height, rn_height) = (ln.height(), rn.height());
959                if ln_height > rn_height + 2 {
960                    Tree::bal(
961                        &ln.left,
962                        ln.elts().clone(),
963                        &Tree::join(&ln.right, elts, r),
964                    )
965                } else if rn_height > ln_height + 2 {
966                    Tree::bal(
967                        &Tree::join(l, elts, &rn.left),
968                        rn.elts().clone(),
969                        &rn.right,
970                    )
971                } else {
972                    Tree::create(l, elts.clone(), r)
973                }
974            }
975        }
976    }
977
978    /// split the tree according to elts, return two balanced trees
979    /// representing all the elements less than and greater than elts,
980    /// if there is a possible intersection return the intersecting
981    /// chunk. In the case of an intersection there may also be an
982    /// intersection at the left and/or right nodes.
983    fn split(&self, vmin: &K, vmax: &K) -> (Self, Option<Chunk<K, V, SIZE>>, Self) {
984        match self {
985            Tree::Empty => (Tree::Empty, None, Tree::Empty),
986            Tree::Node(ref n) => {
987                if vmax < n.min_key() {
988                    let (ll, inter, rl) = n.left.split(vmin, vmax);
989                    (ll, inter, Tree::join(&rl, n.elts(), &n.right))
990                } else if vmin > n.max_key() {
991                    let (lr, inter, rr) = n.right.split(vmin, vmax);
992                    (Tree::join(&n.left, n.elts(), &lr), inter, rr)
993                } else {
994                    (n.left.clone(), Some(n.elts().clone()), n.right.clone())
995                }
996            }
997        }
998    }
999
1000    /// merge all the values in the root node of from into to, and
1001    /// return from with it's current root remove, and to with the
1002    /// elements merged.
1003    fn merge_root_to<F>(
1004        from: &Tree<K, V, SIZE>,
1005        to: &Tree<K, V, SIZE>,
1006        f: &mut F,
1007    ) -> (Self, Self)
1008    where
1009        F: FnMut(&K, &V, &V) -> Option<V>,
1010    {
1011        match (from, to) {
1012            (Tree::Empty, to) => (Tree::Empty, to.clone()),
1013            (Tree::Node(ref n), to) => {
1014                let to =
1015                    to.update_chunk(n.elts().to_vec(), &mut |k0, v0, cur| match cur {
1016                        None => Some((k0, v0)),
1017                        Some((_, v1)) => f(&k0, &v0, v1).map(|v| (k0, v)),
1018                    });
1019                if n.height() == 1 {
1020                    (Tree::Empty, to)
1021                } else {
1022                    match n.right {
1023                        Tree::Empty => (n.left.clone(), to),
1024                        Tree::Node(_) => {
1025                            let elts = n.right.min_elts().unwrap();
1026                            let right = n.right.remove_min_elts();
1027                            (Tree::join(&n.left, elts, &right), to)
1028                        }
1029                    }
1030                }
1031            }
1032        }
1033    }
1034
1035    /// merge two trees, where f is run on the intersection. O(log(n)
1036    /// + m) where n is the size of the largest tree, and m is the number of
1037    /// intersecting chunks.
1038    pub(crate) fn union<F>(
1039        t0: &Tree<K, V, SIZE>,
1040        t1: &Tree<K, V, SIZE>,
1041        f: &mut F,
1042    ) -> Self
1043    where
1044        F: FnMut(&K, &V, &V) -> Option<V>,
1045    {
1046        match (t0, t1) {
1047            (Tree::Empty, Tree::Empty) => Tree::Empty,
1048            (Tree::Empty, t1) => t1.clone(),
1049            (t0, Tree::Empty) => t0.clone(),
1050            (Tree::Node(ref n0), Tree::Node(ref n1)) => {
1051                if n0.height() > n1.height() {
1052                    match t1.split(n0.min_key(), n0.max_key()) {
1053                        (_, Some(_), _) => {
1054                            let (t0, t1) = Tree::merge_root_to(&t0, &t1, f);
1055                            Tree::union(&t0, &t1, f)
1056                        }
1057                        (l1, None, r1) => Tree::join(
1058                            &Tree::union(&n0.left, &l1, f),
1059                            n0.elts(),
1060                            &Tree::union(&n0.right, &r1, f),
1061                        ),
1062                    }
1063                } else {
1064                    match t0.split(n1.min_key(), n1.max_key()) {
1065                        (_, Some(_), _) => {
1066                            let (t1, t0) = Tree::merge_root_to(&t1, &t0, f);
1067                            Tree::union(&t0, &t1, f)
1068                        }
1069                        (l0, None, r0) => Tree::join(
1070                            &Tree::union(&l0, &n1.left, f),
1071                            n1.elts(),
1072                            &Tree::union(&r0, &n1.right, f),
1073                        ),
1074                    }
1075                }
1076            }
1077        }
1078    }
1079
1080    fn intersect_int<F>(
1081        t0: &Tree<K, V, SIZE>,
1082        t1: &Tree<K, V, SIZE>,
1083        r: &mut Vec<(K, V)>,
1084        f: &mut F,
1085    ) where
1086        F: FnMut(&K, &V, &V) -> Option<V>,
1087    {
1088        match (t0, t1) {
1089            (Tree::Empty, _) => (),
1090            (_, Tree::Empty) => (),
1091            (Tree::Node(ref n0), t1) => match t1.split(n0.min_key(), n0.max_key()) {
1092                (l1, None, r1) => {
1093                    Tree::intersect_int(&n0.left, &l1, r, f);
1094                    Tree::intersect_int(&n0.right, &r1, r, f);
1095                }
1096                (l1, Some(elts), r1) if elts.len() == 0 => {
1097                    Tree::intersect_int(&n0.left, &l1, r, f);
1098                    Tree::intersect_int(&n0.right, &r1, r, f);
1099                }
1100                (l1, Some(elts), r1) => {
1101                    let (min_k, max_k) = elts.min_max_key().unwrap();
1102                    Chunk::intersect(n0.elts(), &elts, r, f);
1103                    if n0.min_key() < &min_k && n0.max_key() > &max_k {
1104                        Tree::intersect_int(t0, &Tree::concat(&l1, &r1), r, f)
1105                    } else if n0.min_key() >= &min_k && n0.max_key() <= &max_k {
1106                        let t0 = Tree::concat(&n0.left, &n0.right);
1107                        let t1 = Tree::join(&l1, &elts, &r1);
1108                        Tree::intersect_int(&t0, &t1, r, f);
1109                    } else if n0.min_key() < &min_k {
1110                        let tl = Tree::join(&n0.left, n0.elts(), &Tree::Empty);
1111                        Tree::intersect_int(&tl, &l1, r, f);
1112                        let tr = Tree::join(&Tree::Empty, &elts, &r1);
1113                        Tree::intersect_int(&n0.right, &tr, r, f);
1114                    } else {
1115                        let tl = Tree::join(&l1, &elts, &Tree::Empty);
1116                        Tree::intersect_int(&tl, &n0.left, r, f);
1117                        let tr = Tree::join(&Tree::Empty, n0.elts(), &n0.right);
1118                        Tree::intersect_int(&r1, &tr, r, f);
1119                    }
1120                }
1121            },
1122        }
1123    }
1124
1125    pub(crate) fn intersect<F>(
1126        t0: &Tree<K, V, SIZE>,
1127        t1: &Tree<K, V, SIZE>,
1128        f: &mut F,
1129    ) -> Self
1130    where
1131        F: FnMut(&K, &V, &V) -> Option<V>,
1132    {
1133        let mut r = Vec::new();
1134        Tree::intersect_int(t0, t1, &mut r, f);
1135        Tree::Empty.insert_many(r.into_iter())
1136    }
1137
1138    pub(crate) fn diff<F>(t0: &Tree<K, V, SIZE>, t1: &Tree<K, V, SIZE>, f: &mut F) -> Self
1139    where
1140        F: FnMut(&K, &V, &V) -> Option<V>,
1141    {
1142        let mut actions = Vec::new();
1143        Tree::intersect_int(t0, t1, &mut Vec::new(), &mut |k, v0, v1| {
1144            actions.push((k.clone(), f(k, v0, v1)));
1145            None
1146        });
1147        t0.update_many(actions, &mut |k, v, _| v.map(|v| (k, v)))
1148    }
1149
1150    fn is_empty(&self) -> bool {
1151        match self {
1152            Tree::Empty => true,
1153            Tree::Node(..) => false,
1154        }
1155    }
1156
1157    pub(crate) fn len(&self) -> usize {
1158        match self {
1159            Tree::Empty => 0,
1160            Tree::Node(n) => {
1161                // on a 64 bit platform usize == u64, and on a 32 bit
1162                // platform there can't be enough elements to overflow
1163                // a u32
1164                let size_of_children = (n.height_and_size & 0x00ffffff_ffffffff) as usize;
1165                n.elts().len() + size_of_children
1166            }
1167        }
1168    }
1169
1170    fn height(&self) -> u8 {
1171        match self {
1172            Tree::Empty => 0,
1173            Tree::Node(ref n) => n.height(),
1174        }
1175    }
1176
1177    #[cfg(feature = "pool")]
1178    fn create(
1179        l: &Tree<K, V, SIZE>,
1180        elts: Chunk<K, V, SIZE>,
1181        r: &Tree<K, V, SIZE>,
1182    ) -> Self {
1183        let (min_key, max_key) = elts.min_max_key().unwrap();
1184        let height_and_size =
1185            pack_height_and_size(1 + max(l.height(), r.height()), l.len() + r.len());
1186        let mut t = take::<Node<K, V, SIZE>>();
1187        let inner = Arc::get_mut(&mut t.0).unwrap();
1188        inner.elts = MaybeUninit::new(elts);
1189        inner.min_key = MaybeUninit::new(min_key);
1190        inner.max_key = MaybeUninit::new(max_key);
1191        inner.left = l.clone();
1192        inner.right = r.clone();
1193        inner.height_and_size = height_and_size;
1194        Tree::Node(t)
1195    }
1196
1197    #[cfg(not(feature = "pool"))]
1198    fn create(
1199        l: &Tree<K, V, SIZE>,
1200        elts: Chunk<K, V, SIZE>,
1201        r: &Tree<K, V, SIZE>,
1202    ) -> Self {
1203        let (min_key, max_key) = elts.min_max_key().unwrap();
1204        let height_and_size =
1205            pack_height_and_size(1 + max(l.height(), r.height()), l.len() + r.len());
1206        let n = NodeInner {
1207            elts,
1208            min_key,
1209            max_key,
1210            left: l.clone(),
1211            right: r.clone(),
1212            height_and_size,
1213        };
1214        Tree::Node(Node(Arc::new(n)))
1215    }
1216
1217    fn in_bal(l: &Tree<K, V, SIZE>, r: &Tree<K, V, SIZE>) -> bool {
1218        let (hl, hr) = (l.height(), r.height());
1219        (hl <= hr.saturating_add(2)) && (hr <= hl.saturating_add(2))
1220    }
1221
1222    fn compact(self) -> Self {
1223        match self {
1224            Tree::Empty => self,
1225            Tree::Node(ref tn) => {
1226                let len = tn.elts().len();
1227                if len > SIZE >> 1 {
1228                    self
1229                } else {
1230                    match tn.right.min_elts() {
1231                        None => self,
1232                        Some(chunk) => {
1233                            let n = SIZE - len;
1234                            let to_add =
1235                                chunk.into_iter().map(|(k, v)| (k.clone(), v.clone()));
1236                            let overflow = chunk
1237                                .into_iter()
1238                                .skip(n)
1239                                .map(|(k, v)| (k.clone(), v.clone()));
1240                            let elts = tn.elts().append(to_add);
1241                            let t =
1242                                Tree::bal(&tn.left, elts, &tn.right.remove_min_elts());
1243                            if n >= chunk.len() {
1244                                t
1245                            } else {
1246                                t.insert_many(overflow)
1247                            }
1248                        }
1249                    }
1250                }
1251            }
1252        }
1253    }
1254
1255    fn bal(l: &Tree<K, V, SIZE>, elts: Chunk<K, V, SIZE>, r: &Tree<K, V, SIZE>) -> Self {
1256        let (hl, hr) = (l.height(), r.height());
1257        if hl > hr.saturating_add(2) {
1258            match *l {
1259                Tree::Empty => panic!("tree heights wrong"),
1260                Tree::Node(ref ln) => {
1261                    if ln.left.height() >= ln.right.height() {
1262                        Tree::create(
1263                            &ln.left,
1264                            ln.elts().clone(),
1265                            &Tree::create(&ln.right, elts, r),
1266                        )
1267                        .compact()
1268                    } else {
1269                        match ln.right {
1270                            Tree::Empty => panic!("tree heights wrong"),
1271                            Tree::Node(ref lrn) => Tree::create(
1272                                &Tree::create(&ln.left, ln.elts().clone(), &lrn.left),
1273                                lrn.elts().clone(),
1274                                &Tree::create(&lrn.right, elts, r),
1275                            )
1276                            .compact(),
1277                        }
1278                    }
1279                }
1280            }
1281        } else if hr > hl.saturating_add(2) {
1282            match *r {
1283                Tree::Empty => panic!("tree heights are wrong"),
1284                Tree::Node(ref rn) => {
1285                    if rn.right.height() >= rn.left.height() {
1286                        Tree::create(
1287                            &Tree::create(l, elts, &rn.left),
1288                            rn.elts().clone(),
1289                            &rn.right,
1290                        )
1291                        .compact()
1292                    } else {
1293                        match rn.left {
1294                            Tree::Empty => panic!("tree heights are wrong"),
1295                            Tree::Node(ref rln) => Tree::create(
1296                                &Tree::create(l, elts, &rln.left),
1297                                rln.elts().clone(),
1298                                &Tree::create(&rln.right, rn.elts().clone(), &rn.right),
1299                            )
1300                            .compact(),
1301                        }
1302                    }
1303                }
1304            }
1305        } else {
1306            Tree::create(l, elts, r).compact()
1307        }
1308    }
1309
1310    fn update_chunk<Q, D, F>(&self, chunk: Vec<(Q, D)>, f: &mut F) -> Self
1311    where
1312        Q: Ord,
1313        K: Borrow<Q>,
1314        F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1315    {
1316        if chunk.len() == 0 {
1317            return self.clone();
1318        }
1319        match self {
1320            &Tree::Empty => {
1321                let chunk = Chunk::create_with(chunk, f);
1322                if chunk.len() == 0 {
1323                    Tree::Empty
1324                } else {
1325                    Tree::create(&Tree::Empty, chunk, &Tree::Empty)
1326                }
1327            }
1328            &Tree::Node(ref tn) => {
1329                let leaf = match (&tn.left, &tn.right) {
1330                    (&Tree::Empty, &Tree::Empty) => true,
1331                    (_, _) => false,
1332                };
1333                match tn.elts().update_chunk(chunk, leaf, f) {
1334                    UpdateChunk::Updated {
1335                        elts,
1336                        update_left,
1337                        update_right,
1338                        overflow_right,
1339                    } => {
1340                        let l = tn.left.update_chunk(update_left, f);
1341                        let r = tn.right.insert_chunk(overflow_right);
1342                        let r = r.update_chunk(update_right, f);
1343                        Tree::bal(&l, elts, &r)
1344                    }
1345                    UpdateChunk::Removed {
1346                        not_done,
1347                        update_left,
1348                        update_right,
1349                    } => {
1350                        let l = tn.left.update_chunk(update_left, f);
1351                        let r = tn.right.update_chunk(update_right, f);
1352                        let t = Tree::concat(&l, &r);
1353                        t.update_chunk(not_done, f)
1354                    }
1355                    UpdateChunk::UpdateLeft(chunk) => {
1356                        let l = tn.left.update_chunk(chunk, f);
1357                        Tree::bal(&l, tn.elts().clone(), &tn.right)
1358                    }
1359                    UpdateChunk::UpdateRight(chunk) => {
1360                        let r = tn.right.update_chunk(chunk, f);
1361                        Tree::bal(&tn.left, tn.elts().clone(), &r)
1362                    }
1363                }
1364            }
1365        }
1366    }
1367
1368    fn insert_chunk(&self, chunk: Vec<(K, V)>) -> Self {
1369        self.update_chunk(chunk, &mut |k, v, _| Some((k, v)))
1370    }
1371
1372    pub(crate) fn update_many<Q, D, E, F>(&self, elts: E, f: &mut F) -> Self
1373    where
1374        E: IntoIterator<Item = (Q, D)>,
1375        Q: Ord,
1376        K: Borrow<Q>,
1377        F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1378    {
1379        let mut elts = {
1380            let mut v = elts.into_iter().collect::<Vec<(Q, D)>>();
1381            v.sort_by(|(ref k0, _), (ref k1, _)| k0.cmp(k1));
1382            v.dedup_by(|t0, t1| t0.0 == t1.0);
1383            v
1384        };
1385        let mut t = self.clone();
1386        while elts.len() > 0 {
1387            let chunk = elts.drain(0..min(SIZE, elts.len())).collect::<Vec<_>>();
1388            t = t.update_chunk(chunk, f)
1389        }
1390        t
1391    }
1392
1393    pub(crate) fn insert_many<E: IntoIterator<Item = (K, V)>>(&self, elts: E) -> Self {
1394        self.update_many(elts, &mut |k, v, _| Some((k, v)))
1395    }
1396
1397    pub(crate) fn update_cow<Q, D, F>(&mut self, q: Q, d: D, f: &mut F) -> Option<V>
1398    where
1399        Q: Ord,
1400        K: Borrow<Q>,
1401        F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1402    {
1403        match self {
1404            Tree::Empty => match f(q, d, None) {
1405                None => None,
1406                Some((k, v)) => {
1407                    *self =
1408                        Tree::create(&Tree::Empty, Chunk::singleton(k, v), &Tree::Empty);
1409                    None
1410                }
1411            },
1412            Tree::Node(ref mut tn) => {
1413                // CR estokes: problem? doesn't use the pool. check chunk as well.
1414                let tn = tn.make_mut();
1415                let leaf = match (&tn.left, &tn.right) {
1416                    (&Tree::Empty, &Tree::Empty) => true,
1417                    (_, _) => false,
1418                };
1419                match tn.elts_mut().update_mut(q, d, leaf, f) {
1420                    MutUpdate::UpdateLeft(k, d) => {
1421                        let prev = tn.left.update_cow(k, d, f);
1422                        if !Tree::in_bal(&tn.left, &tn.right) {
1423                            *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right)
1424                        } else {
1425                            tn.mutated();
1426                        }
1427                        prev
1428                    }
1429                    MutUpdate::UpdateRight(k, d) => {
1430                        let prev = tn.right.update_cow(k, d, f);
1431                        if !Tree::in_bal(&tn.left, &tn.right) {
1432                            *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right)
1433                        } else {
1434                            tn.mutated();
1435                        }
1436                        prev
1437                    }
1438                    MutUpdate::Updated { overflow, previous } => match overflow {
1439                        None => {
1440                            if tn.elts().len() > 0 {
1441                                tn.mutated();
1442                                previous
1443                            } else {
1444                                *self = Tree::concat(&tn.left, &tn.right);
1445                                previous
1446                            }
1447                        }
1448                        Some((ovk, ovv)) => {
1449                            let _ = tn.right.insert_cow(ovk, ovv);
1450                            if tn.elts().len() > 0 {
1451                                if !Tree::in_bal(&tn.left, &tn.right) {
1452                                    *self =
1453                                        Tree::bal(&tn.left, tn.elts().clone(), &tn.right);
1454                                    previous
1455                                } else {
1456                                    tn.mutated();
1457                                    previous
1458                                }
1459                            } else {
1460                                // this should be impossible
1461                                *self = Tree::concat(&tn.left, &tn.right);
1462                                previous
1463                            }
1464                        }
1465                    },
1466                }
1467            }
1468        }
1469    }
1470
1471    pub(crate) fn update<Q, D, F>(&self, q: Q, d: D, f: &mut F) -> (Self, Option<V>)
1472    where
1473        Q: Ord,
1474        K: Borrow<Q>,
1475        F: FnMut(Q, D, Option<(&K, &V)>) -> Option<(K, V)>,
1476    {
1477        match self {
1478            Tree::Empty => match f(q, d, None) {
1479                None => (self.clone(), None),
1480                Some((k, v)) => (
1481                    Tree::create(&Tree::Empty, Chunk::singleton(k, v), &Tree::Empty),
1482                    None,
1483                ),
1484            },
1485            Tree::Node(ref tn) => {
1486                let leaf = match (&tn.left, &tn.right) {
1487                    (&Tree::Empty, &Tree::Empty) => true,
1488                    (_, _) => false,
1489                };
1490                match tn.elts().update(q, d, leaf, f) {
1491                    Update::UpdateLeft(k, d) => {
1492                        let (l, prev) = tn.left.update(k, d, f);
1493                        (Tree::bal(&l, tn.elts().clone(), &tn.right), prev)
1494                    }
1495                    Update::UpdateRight(k, d) => {
1496                        let (r, prev) = tn.right.update(k, d, f);
1497                        (Tree::bal(&tn.left, tn.elts().clone(), &r), prev)
1498                    }
1499                    Update::Updated {
1500                        elts,
1501                        overflow,
1502                        previous,
1503                    } => match overflow {
1504                        None => {
1505                            if elts.len() == 0 {
1506                                (Tree::concat(&tn.left, &tn.right), previous)
1507                            } else {
1508                                (Tree::create(&tn.left, elts, &tn.right), previous)
1509                            }
1510                        }
1511                        Some((ovk, ovv)) => {
1512                            let (r, _) = tn.right.insert(ovk, ovv);
1513                            if elts.len() == 0 {
1514                                (Tree::concat(&tn.left, &r), previous)
1515                            } else {
1516                                (Tree::bal(&tn.left, elts, &r), previous)
1517                            }
1518                        }
1519                    },
1520                }
1521            }
1522        }
1523    }
1524
1525    pub(crate) fn insert(&self, k: K, v: V) -> (Self, Option<V>) {
1526        self.update(k, v, &mut |k, v, _| Some((k, v)))
1527    }
1528
1529    pub(crate) fn insert_cow(&mut self, k: K, v: V) -> Option<V> {
1530        self.update_cow(k, v, &mut |k, v, _| Some((k, v)))
1531    }
1532
1533    fn min_elts<'a>(&'a self) -> Option<&'a Chunk<K, V, SIZE>> {
1534        match self {
1535            Tree::Empty => None,
1536            Tree::Node(ref tn) => match tn.left {
1537                Tree::Empty => Some(tn.elts()),
1538                Tree::Node(_) => tn.left.min_elts(),
1539            },
1540        }
1541    }
1542
1543    fn remove_min_elts(&self) -> Self {
1544        match self {
1545            Tree::Empty => panic!("remove min elt"),
1546            Tree::Node(ref tn) => match tn.left {
1547                Tree::Empty => tn.right.clone(),
1548                Tree::Node(_) => {
1549                    Tree::bal(&tn.left.remove_min_elts(), tn.elts().clone(), &tn.right)
1550                }
1551            },
1552        }
1553    }
1554
1555    fn concat(l: &Tree<K, V, SIZE>, r: &Tree<K, V, SIZE>) -> Tree<K, V, SIZE> {
1556        match (l, r) {
1557            (Tree::Empty, _) => r.clone(),
1558            (_, Tree::Empty) => l.clone(),
1559            (_, _) => {
1560                let elts = match r.min_elts() {
1561                    Some(e) => e,
1562                    None => &Chunk::empty(), // this shouldn't happen
1563                };
1564                Tree::bal(l, elts.clone(), &r.remove_min_elts())
1565            }
1566        }
1567    }
1568
1569    pub(crate) fn remove<Q: ?Sized + Ord>(&self, k: &Q) -> (Self, Option<V>)
1570    where
1571        K: Borrow<Q>,
1572    {
1573        match self {
1574            &Tree::Empty => (Tree::Empty, None),
1575            &Tree::Node(ref tn) => match tn.elts().get(k) {
1576                Loc::NotPresent(_) => (self.clone(), None),
1577                Loc::Here(i) => {
1578                    let p = tn.elts().val(i).clone();
1579                    let elts = tn.elts().remove_elt_at(i);
1580                    if elts.len() == 0 {
1581                        (Tree::concat(&tn.left, &tn.right), Some(p))
1582                    } else {
1583                        (Tree::create(&tn.left, elts, &tn.right), Some(p))
1584                    }
1585                }
1586                Loc::InLeft => {
1587                    let (l, p) = tn.left.remove(k);
1588                    (Tree::bal(&l, tn.elts().clone(), &tn.right), p)
1589                }
1590                Loc::InRight => {
1591                    let (r, p) = tn.right.remove(k);
1592                    (Tree::bal(&tn.left, tn.elts().clone(), &r), p)
1593                }
1594            },
1595        }
1596    }
1597
1598    pub(crate) fn remove_cow<Q: ?Sized + Ord>(&mut self, k: &Q) -> Option<V>
1599    where
1600        K: Borrow<Q>,
1601    {
1602        match self {
1603            Tree::Empty => None,
1604            Tree::Node(ref mut tn) => {
1605                // CR estokes: validate this
1606                let tn = tn.make_mut();
1607                match tn.elts().get(k) {
1608                    Loc::NotPresent(_) => None,
1609                    Loc::Here(i) => {
1610                        let (_, p) = tn.elts_mut().remove_elt_at_mut(i);
1611                        if tn.elts().len() == 0 {
1612                            *self = Tree::concat(&tn.left, &tn.right);
1613                            Some(p)
1614                        } else {
1615                            tn.mutated();
1616                            Some(p)
1617                        }
1618                    }
1619                    Loc::InLeft => {
1620                        let p = tn.left.remove_cow(k);
1621                        if !Tree::in_bal(&tn.left, &tn.right) {
1622                            *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right);
1623                        } else {
1624                            tn.mutated()
1625                        }
1626                        p
1627                    }
1628                    Loc::InRight => {
1629                        let p = tn.right.remove_cow(k);
1630                        if !Tree::in_bal(&tn.left, &tn.right) {
1631                            *self = Tree::bal(&tn.left, tn.elts().clone(), &tn.right);
1632                        } else {
1633                            tn.mutated()
1634                        }
1635                        p
1636                    }
1637                }
1638            }
1639        }
1640    }
1641
1642    // this is structured as a loop so that the optimizer can inline
1643    // the closure argument. Sadly it doesn't do that if get_gen is a
1644    // recursive function, and the difference is >10%. True as of
1645    // 2018-07-19
1646    fn get_gen<'a, Q, F, R>(&'a self, k: &Q, f: F) -> Option<R>
1647    where
1648        Q: ?Sized + Ord,
1649        K: Borrow<Q>,
1650        F: FnOnce(&'a Chunk<K, V, SIZE>, usize) -> R,
1651        R: 'a,
1652    {
1653        match self {
1654            Tree::Empty => None,
1655            Tree::Node(n) => {
1656                let mut tn = n;
1657                loop {
1658                    match (k.cmp(tn.min_key().borrow()), k.cmp(tn.max_key().borrow())) {
1659                        (Ordering::Less, _) => match tn.left {
1660                            Tree::Empty => break None,
1661                            Tree::Node(ref n) => tn = n,
1662                        },
1663                        (_, Ordering::Greater) => match tn.right {
1664                            Tree::Empty => break None,
1665                            Tree::Node(ref n) => tn = n,
1666                        },
1667                        (_, _) => {
1668                            let e = tn.elts();
1669                            break e.get_local(k).map(|i| f(e, i));
1670                        }
1671                    }
1672                }
1673            }
1674        }
1675    }
1676
1677    pub(crate) fn get<'a, Q>(&'a self, k: &Q) -> Option<&'a V>
1678    where
1679        Q: ?Sized + Ord,
1680        K: Borrow<Q>,
1681    {
1682        self.get_gen(k, |e, i| e.val(i))
1683    }
1684
1685    pub(crate) fn get_key<'a, Q>(&'a self, k: &Q) -> Option<&'a K>
1686    where
1687        Q: ?Sized + Ord,
1688        K: Borrow<Q>,
1689    {
1690        self.get_gen(k, |e, i| e.key(i))
1691    }
1692
1693    pub(crate) fn get_full<'a, Q>(&'a self, k: &Q) -> Option<(&'a K, &'a V)>
1694    where
1695        Q: ?Sized + Ord,
1696        K: Borrow<Q>,
1697    {
1698        self.get_gen(k, |e, i| e.kv(i))
1699    }
1700
1701    pub(crate) fn get_mut_cow<'a, Q>(&'a mut self, k: &Q) -> Option<&'a mut V>
1702    where
1703        Q: ?Sized + Ord,
1704        K: Borrow<Q>,
1705    {
1706        match self {
1707            Tree::Empty => None,
1708            Tree::Node(tn) => {
1709                let tn = tn.make_mut();
1710                match (k.cmp(tn.min_key().borrow()), k.cmp(tn.max_key().borrow())) {
1711                    (Ordering::Less, _) => tn.left.get_mut_cow(k),
1712                    (_, Ordering::Greater) => tn.right.get_mut_cow(k),
1713                    (_, _) => match tn.elts().get_local(k) {
1714                        Some(i) => Some(tn.elts_mut().val_mut(i)),
1715                        None => None,
1716                    },
1717                }
1718            }
1719        }
1720    }
1721
1722    pub(crate) fn get_or_insert_cow<'a, F>(&'a mut self, k: K, f: F) -> &'a mut V
1723    where
1724        F: FnOnce() -> V,
1725    {
1726        match self.get_mut_cow(&k).map(|v| v as *mut V) {
1727            Some(v) => unsafe { &mut *v },
1728            None => {
1729                self.insert_cow(k.clone(), f());
1730                self.get_mut_cow(&k).unwrap()
1731            }
1732        }
1733    }
1734}
1735
1736impl<K, V, const SIZE: usize> Tree<K, V, SIZE>
1737where
1738    K: Ord + Clone + Debug,
1739    V: Clone + Debug,
1740{
1741    #[allow(dead_code)]
1742    pub(crate) fn invariant(&self) -> () {
1743        fn in_range<K, V, const SIZE: usize>(
1744            lower: Option<&K>,
1745            upper: Option<&K>,
1746            elts: &Chunk<K, V, SIZE>,
1747        ) -> bool
1748        where
1749            K: Ord + Clone + Debug,
1750            V: Clone + Debug,
1751        {
1752            (match lower {
1753                None => true,
1754                Some(lower) => elts
1755                    .into_iter()
1756                    .all(|(k, _)| lower.cmp(k) == Ordering::Less),
1757            }) && (match upper {
1758                None => true,
1759                Some(upper) => elts
1760                    .into_iter()
1761                    .all(|(k, _)| upper.cmp(k) == Ordering::Greater),
1762            })
1763        }
1764
1765        fn sorted<K, V, const SIZE: usize>(elts: &Chunk<K, V, SIZE>) -> bool
1766        where
1767            K: Ord + Clone + Debug,
1768            V: Clone + Debug,
1769        {
1770            if elts.len() == 1 {
1771                true
1772            } else {
1773                for i in 0..(elts.len() - 1) {
1774                    match elts.key(i).cmp(&elts.key(i + 1)) {
1775                        Ordering::Greater => return false,
1776                        Ordering::Less => (),
1777                        Ordering::Equal => panic!("duplicates found: {:#?}", elts),
1778                    }
1779                }
1780                true
1781            }
1782        }
1783
1784        fn check<K, V, const SIZE: usize>(
1785            t: &Tree<K, V, SIZE>,
1786            lower: Option<&K>,
1787            upper: Option<&K>,
1788            len: usize,
1789        ) -> (u8, usize)
1790        where
1791            K: Ord + Clone + Debug,
1792            V: Clone + Debug,
1793        {
1794            match *t {
1795                Tree::Empty => (0, len),
1796                Tree::Node(ref tn) => {
1797                    if !in_range(lower, upper, tn.elts()) {
1798                        panic!("tree invariant violated lower\n{:#?}\n\nupper\n{:#?}\n\nelts\n{:#?}\n\ntree\n{:#?}",
1799                               lower, upper, &tn.elts, t)
1800                    };
1801                    if !sorted(tn.elts()) {
1802                        panic!("elements isn't sorted")
1803                    };
1804                    let (thl, len) =
1805                        check(&tn.left, lower, tn.elts().min_elt().map(|(k, _)| k), len);
1806                    let (thr, len) =
1807                        check(&tn.right, tn.elts().max_elt().map(|(k, _)| k), upper, len);
1808                    let th = max(thl, thr).saturating_add(1);
1809                    let (hl, hr) = (tn.left.height(), tn.right.height());
1810                    let ub = max(hl, hr) - min(hl, hr);
1811                    if thl != hl {
1812                        panic!("left node height is wrong")
1813                    };
1814                    if thr != hr {
1815                        panic!("right node height is wrong")
1816                    };
1817                    let h = t.height();
1818                    if th != h {
1819                        panic!("node height is wrong {} vs {}", th, h)
1820                    };
1821                    if ub > 2 {
1822                        panic!("tree is unbalanced {:#?} tree: {:#?}", ub, t)
1823                    };
1824                    (th, len + tn.elts().len())
1825                }
1826            }
1827        }
1828
1829        //println!("{:#?}", self);
1830        let (_height, tlen) = check(self, None, None, 0);
1831        let len = self.len();
1832        if len != tlen {
1833            panic!("len is wrong {} vs {}", len, tlen)
1834        }
1835    }
1836}