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

imbl/nodes/
btree.rs
