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