im_rc/nodes/
btree.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5use std::borrow::Borrow;
6use std::cmp::Ordering;
7use std::mem;
8use std::ops::{Bound, RangeBounds};
9
10use sized_chunks::Chunk;
11use typenum::{Add1, Unsigned};
12
13use crate::config::OrdChunkSize as NodeSize;
14use crate::util::{Pool, PoolClone, PoolDefault, PoolRef};
15
16use self::Insert::*;
17use self::InsertAction::*;
18
19pub(crate) const NODE_SIZE: usize = NodeSize::USIZE;
20const MEDIAN: usize = (NODE_SIZE + 1) >> 1;
21
22pub trait BTreeValue {
23    type Key;
24    fn ptr_eq(&self, other: &Self) -> bool;
25    fn search_key<BK>(slice: &[Self], key: &BK) -> Result<usize, usize>
26    where
27        BK: Ord + ?Sized,
28        Self: Sized,
29        Self::Key: Borrow<BK>;
30    fn search_value(slice: &[Self], value: &Self) -> Result<usize, usize>
31    where
32        Self: Sized;
33    fn cmp_keys<BK>(&self, other: &BK) -> Ordering
34    where
35        BK: Ord + ?Sized,
36        Self::Key: Borrow<BK>;
37    fn cmp_values(&self, other: &Self) -> Ordering;
38}
39
40pub(crate) struct Node<A> {
41    keys: Chunk<A, NodeSize>,
42    children: Chunk<Option<PoolRef<Node<A>>>, Add1<NodeSize>>,
43}
44
45#[cfg(feature = "pool")]
46#[allow(unsafe_code)]
47unsafe fn cast_uninit<A>(target: &mut A) -> &mut mem::MaybeUninit<A> {
48    &mut *(target as *mut A as *mut mem::MaybeUninit<A>)
49}
50
51#[allow(unsafe_code)]
52impl<A> PoolDefault for Node<A> {
53    #[cfg(feature = "pool")]
54    unsafe fn default_uninit(target: &mut mem::MaybeUninit<Self>) {
55        let ptr: *mut Self = target.as_mut_ptr();
56        Chunk::default_uninit(cast_uninit(&mut (*ptr).keys));
57        Chunk::default_uninit(cast_uninit(&mut (*ptr).children));
58        (*ptr).children.push_back(None);
59    }
60}
61
62#[allow(unsafe_code)]
63impl<A> PoolClone for Node<A>
64where
65    A: Clone,
66{
67    #[cfg(feature = "pool")]
68    unsafe fn clone_uninit(&self, target: &mut mem::MaybeUninit<Self>) {
69        self.keys
70            .clone_uninit(cast_uninit(&mut (*target.as_mut_ptr()).keys));
71        self.children
72            .clone_uninit(cast_uninit(&mut (*target.as_mut_ptr()).children));
73    }
74}
75
76pub(crate) enum Insert<A> {
77    Added,
78    Replaced(A),
79    Split(Node<A>, A, Node<A>),
80}
81
82enum InsertAction<A> {
83    AddedAction,
84    ReplacedAction(A),
85    InsertAt,
86    InsertSplit(Node<A>, A, Node<A>),
87}
88
89pub(crate) enum Remove<A> {
90    NoChange,
91    Removed(A),
92    Update(A, Node<A>),
93}
94
95enum Boundary {
96    Lowest,
97    Highest,
98}
99
100enum RemoveAction {
101    DeleteAt(usize),
102    PullUp(Boundary, usize, usize),
103    Merge(usize),
104    StealFromLeft(usize),
105    StealFromRight(usize),
106    MergeFirst(usize),
107    ContinueDown(usize),
108}
109
110impl<A> Clone for Node<A>
111where
112    A: Clone,
113{
114    fn clone(&self) -> Self {
115        Node {
116            keys: self.keys.clone(),
117            children: self.children.clone(),
118        }
119    }
120}
121
122impl<A> Default for Node<A> {
123    fn default() -> Self {
124        Node {
125            keys: Chunk::new(),
126            children: Chunk::unit(None),
127        }
128    }
129}
130
131impl<A> Node<A> {
132    #[inline]
133    fn has_room(&self) -> bool {
134        self.keys.len() < NODE_SIZE
135    }
136
137    #[inline]
138    fn too_small(&self) -> bool {
139        self.keys.len() < MEDIAN
140    }
141
142    #[inline]
143    pub(crate) fn unit(value: A) -> Self {
144        Node {
145            keys: Chunk::unit(value),
146            children: Chunk::pair(None, None),
147        }
148    }
149
150    #[inline]
151    pub(crate) fn new_from_split(
152        pool: &Pool<Node<A>>,
153        left: Node<A>,
154        median: A,
155        right: Node<A>,
156    ) -> Self {
157        Node {
158            keys: Chunk::unit(median),
159            children: Chunk::pair(
160                Some(PoolRef::new(pool, left)),
161                Some(PoolRef::new(pool, right)),
162            ),
163        }
164    }
165
166    pub(crate) fn min(&self) -> Option<&A> {
167        match self.children.first().unwrap() {
168            None => self.keys.first(),
169            Some(ref child) => child.min(),
170        }
171    }
172
173    pub(crate) fn max(&self) -> Option<&A> {
174        match self.children.last().unwrap() {
175            None => self.keys.last(),
176            Some(ref child) => child.max(),
177        }
178    }
179}
180
181impl<A: BTreeValue> Node<A> {
182    fn child_contains<BK>(&self, index: usize, key: &BK) -> bool
183    where
184        BK: Ord + ?Sized,
185        A::Key: Borrow<BK>,
186    {
187        if let Some(Some(ref child)) = self.children.get(index) {
188            child.lookup(key).is_some()
189        } else {
190            false
191        }
192    }
193
194    pub(crate) fn lookup<BK>(&self, key: &BK) -> Option<&A>
195    where
196        BK: Ord + ?Sized,
197        A::Key: Borrow<BK>,
198    {
199        if self.keys.is_empty() {
200            return None;
201        }
202        // Perform a binary search, resulting in either a match or
203        // the index of the first higher key, meaning we search the
204        // child to the left of it.
205        match A::search_key(&self.keys, key) {
206            Ok(index) => Some(&self.keys[index]),
207            Err(index) => match self.children[index] {
208                None => None,
209                Some(ref node) => node.lookup(key),
210            },
211        }
212    }
213
214    pub(crate) fn lookup_mut<BK>(&mut self, pool: &Pool<Node<A>>, key: &BK) -> Option<&mut A>
215    where
216        A: Clone,
217        BK: Ord + ?Sized,
218        A::Key: Borrow<BK>,
219    {
220        if self.keys.is_empty() {
221            return None;
222        }
223        // Perform a binary search, resulting in either a match or
224        // the index of the first higher key, meaning we search the
225        // child to the left of it.
226        match A::search_key(&self.keys, key) {
227            Ok(index) => Some(&mut self.keys[index]),
228            Err(index) => match self.children[index] {
229                None => None,
230                Some(ref mut child_ref) => {
231                    let child = PoolRef::make_mut(pool, child_ref);
232                    child.lookup_mut(pool, key)
233                }
234            },
235        }
236    }
237
238    pub(crate) fn lookup_prev<'a, BK>(&'a self, key: &BK) -> Option<&A>
239    where
240        BK: Ord + ?Sized,
241        A::Key: Borrow<BK>,
242    {
243        if self.keys.is_empty() {
244            return None;
245        }
246        match A::search_key(&self.keys, key) {
247            Ok(index) => Some(&self.keys[index]),
248            Err(index) => match self.children[index] {
249                None if index == 0 => None,
250                None => self.keys.get(index - 1).map(|_| &self.keys[index - 1]),
251                Some(ref node) => node.lookup_prev(key),
252            },
253        }
254    }
255
256    pub(crate) fn lookup_next<'a, BK>(&'a self, key: &BK) -> Option<&A>
257    where
258        BK: Ord + ?Sized,
259        A::Key: Borrow<BK>,
260    {
261        if self.keys.is_empty() {
262            return None;
263        }
264        match A::search_key(&self.keys, key) {
265            Ok(index) => Some(&self.keys[index]),
266            Err(index) => match self.children[index] {
267                None => self.keys.get(index).map(|_| &self.keys[index]),
268                Some(ref node) => node.lookup_next(key),
269            },
270        }
271    }
272
273    pub(crate) fn lookup_prev_mut<'a, BK>(
274        &'a mut self,
275        pool: &Pool<Node<A>>,
276        key: &BK,
277    ) -> Option<&mut A>
278    where
279        A: Clone,
280        BK: Ord + ?Sized,
281        A::Key: Borrow<BK>,
282    {
283        if self.keys.is_empty() {
284            return None;
285        }
286        match A::search_key(&self.keys, key) {
287            Ok(index) => Some(&mut self.keys[index]),
288            Err(index) => match self.children[index] {
289                None if index == 0 => None,
290                None => match self.keys.get(index - 1) {
291                    Some(_) => Some(&mut self.keys[index - 1]),
292                    None => None,
293                },
294                Some(ref mut node) => PoolRef::make_mut(pool, node).lookup_prev_mut(pool, key),
295            },
296        }
297    }
298
299    pub(crate) fn lookup_next_mut<'a, BK>(
300        &'a mut self,
301        pool: &Pool<Node<A>>,
302        key: &BK,
303    ) -> Option<&mut A>
304    where
305        A: Clone,
306        BK: Ord + ?Sized,
307        A::Key: Borrow<BK>,
308    {
309        if self.keys.is_empty() {
310            return None;
311        }
312        match A::search_key(&self.keys, key) {
313            Ok(index) => Some(&mut self.keys[index]),
314            Err(index) => match self.children[index] {
315                None => match self.keys.get(index) {
316                    Some(_) => Some(&mut self.keys[index]),
317                    None => None,
318                },
319                Some(ref mut node) => PoolRef::make_mut(pool, node).lookup_next_mut(pool, key),
320            },
321        }
322    }
323
324    pub(crate) fn path_first<'a, BK>(
325        &'a self,
326        mut path: Vec<(&'a Node<A>, usize)>,
327    ) -> Vec<(&'a Node<A>, usize)>
328    where
329        A: 'a,
330        BK: Ord + ?Sized,
331        A::Key: Borrow<BK>,
332    {
333        if self.keys.is_empty() {
334            return Vec::new();
335        }
336        match self.children[0] {
337            None => {
338                path.push((self, 0));
339                path
340            }
341            Some(ref node) => {
342                path.push((self, 0));
343                node.path_first(path)
344            }
345        }
346    }
347
348    pub(crate) fn path_last<'a, BK>(
349        &'a self,
350        mut path: Vec<(&'a Node<A>, usize)>,
351    ) -> Vec<(&'a Node<A>, usize)>
352    where
353        A: 'a,
354        BK: Ord + ?Sized,
355        A::Key: Borrow<BK>,
356    {
357        if self.keys.is_empty() {
358            return Vec::new();
359        }
360        let end = self.children.len() - 1;
361        match self.children[end] {
362            None => {
363                path.push((self, end - 1));
364                path
365            }
366            Some(ref node) => {
367                path.push((self, end));
368                node.path_last(path)
369            }
370        }
371    }
372
373    pub(crate) fn path_next<'a, BK>(
374        &'a self,
375        key: &BK,
376        mut path: Vec<(&'a Node<A>, usize)>,
377    ) -> Vec<(&'a Node<A>, usize)>
378    where
379        A: 'a,
380        BK: Ord + ?Sized,
381        A::Key: Borrow<BK>,
382    {
383        if self.keys.is_empty() {
384            return Vec::new();
385        }
386        match A::search_key(&self.keys, key) {
387            Ok(index) => {
388                path.push((self, index));
389                path
390            }
391            Err(index) => match self.children[index] {
392                None => match self.keys.get(index) {
393                    Some(_) => {
394                        path.push((self, index));
395                        path
396                    }
397                    None => {
398                        // go back up to find next
399                        while let Some((node, idx)) = path.last() {
400                            if node.keys.len() == *idx {
401                                path.pop();
402                            } else {
403                                break;
404                            }
405                        }
406                        path
407                    }
408                },
409                Some(ref node) => {
410                    path.push((self, index));
411                    node.path_next(key, path)
412                }
413            },
414        }
415    }
416
417    pub(crate) fn path_prev<'a, BK>(
418        &'a self,
419        key: &BK,
420        mut path: Vec<(&'a Node<A>, usize)>,
421    ) -> Vec<(&'a Node<A>, usize)>
422    where
423        A: 'a,
424        BK: Ord + ?Sized,
425        A::Key: Borrow<BK>,
426    {
427        if self.keys.is_empty() {
428            return Vec::new();
429        }
430        match A::search_key(&self.keys, key) {
431            Ok(index) => {
432                path.push((self, index));
433                path
434            }
435            Err(index) => match self.children[index] {
436                None if index == 0 => {
437                    // go back up to find prev
438                    while let Some((_, idx)) = path.last_mut() {
439                        if *idx == 0 {
440                            path.pop();
441                        } else {
442                            *idx -= 1;
443                            break;
444                        }
445                    }
446                    path
447                }
448                None => {
449                    path.push((self, index - 1));
450                    path
451                }
452                Some(ref node) => {
453                    path.push((self, index));
454                    node.path_prev(key, path)
455                }
456            },
457        }
458    }
459
460    fn split(
461        &mut self,
462        pool: &Pool<Node<A>>,
463        value: A,
464        ins_left: Option<Node<A>>,
465        ins_right: Option<Node<A>>,
466    ) -> Insert<A> {
467        let left_child = ins_left.map(|node| PoolRef::new(pool, node));
468        let right_child = ins_right.map(|node| PoolRef::new(pool, node));
469        let index = A::search_value(&self.keys, &value).unwrap_err();
470        let mut left_keys;
471        let mut left_children;
472        let mut right_keys;
473        let mut right_children;
474        let median;
475        match index.cmp(&MEDIAN) {
476            Ordering::Less => {
477                self.children[index] = left_child;
478
479                left_keys = Chunk::from_front(&mut self.keys, index);
480                left_keys.push_back(value);
481                left_keys.drain_from_front(&mut self.keys, MEDIAN - index - 1);
482
483                left_children = Chunk::from_front(&mut self.children, index + 1);
484                left_children.push_back(right_child);
485                left_children.drain_from_front(&mut self.children, MEDIAN - index - 1);
486
487                median = self.keys.pop_front();
488
489                right_keys = Chunk::drain_from(&mut self.keys);
490                right_children = Chunk::drain_from(&mut self.children);
491            }
492            Ordering::Greater => {
493                self.children[index] = left_child;
494
495                left_keys = Chunk::from_front(&mut self.keys, MEDIAN);
496                left_children = Chunk::from_front(&mut self.children, MEDIAN + 1);
497
498                median = self.keys.pop_front();
499
500                right_keys = Chunk::from_front(&mut self.keys, index - MEDIAN - 1);
501                right_keys.push_back(value);
502                right_keys.append(&mut self.keys);
503
504                right_children = Chunk::from_front(&mut self.children, index - MEDIAN);
505                right_children.push_back(right_child);
506                right_children.append(&mut self.children);
507            }
508            Ordering::Equal => {
509                left_keys = Chunk::from_front(&mut self.keys, MEDIAN);
510                left_children = Chunk::from_front(&mut self.children, MEDIAN);
511                left_children.push_back(left_child);
512
513                median = value;
514
515                right_keys = Chunk::drain_from(&mut self.keys);
516                right_children = Chunk::drain_from(&mut self.children);
517                right_children[0] = right_child;
518            }
519        }
520
521        debug_assert!(left_keys.len() == MEDIAN);
522        debug_assert!(left_children.len() == MEDIAN + 1);
523        debug_assert!(right_keys.len() == MEDIAN);
524        debug_assert!(right_children.len() == MEDIAN + 1);
525
526        Split(
527            Node {
528                keys: left_keys,
529                children: left_children,
530            },
531            median,
532            Node {
533                keys: right_keys,
534                children: right_children,
535            },
536        )
537    }
538
539    fn merge(middle: A, left: Node<A>, mut right: Node<A>) -> Node<A> {
540        let mut keys = left.keys;
541        keys.push_back(middle);
542        keys.append(&mut right.keys);
543        let mut children = left.children;
544        children.append(&mut right.children);
545        Node { keys, children }
546    }
547
548    fn pop_min(&mut self) -> (A, Option<PoolRef<Node<A>>>) {
549        let value = self.keys.pop_front();
550        let child = self.children.pop_front();
551        (value, child)
552    }
553
554    fn pop_max(&mut self) -> (A, Option<PoolRef<Node<A>>>) {
555        let value = self.keys.pop_back();
556        let child = self.children.pop_back();
557        (value, child)
558    }
559
560    fn push_min(&mut self, child: Option<PoolRef<Node<A>>>, value: A) {
561        self.keys.push_front(value);
562        self.children.push_front(child);
563    }
564
565    fn push_max(&mut self, child: Option<PoolRef<Node<A>>>, value: A) {
566        self.keys.push_back(value);
567        self.children.push_back(child);
568    }
569
570    pub(crate) fn insert(&mut self, pool: &Pool<Node<A>>, value: A) -> Insert<A>
571    where
572        A: Clone,
573    {
574        if self.keys.is_empty() {
575            self.keys.push_back(value);
576            self.children.push_back(None);
577            return Insert::Added;
578        }
579        let (median, left, right) = match A::search_value(&self.keys, &value) {
580            // Key exists in node
581            Ok(index) => {
582                return Insert::Replaced(mem::replace(&mut self.keys[index], value));
583            }
584            // Key is adjacent to some key in node
585            Err(index) => {
586                let has_room = self.has_room();
587                let action = match self.children[index] {
588                    // No child at location, this is the target node.
589                    None => InsertAt,
590                    // Child at location, pass it on.
591                    Some(ref mut child_ref) => {
592                        let child = PoolRef::make_mut(pool, child_ref);
593                        match child.insert(pool, value.clone()) {
594                            Insert::Added => AddedAction,
595                            Insert::Replaced(value) => ReplacedAction(value),
596                            Insert::Split(left, median, right) => InsertSplit(left, median, right),
597                        }
598                    }
599                };
600                match action {
601                    ReplacedAction(value) => return Insert::Replaced(value),
602                    AddedAction => {
603                        return Insert::Added;
604                    }
605                    InsertAt => {
606                        if has_room {
607                            self.keys.insert(index, value);
608                            self.children.insert(index + 1, None);
609                            return Insert::Added;
610                        } else {
611                            (value, None, None)
612                        }
613                    }
614                    InsertSplit(left, median, right) => {
615                        if has_room {
616                            self.children[index] = Some(PoolRef::new(pool, left));
617                            self.keys.insert(index, median);
618                            self.children
619                                .insert(index + 1, Some(PoolRef::new(pool, right)));
620                            return Insert::Added;
621                        } else {
622                            (median, Some(left), Some(right))
623                        }
624                    }
625                }
626            }
627        };
628        self.split(pool, median, left, right)
629    }
630
631    pub(crate) fn remove<BK>(&mut self, pool: &Pool<Node<A>>, key: &BK) -> Remove<A>
632    where
633        A: Clone,
634        BK: Ord + ?Sized,
635        A::Key: Borrow<BK>,
636    {
637        let index = A::search_key(&self.keys, key);
638        self.remove_index(pool, index, Ok(key))
639    }
640
641    fn remove_target<BK>(
642        &mut self,
643        pool: &Pool<Node<A>>,
644        target: Result<&BK, Boundary>,
645    ) -> Remove<A>
646    where
647        A: Clone,
648        BK: Ord + ?Sized,
649        A::Key: Borrow<BK>,
650    {
651        let index = match target {
652            Ok(key) => A::search_key(&self.keys, key),
653            Err(Boundary::Lowest) => Err(0),
654            Err(Boundary::Highest) => Err(self.keys.len()),
655        };
656        self.remove_index(pool, index, target)
657    }
658
659    fn remove_index<BK>(
660        &mut self,
661        pool: &Pool<Node<A>>,
662        index: Result<usize, usize>,
663        target: Result<&BK, Boundary>,
664    ) -> Remove<A>
665    where
666        A: Clone,
667        BK: Ord + ?Sized,
668        A::Key: Borrow<BK>,
669    {
670        let action = match index {
671            // Key exists in node, remove it.
672            Ok(index) => {
673                match (&self.children[index], &self.children[index + 1]) {
674                    // If we're a leaf, just delete the entry.
675                    (&None, &None) => RemoveAction::DeleteAt(index),
676                    // First consider pulling either predecessor (from left) or successor (from right).
677                    // otherwise just merge the two small children.
678                    (&Some(ref left), &Some(ref right)) => {
679                        if !left.too_small() {
680                            RemoveAction::PullUp(Boundary::Highest, index, index)
681                        } else if !right.too_small() {
682                            RemoveAction::PullUp(Boundary::Lowest, index, index + 1)
683                        } else {
684                            RemoveAction::Merge(index)
685                        }
686                    }
687                    _ => unreachable!("Branch missing children"),
688                }
689            }
690            // Target is adjacent to some key in node
691            Err(index) => match self.children[index] {
692                // We're deading with a leaf node
693                None => match target {
694                    // No child at location means key isn't in map.
695                    Ok(_key) => return Remove::NoChange,
696                    // Looking for the lowest or highest key
697                    Err(Boundary::Lowest) => RemoveAction::DeleteAt(0),
698                    Err(Boundary::Highest) => RemoveAction::DeleteAt(self.keys.len() - 1),
699                },
700                // Child at location, but it's at minimum capacity.
701                Some(ref child) if child.too_small() => {
702                    let left = if index > 0 {
703                        self.children.get(index - 1)
704                    } else {
705                        None
706                    }; // index is usize and can't be negative, best make sure it never is.
707                    match (left, self.children.get(index + 1)) {
708                        // If it has a left sibling with capacity, steal a key from it.
709                        (Some(&Some(ref old_left)), _) if !old_left.too_small() => {
710                            RemoveAction::StealFromLeft(index)
711                        }
712                        // If it has a right sibling with capacity, same as above.
713                        (_, Some(&Some(ref old_right))) if !old_right.too_small() => {
714                            RemoveAction::StealFromRight(index)
715                        }
716                        // If it has neither, we'll have to merge it with a sibling.
717                        // If we have a right sibling, we'll merge with that.
718                        (_, Some(&Some(_))) => RemoveAction::MergeFirst(index),
719                        // If we have a left sibling, we'll merge with that.
720                        (Some(&Some(_)), _) => RemoveAction::MergeFirst(index - 1),
721                        // If none of the above, we're in a bad state.
722                        _ => unreachable!(),
723                    }
724                }
725                // Child at location, and it's big enough, we can recurse down.
726                Some(_) => RemoveAction::ContinueDown(index),
727            },
728        };
729        match action {
730            RemoveAction::DeleteAt(index) => {
731                let pair = self.keys.remove(index);
732                self.children.remove(index);
733                Remove::Removed(pair)
734            }
735            RemoveAction::PullUp(boundary, pull_to, child_index) => {
736                let children = &mut self.children;
737                let mut update = None;
738                let value;
739                if let Some(&mut Some(ref mut child_ref)) = children.get_mut(child_index) {
740                    let child = PoolRef::make_mut(pool, child_ref);
741                    match child.remove_target(pool, Err(boundary)) {
742                        Remove::NoChange => unreachable!(),
743                        Remove::Removed(pulled_value) => {
744                            value = self.keys.set(pull_to, pulled_value);
745                        }
746                        Remove::Update(pulled_value, new_child) => {
747                            value = self.keys.set(pull_to, pulled_value);
748                            update = Some(new_child);
749                        }
750                    }
751                } else {
752                    unreachable!()
753                }
754                if let Some(new_child) = update {
755                    children[child_index] = Some(PoolRef::new(pool, new_child));
756                }
757                Remove::Removed(value)
758            }
759            RemoveAction::Merge(index) => {
760                let left = self.children.remove(index).unwrap();
761                let right = mem::replace(&mut self.children[index], None).unwrap();
762                let value = self.keys.remove(index);
763                let mut merged_child = Node::merge(
764                    value,
765                    PoolRef::unwrap_or_clone(left),
766                    PoolRef::unwrap_or_clone(right),
767                );
768                let (removed, new_child) = match merged_child.remove_target(pool, target) {
769                    Remove::NoChange => unreachable!(),
770                    Remove::Removed(removed) => (removed, merged_child),
771                    Remove::Update(removed, updated_child) => (removed, updated_child),
772                };
773                if self.keys.is_empty() {
774                    // If we've depleted the root node, the merged child becomes the root.
775                    Remove::Update(removed, new_child)
776                } else {
777                    self.children[index] = Some(PoolRef::new(pool, new_child));
778                    Remove::Removed(removed)
779                }
780            }
781            RemoveAction::StealFromLeft(index) => {
782                let mut update = None;
783                let out_value;
784                {
785                    let mut children = self.children.as_mut_slice()[index - 1..=index]
786                        .iter_mut()
787                        .map(|n| n.as_mut().unwrap());
788                    let left = PoolRef::make_mut(pool, children.next().unwrap());
789                    let child = PoolRef::make_mut(pool, children.next().unwrap());
790                    // Prepare the rebalanced node.
791                    child.push_min(
792                        left.children.last().unwrap().clone(),
793                        self.keys[index - 1].clone(),
794                    );
795                    match child.remove_target(pool, target) {
796                        Remove::NoChange => {
797                            // Key wasn't there, we need to revert the steal.
798                            child.pop_min();
799                            return Remove::NoChange;
800                        }
801                        Remove::Removed(value) => {
802                            // If we did remove something, we complete the rebalancing.
803                            let (left_value, _) = left.pop_max();
804                            self.keys[index - 1] = left_value;
805                            out_value = value;
806                        }
807                        Remove::Update(value, new_child) => {
808                            // If we did remove something, we complete the rebalancing.
809                            let (left_value, _) = left.pop_max();
810                            self.keys[index - 1] = left_value;
811                            update = Some(new_child);
812                            out_value = value;
813                        }
814                    }
815                }
816                if let Some(new_child) = update {
817                    self.children[index] = Some(PoolRef::new(pool, new_child));
818                }
819                Remove::Removed(out_value)
820            }
821            RemoveAction::StealFromRight(index) => {
822                let mut update = None;
823                let out_value;
824                {
825                    let mut children = self.children.as_mut_slice()[index..index + 2]
826                        .iter_mut()
827                        .map(|n| n.as_mut().unwrap());
828                    let child = PoolRef::make_mut(pool, children.next().unwrap());
829                    let right = PoolRef::make_mut(pool, children.next().unwrap());
830                    // Prepare the rebalanced node.
831                    child.push_max(right.children[0].clone(), self.keys[index].clone());
832                    match child.remove_target(pool, target) {
833                        Remove::NoChange => {
834                            // Key wasn't there, we need to revert the steal.
835                            child.pop_max();
836                            return Remove::NoChange;
837                        }
838                        Remove::Removed(value) => {
839                            // If we did remove something, we complete the rebalancing.
840                            let (right_value, _) = right.pop_min();
841                            self.keys[index] = right_value;
842                            out_value = value;
843                        }
844                        Remove::Update(value, new_child) => {
845                            // If we did remove something, we complete the rebalancing.
846                            let (right_value, _) = right.pop_min();
847                            self.keys[index] = right_value;
848                            update = Some(new_child);
849                            out_value = value;
850                        }
851                    }
852                }
853                if let Some(new_child) = update {
854                    self.children[index] = Some(PoolRef::new(pool, new_child));
855                }
856                Remove::Removed(out_value)
857            }
858            RemoveAction::MergeFirst(index) => {
859                if let Ok(key) = target {
860                    // Bail early if we're looking for a not existing key
861                    match self.keys[index].cmp_keys(key) {
862                        Ordering::Less if !self.child_contains(index + 1, key) => {
863                            return Remove::NoChange
864                        }
865                        Ordering::Greater if !self.child_contains(index, key) => {
866                            return Remove::NoChange
867                        }
868                        _ => (),
869                    }
870                }
871                let left = self.children.remove(index).unwrap();
872                let right = mem::replace(&mut self.children[index], None).unwrap();
873                let middle = self.keys.remove(index);
874                let mut merged = Node::merge(
875                    middle,
876                    PoolRef::unwrap_or_clone(left),
877                    PoolRef::unwrap_or_clone(right),
878                );
879                let update;
880                let out_value;
881                match merged.remove_target(pool, target) {
882                    Remove::NoChange => {
883                        panic!("nodes::btree::Node::remove: caught an absent key too late while merging");
884                    }
885                    Remove::Removed(value) => {
886                        if self.keys.is_empty() {
887                            return Remove::Update(value, merged);
888                        }
889                        update = merged;
890                        out_value = value;
891                    }
892                    Remove::Update(value, new_child) => {
893                        if self.keys.is_empty() {
894                            return Remove::Update(value, new_child);
895                        }
896                        update = new_child;
897                        out_value = value;
898                    }
899                }
900                self.children[index] = Some(PoolRef::new(pool, update));
901                Remove::Removed(out_value)
902            }
903            RemoveAction::ContinueDown(index) => {
904                let mut update = None;
905                let out_value;
906                if let Some(&mut Some(ref mut child_ref)) = self.children.get_mut(index) {
907                    let child = PoolRef::make_mut(pool, child_ref);
908                    match child.remove_target(pool, target) {
909                        Remove::NoChange => return Remove::NoChange,
910                        Remove::Removed(value) => {
911                            out_value = value;
912                        }
913                        Remove::Update(value, new_child) => {
914                            update = Some(new_child);
915                            out_value = value;
916                        }
917                    }
918                } else {
919                    unreachable!()
920                }
921                if let Some(new_child) = update {
922                    self.children[index] = Some(PoolRef::new(pool, new_child));
923                }
924                Remove::Removed(out_value)
925            }
926        }
927    }
928}
929
930// Iterator
931
932/// An iterator over an ordered set.
933pub struct Iter<'a, A> {
934    fwd_path: Vec<(&'a Node<A>, usize)>,
935    back_path: Vec<(&'a Node<A>, usize)>,
936    pub(crate) remaining: usize,
937}
938
939impl<'a, A: BTreeValue> Iter<'a, A> {
940    pub(crate) fn new<R, BK>(root: &'a Node<A>, size: usize, range: R) -> Self
941    where
942        R: RangeBounds<BK>,
943        A::Key: Borrow<BK>,
944        BK: Ord + ?Sized,
945    {
946        let fwd_path = match range.start_bound() {
947            Bound::Included(key) => root.path_next(key, Vec::new()),
948            Bound::Excluded(key) => {
949                let mut path = root.path_next(key, Vec::new());
950                if let Some(value) = Self::get(&path) {
951                    if value.cmp_keys(key) == Ordering::Equal {
952                        Self::step_forward(&mut path);
953                    }
954                }
955                path
956            }
957            Bound::Unbounded => root.path_first(Vec::new()),
958        };
959        let back_path = match range.end_bound() {
960            Bound::Included(key) => root.path_prev(key, Vec::new()),
961            Bound::Excluded(key) => {
962                let mut path = root.path_prev(key, Vec::new());
963                if let Some(value) = Self::get(&path) {
964                    if value.cmp_keys(key) == Ordering::Equal {
965                        Self::step_back(&mut path);
966                    }
967                }
968                path
969            }
970            Bound::Unbounded => root.path_last(Vec::new()),
971        };
972        Iter {
973            fwd_path,
974            back_path,
975            remaining: size,
976        }
977    }
978
979    fn get(path: &[(&'a Node<A>, usize)]) -> Option<&'a A> {
980        match path.last() {
981            Some((node, index)) => Some(&node.keys[*index]),
982            None => None,
983        }
984    }
985
986    fn step_forward(path: &mut Vec<(&'a Node<A>, usize)>) -> Option<&'a A> {
987        match path.pop() {
988            Some((node, index)) => {
989                let index = index + 1;
990                match node.children[index] {
991                    // Child between current and next key -> step down
992                    Some(ref child) => {
993                        path.push((node, index));
994                        path.push((child, 0));
995                        let mut node = child;
996                        while let Some(ref left_child) = node.children[0] {
997                            path.push((left_child, 0));
998                            node = left_child;
999                        }
1000                        Some(&node.keys[0])
1001                    }
1002                    None => match node.keys.get(index) {
1003                        // Yield next key
1004                        value @ Some(_) => {
1005                            path.push((node, index));
1006                            value
1007                        }
1008                        // No more keys -> exhausted level, step up and yield
1009                        None => loop {
1010                            match path.pop() {
1011                                None => {
1012                                    return None;
1013                                }
1014                                Some((node, index)) => {
1015                                    if let value @ Some(_) = node.keys.get(index) {
1016                                        path.push((node, index));
1017                                        return value;
1018                                    }
1019                                }
1020                            }
1021                        },
1022                    },
1023                }
1024            }
1025            None => None,
1026        }
1027    }
1028
1029    fn step_back(path: &mut Vec<(&'a Node<A>, usize)>) -> Option<&'a A> {
1030        match path.pop() {
1031            Some((node, index)) => match node.children[index] {
1032                Some(ref child) => {
1033                    path.push((node, index));
1034                    let mut end = child.keys.len() - 1;
1035                    path.push((child, end));
1036                    let mut node = child;
1037                    while let Some(ref right_child) = node.children[end + 1] {
1038                        end = right_child.keys.len() - 1;
1039                        path.push((right_child, end));
1040                        node = right_child;
1041                    }
1042                    Some(&node.keys[end])
1043                }
1044                None => {
1045                    if index == 0 {
1046                        loop {
1047                            match path.pop() {
1048                                None => {
1049                                    return None;
1050                                }
1051                                Some((node, index)) => {
1052                                    if index > 0 {
1053                                        let index = index - 1;
1054                                        path.push((node, index));
1055                                        return Some(&node.keys[index]);
1056                                    }
1057                                }
1058                            }
1059                        }
1060                    } else {
1061                        let index = index - 1;
1062                        path.push((node, index));
1063                        Some(&node.keys[index])
1064                    }
1065                }
1066            },
1067            None => None,
1068        }
1069    }
1070}
1071
1072impl<'a, A: 'a + BTreeValue> Iterator for Iter<'a, A> {
1073    type Item = &'a A;
1074
1075    fn next(&mut self) -> Option<Self::Item> {
1076        match Iter::get(&self.fwd_path) {
1077            None => None,
1078            Some(value) => match Iter::get(&self.back_path) {
1079                Some(last_value) if value.cmp_values(last_value) == Ordering::Greater => None,
1080                None => None,
1081                Some(_) => {
1082                    Iter::step_forward(&mut self.fwd_path);
1083                    self.remaining -= 1;
1084                    Some(value)
1085                }
1086            },
1087        }
1088    }
1089
1090    fn size_hint(&self) -> (usize, Option<usize>) {
1091        // (0, Some(self.remaining))
1092        (0, None)
1093    }
1094}
1095
1096impl<'a, A: 'a + BTreeValue> DoubleEndedIterator for Iter<'a, A> {
1097    fn next_back(&mut self) -> Option<Self::Item> {
1098        match Iter::get(&self.back_path) {
1099            None => None,
1100            Some(value) => match Iter::get(&self.fwd_path) {
1101                Some(last_value) if value.cmp_values(last_value) == Ordering::Less => None,
1102                None => None,
1103                Some(_) => {
1104                    Iter::step_back(&mut self.back_path);
1105                    self.remaining -= 1;
1106                    Some(value)
1107                }
1108            },
1109        }
1110    }
1111}
1112
1113// Consuming iterator
1114
1115enum ConsumingIterItem<A> {
1116    Consider(Node<A>),
1117    Yield(A),
1118}
1119
1120/// A consuming iterator over an ordered set.
1121pub struct ConsumingIter<A> {
1122    fwd_last: Option<A>,
1123    fwd_stack: Vec<ConsumingIterItem<A>>,
1124    back_last: Option<A>,
1125    back_stack: Vec<ConsumingIterItem<A>>,
1126    remaining: usize,
1127}
1128
1129impl<A: Clone> ConsumingIter<A> {
1130    pub(crate) fn new(root: &Node<A>, total: usize) -> Self {
1131        ConsumingIter {
1132            fwd_last: None,
1133            fwd_stack: vec![ConsumingIterItem::Consider(root.clone())],
1134            back_last: None,
1135            back_stack: vec![ConsumingIterItem::Consider(root.clone())],
1136            remaining: total,
1137        }
1138    }
1139
1140    fn push_node(stack: &mut Vec<ConsumingIterItem<A>>, maybe_node: Option<PoolRef<Node<A>>>) {
1141        if let Some(node) = maybe_node {
1142            stack.push(ConsumingIterItem::Consider(PoolRef::unwrap_or_clone(node)))
1143        }
1144    }
1145
1146    fn push(stack: &mut Vec<ConsumingIterItem<A>>, mut node: Node<A>) {
1147        for _n in 0..node.keys.len() {
1148            ConsumingIter::push_node(stack, node.children.pop_back());
1149            stack.push(ConsumingIterItem::Yield(node.keys.pop_back()));
1150        }
1151        ConsumingIter::push_node(stack, node.children.pop_back());
1152    }
1153
1154    fn push_fwd(&mut self, node: Node<A>) {
1155        ConsumingIter::push(&mut self.fwd_stack, node)
1156    }
1157
1158    fn push_node_back(&mut self, maybe_node: Option<PoolRef<Node<A>>>) {
1159        if let Some(node) = maybe_node {
1160            self.back_stack
1161                .push(ConsumingIterItem::Consider(PoolRef::unwrap_or_clone(node)))
1162        }
1163    }
1164
1165    fn push_back(&mut self, mut node: Node<A>) {
1166        for _i in 0..node.keys.len() {
1167            self.push_node_back(node.children.pop_front());
1168            self.back_stack
1169                .push(ConsumingIterItem::Yield(node.keys.pop_front()));
1170        }
1171        self.push_node_back(node.children.pop_back());
1172    }
1173}
1174
1175impl<A> Iterator for ConsumingIter<A>
1176where
1177    A: BTreeValue + Clone,
1178{
1179    type Item = A;
1180
1181    fn next(&mut self) -> Option<Self::Item> {
1182        loop {
1183            match self.fwd_stack.pop() {
1184                None => {
1185                    self.remaining = 0;
1186                    return None;
1187                }
1188                Some(ConsumingIterItem::Consider(node)) => self.push_fwd(node),
1189                Some(ConsumingIterItem::Yield(value)) => {
1190                    if let Some(ref last) = self.back_last {
1191                        if value.cmp_values(last) != Ordering::Less {
1192                            self.fwd_stack.clear();
1193                            self.back_stack.clear();
1194                            self.remaining = 0;
1195                            return None;
1196                        }
1197                    }
1198                    self.remaining -= 1;
1199                    self.fwd_last = Some(value.clone());
1200                    return Some(value);
1201                }
1202            }
1203        }
1204    }
1205
1206    fn size_hint(&self) -> (usize, Option<usize>) {
1207        (self.remaining, Some(self.remaining))
1208    }
1209}
1210
1211impl<A> DoubleEndedIterator for ConsumingIter<A>
1212where
1213    A: BTreeValue + Clone,
1214{
1215    fn next_back(&mut self) -> Option<Self::Item> {
1216        loop {
1217            match self.back_stack.pop() {
1218                None => {
1219                    self.remaining = 0;
1220                    return None;
1221                }
1222                Some(ConsumingIterItem::Consider(node)) => self.push_back(node),
1223                Some(ConsumingIterItem::Yield(value)) => {
1224                    if let Some(ref last) = self.fwd_last {
1225                        if value.cmp_values(last) != Ordering::Greater {
1226                            self.fwd_stack.clear();
1227                            self.back_stack.clear();
1228                            self.remaining = 0;
1229                            return None;
1230                        }
1231                    }
1232                    self.remaining -= 1;
1233                    self.back_last = Some(value.clone());
1234                    return Some(value);
1235                }
1236            }
1237        }
1238    }
1239}
1240
1241impl<A: BTreeValue + Clone> ExactSizeIterator for ConsumingIter<A> {}
1242
1243// DiffIter
1244
1245/// An iterator over the differences between two ordered sets.
1246pub struct DiffIter<'a, A> {
1247    old_stack: Vec<IterItem<'a, A>>,
1248    new_stack: Vec<IterItem<'a, A>>,
1249}
1250
1251/// A description of a difference between two ordered sets.
1252#[derive(PartialEq, Eq, Debug)]
1253pub enum DiffItem<'a, A> {
1254    /// This value has been added to the new set.
1255    Add(&'a A),
1256    /// This value has been changed between the two sets.
1257    Update {
1258        /// The old value.
1259        old: &'a A,
1260        /// The new value.
1261        new: &'a A,
1262    },
1263    /// This value has been removed from the new set.
1264    Remove(&'a A),
1265}
1266
1267enum IterItem<'a, A> {
1268    Consider(&'a Node<A>),
1269    Yield(&'a A),
1270}
1271
1272impl<'a, A: 'a> DiffIter<'a, A> {
1273    pub(crate) fn new(old: &'a Node<A>, new: &'a Node<A>) -> Self {
1274        DiffIter {
1275            old_stack: if old.keys.is_empty() {
1276                Vec::new()
1277            } else {
1278                vec![IterItem::Consider(old)]
1279            },
1280            new_stack: if new.keys.is_empty() {
1281                Vec::new()
1282            } else {
1283                vec![IterItem::Consider(new)]
1284            },
1285        }
1286    }
1287
1288    fn push_node(stack: &mut Vec<IterItem<'a, A>>, maybe_node: &'a Option<PoolRef<Node<A>>>) {
1289        if let Some(ref node) = *maybe_node {
1290            stack.push(IterItem::Consider(node))
1291        }
1292    }
1293
1294    fn push(stack: &mut Vec<IterItem<'a, A>>, node: &'a Node<A>) {
1295        for n in 0..node.keys.len() {
1296            let i = node.keys.len() - n;
1297            Self::push_node(stack, &node.children[i]);
1298            stack.push(IterItem::Yield(&node.keys[i - 1]));
1299        }
1300        Self::push_node(stack, &node.children[0]);
1301    }
1302}
1303
1304impl<'a, A> Iterator for DiffIter<'a, A>
1305where
1306    A: 'a + BTreeValue + PartialEq,
1307{
1308    type Item = DiffItem<'a, A>;
1309
1310    fn next(&mut self) -> Option<Self::Item> {
1311        loop {
1312            match (self.old_stack.pop(), self.new_stack.pop()) {
1313                (None, None) => return None,
1314                (None, Some(new)) => match new {
1315                    IterItem::Consider(new) => Self::push(&mut self.new_stack, new),
1316                    IterItem::Yield(new) => return Some(DiffItem::Add(new)),
1317                },
1318                (Some(old), None) => match old {
1319                    IterItem::Consider(old) => Self::push(&mut self.old_stack, old),
1320                    IterItem::Yield(old) => return Some(DiffItem::Remove(old)),
1321                },
1322                (Some(old), Some(new)) => match (old, new) {
1323                    (IterItem::Consider(old), IterItem::Consider(new)) => {
1324                        if !std::ptr::eq(old, new) {
1325                            match old.keys[0].cmp_values(&new.keys[0]) {
1326                                Ordering::Less => {
1327                                    Self::push(&mut self.old_stack, old);
1328                                    self.new_stack.push(IterItem::Consider(new));
1329                                }
1330                                Ordering::Greater => {
1331                                    self.old_stack.push(IterItem::Consider(old));
1332                                    Self::push(&mut self.new_stack, new);
1333                                }
1334                                Ordering::Equal => {
1335                                    Self::push(&mut self.old_stack, old);
1336                                    Self::push(&mut self.new_stack, new);
1337                                }
1338                            }
1339                        }
1340                    }
1341                    (IterItem::Consider(old), IterItem::Yield(new)) => {
1342                        Self::push(&mut self.old_stack, old);
1343                        self.new_stack.push(IterItem::Yield(new));
1344                    }
1345                    (IterItem::Yield(old), IterItem::Consider(new)) => {
1346                        self.old_stack.push(IterItem::Yield(old));
1347                        Self::push(&mut self.new_stack, new);
1348                    }
1349                    (IterItem::Yield(old), IterItem::Yield(new)) => match old.cmp_values(new) {
1350                        Ordering::Less => {
1351                            self.new_stack.push(IterItem::Yield(new));
1352                            return Some(DiffItem::Remove(old));
1353                        }
1354                        Ordering::Equal => {
1355                            if old != new {
1356                                return Some(DiffItem::Update { old, new });
1357                            }
1358                        }
1359                        Ordering::Greater => {
1360                            self.old_stack.push(IterItem::Yield(old));
1361                            return Some(DiffItem::Add(new));
1362                        }
1363                    },
1364                },
1365            }
1366        }
1367    }
1368}
im_rc/nodes/btree.rs

im_rc/nodes/
btree.rs
