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
32const 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); }
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 #[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 #[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 #[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 #[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 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 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 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 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 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 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 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 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 *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(), };
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 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 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 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}