grove/trees/
avl.rs

1//! Implementation of AVL trees.
2//! Balanced by keeping track of node ranks, this is a worst-case balancing
3//! Algorithm that has a small memory overhead per node.
4
5use crate::locators;
6
7use super::basic_tree::*;
8use super::*;
9
10/// The type that is used for rank bookkeeping.
11/// `u8` is definitely enough, since the rank of the tree is logarithmic in the tree size.
12type T = u8;
13/// Used for rank differences
14type TD = i8;
15
16/// An AVL tree. Balanced by keeping track of node ranks, this is a worst-case balancing
17/// Algorithm that has a small memory overhead per node.
18pub struct AVLTree<D: Data> {
19    tree: BasicTree<D, T>,
20}
21
22/// For implementing `rank`, `rank_diff` and `rebuild_ranks` for
23/// trees, nodes and walkers alike.
24trait Rankable {
25    fn rank(&self) -> T;
26
27    /// Returns `true` if the rank of the current node had to be updated,
28    /// `false` if it was correct.
29    fn rebuild_ranks(&mut self) -> bool;
30
31    /// Returns `right.rank() - left.rank()`
32    fn rank_diff(&self) -> TD;
33}
34
35impl<D: Data> Rankable for BasicTree<D, T> {
36    fn rank(&self) -> T {
37        match self.node() {
38            None => 0,
39            Some(node) => node.rank(),
40        }
41    }
42
43    fn rebuild_ranks(&mut self) -> bool {
44        if let Some(node) = self.node_mut() {
45            node.rebuild_ranks()
46        } else {
47            true
48        }
49    }
50
51    /// Returns `right.rank() - left.rank()`
52    fn rank_diff(&self) -> TD {
53        match self.node() {
54            None => 0,
55            Some(node) => node.rank_diff(),
56        }
57    }
58}
59
60impl<D: Data> Rankable for BasicNode<D, T> {
61    fn rank(&self) -> T {
62        *self.alg_data()
63    }
64
65    /// Returns `right.rank() - left.rank()`
66    fn rank_diff(&self) -> TD {
67        let diff = self.right.rank() as TD - self.left.rank() as TD;
68        if self.action().to_reverse() {
69            -diff
70        } else {
71            diff
72        }
73    }
74
75    fn rebuild_ranks(&mut self) -> bool {
76        let new_rank = std::cmp::max(self.left.rank(), self.right.rank()) + 1;
77        let changed = self.rank() != new_rank;
78        self.alg_data = new_rank;
79        changed
80    }
81}
82
83impl<D: Data> AVLTree<D> {
84    /// Creates an empty [`AVLTree`].
85    pub fn new() -> Self {
86        AVLTree {
87            tree: BasicTree::Empty,
88        }
89    }
90
91    /// Asserts that the ranks at the current node are correct.
92    /// Otherwise, panics.
93    pub fn assert_ranks_locally(&self) {
94        if let Some(node) = self.tree.node() {
95            Self::assert_ranks_locally_internal(node);
96        }
97    }
98
99    fn assert_ranks_locally_internal(node: &BasicNode<D, T>) {
100        assert!(node.rank() == node.left.rank() + 1 || node.rank() == node.right.rank() + 1);
101        assert!(node.left.rank() == node.rank() - 1 || node.left.rank() == node.rank() - 2);
102        assert!(node.right.rank() == node.rank() - 1 || node.right.rank() == node.rank() - 2);
103    }
104
105    /// Asserts that the tree's ranks are correct.
106    /// Otherwise, panics.
107    pub fn assert_ranks(&self) {
108        self.tree
109            .assert_correctness_with(Self::assert_ranks_locally_internal);
110    }
111}
112
113impl<D: Data> Rankable for AVLTree<D> {
114    fn rank(&self) -> T {
115        self.tree.rank()
116    }
117
118    /// Returns `right.rank() - left.rank()`
119    fn rank_diff(&self) -> TD {
120        self.tree.rank_diff()
121    }
122
123    fn rebuild_ranks(&mut self) -> bool {
124        self.tree.rebuild_ranks()
125    }
126}
127
128impl<D: Data> Default for AVLTree<D> {
129    fn default() -> Self {
130        AVLTree::new()
131    }
132}
133
134impl<D: Data> SomeTree<D> for AVLTree<D> {
135    fn segment_summary_imm<L>(&self, locator: L) -> D::Summary
136    where
137        L: crate::Locator<D>,
138        D::Value: Clone,
139    {
140        segment_algorithms::segment_summary_imm(&self.tree, locator)
141    }
142
143    fn segment_summary<L>(&mut self, locator: L) -> D::Summary
144    where
145        L: crate::Locator<D>,
146    {
147        segment_algorithms::segment_summary(self, locator)
148    }
149
150    fn act_segment<L>(&mut self, action: D::Action, locator: L)
151    where
152        L: crate::Locator<D>,
153    {
154        if !action.to_reverse() {
155            segment_algorithms::act_segment(self, action, locator)
156        } else {
157            // split out the middle
158            let mut mid: AVLTree<D> = self
159                .slice(locators::LeftEdgeOf(locator.clone()))
160                .split_right()
161                .unwrap();
162
163            let mut walker2 = AVLWalker {
164                walker: BasicWalker::new_with_context(
165                    &mut mid.tree,
166                    self.subtree_summary(),
167                    Default::default(),
168                ),
169            };
170            walker2.search_subtree(locators::RightEdgeOf(locator));
171            let right = walker2.split_right().unwrap();
172            drop(walker2);
173
174            // apply action
175            mid.act_subtree(action);
176
177            // glue back together
178            mid.concatenate_right(right);
179            self.concatenate_right(mid);
180        }
181    }
182
183    type TreeData = u8;
184    fn iter_locator<'a, L: locators::Locator<D>>(
185        &'a mut self,
186        locator: L,
187    ) -> basic_tree::iterators::IterLocator<'a, D, L, u8> {
188        iterators::IterLocator::new(&mut self.tree, locator)
189    }
190
191    fn assert_correctness(&self)
192    where
193        D::Summary: Eq,
194    {
195        self.tree.assert_correctness_with(|node| {
196            node.assert_correctness_locally();
197            Self::assert_ranks_locally_internal(node);
198        });
199    }
200}
201
202impl<'a, D: Data> SomeTreeRef<D> for &'a mut AVLTree<D> {
203    type Walker = AVLWalker<'a, D>;
204
205    fn walker(self) -> Self::Walker {
206        AVLWalker {
207            walker: self.tree.walker(),
208        }
209    }
210}
211
212impl<'a, D: Data> ModifiableTreeRef<D> for &'a mut AVLTree<D> {
213    type ModifiableWalker = AVLWalker<'a, D>;
214}
215
216impl<'a, D: Data> SplittableTreeRef<D> for &'a mut AVLTree<D> {
217    type T = AVLTree<D>;
218
219    type SplittableWalker = AVLWalker<'a, D>;
220}
221
222derive_SomeEntry! {tree, T,
223    impl<D: Data> SomeEntry<D> for AVLTree<D> {
224        fn assert_correctness_locally(&self)
225        where
226            D::Summary: Eq,
227        {
228            if let Some(node) = self.tree.node() {
229                Self::assert_ranks_locally_internal(node);
230                node.assert_correctness_locally();
231            }
232        }
233    }
234}
235
236impl<D: Data> std::iter::FromIterator<D::Value> for AVLTree<D> {
237    /// This takes `O(n)` worst-case time.
238    fn from_iter<T: IntoIterator<Item = D::Value>>(iter: T) -> Self {
239        // TODO: check if inserting is O(1) amortized. if it is, we can do this by
240        // just calling insert.
241        // if not, than this is `O(n log n)` worst-case time.
242
243        let mut tree: AVLTree<D> = Default::default();
244        let mut walker = tree.walker();
245        for val in iter.into_iter() {
246            // note: this relies on the assumption, that after we insert a node, the new position of the locator
247            // will be an ancestor of the location where the value was inserted.
248            while walker.go_right().is_ok() {}
249            walker.insert(val);
250        }
251        drop(walker);
252        tree
253    }
254}
255
256impl<D: Data> IntoIterator for AVLTree<D> {
257    type Item = D::Value;
258    type IntoIter = iterators::IntoIter<D, std::ops::RangeFull, T>;
259
260    fn into_iter(self) -> Self::IntoIter {
261        iterators::IntoIter::new(self.tree, ..)
262    }
263}
264
265/// A walker struct for [`AVLTree`].
266pub struct AVLWalker<'a, D: Data> {
267    walker: BasicWalker<'a, D, T>,
268}
269
270impl<'a, D: Data> std::ops::Drop for AVLWalker<'a, D> {
271    fn drop(&mut self) {
272        self.go_to_root()
273    }
274}
275
276derive_SomeWalker! {walker,
277    impl<'a, D: Data> SomeWalker<D> for AVLWalker<'a, D> {
278        fn go_up(&mut self) -> Result<Side, ()> {
279            let res = self.walker.go_up()?;
280            let changed = self.inner_mut().rebuild_ranks();
281            assert!(!changed); // it shouldn't have changed without being rebalanced already
282            Ok(res)
283        }
284    }
285}
286
287derive_SomeEntry! {walker, T,
288    impl<'a, D: Data> SomeEntry<D> for AVLWalker<'a, D> {
289        fn assert_correctness_locally(&self)
290        where
291            D::Summary: Eq,
292        {
293            self.walker.assert_correctness_locally();
294            if let Some(node) = self.walker.node() {
295                AVLTree::assert_ranks_locally_internal(node);
296            }
297        }
298    }
299}
300
301impl<'a, D: Data> Rankable for AVLWalker<'a, D> {
302    /// Returns the priority of the current node. Lower numbers means
303    /// The node is closer to the root.
304    fn rank(&self) -> T {
305        match self.walker.node() {
306            None => 0,
307            Some(node) => *node.alg_data(),
308        }
309    }
310
311    /// Returns `right.rank() - left.rank()`
312    fn rank_diff(&self) -> TD {
313        self.walker.inner().rank_diff()
314    }
315
316    fn rebuild_ranks(&mut self) -> bool {
317        self.inner_mut().rebuild_ranks()
318    }
319}
320
321impl<'a, D: Data> AVLWalker<'a, D> {
322    fn inner(&self) -> &BasicTree<D, T> {
323        self.walker.inner()
324    }
325
326    fn inner_mut(&mut self) -> &mut BasicTree<D, T> {
327        self.walker.inner_mut()
328    }
329
330    fn rot_left(&mut self) -> Option<()> {
331        let rebuilder = |node: &mut BasicNode<D, T>| {
332            node.rebuild_ranks();
333        };
334        self.walker.rot_left_with_custom_rebuilder(rebuilder)
335    }
336
337    fn rot_right(&mut self) -> Option<()> {
338        let rebuilder = |node: &mut BasicNode<D, T>| {
339            node.rebuild_ranks();
340        };
341        self.walker.rot_right_with_custom_rebuilder(rebuilder)
342    }
343
344    fn rot_up(&mut self) -> Result<Side, ()> {
345        let rebuilder = |node: &mut BasicNode<D, T>| {
346            node.rebuild_ranks();
347        };
348        self.walker.rot_up_with_custom_rebuilder(rebuilder)
349    }
350
351    // For completeness this function is still here. It might be used in future versions.
352    #[allow(dead_code)]
353    fn rot_side(&mut self, side: Side) -> Option<()> {
354        let rebuilder = |node: &mut BasicNode<D, T>| {
355            node.rebuild_ranks();
356        };
357        self.walker.rot_side_with_custom_rebuilder(side, rebuilder)
358    }
359
360    /// This function gets called when a node is deleted or inserted,
361    /// at the current position.
362    fn rebalance(&mut self) {
363        if self.is_empty() {
364            let res = self.walker.go_up(); // ranks may be incorrect, so go up with the inner walker
365            if res.is_err() {
366                return;
367            }
368        }
369
370        self.rebuild_ranks();
371
372        loop {
373            let node = self.inner().node().unwrap();
374            match self.rank_diff() {
375                -2 => {
376                    // -2, left is deeper
377                    if node.left.rank_diff() <= 0 {
378                        // left left case
379                        self.rot_right().unwrap();
380                    } else {
381                        // left.rank() = 1, left right case
382                        self.go_left().unwrap();
383                        self.rot_left().unwrap();
384                        let res = self.rot_up();
385                        assert!(res == Ok(Side::Left));
386                    }
387                }
388
389                -1..=1 => {} // do nothing, the current node is now balanced.
390
391                2 => {
392                    // 2, left is shallower
393                    if node.right.rank_diff() >= 0 {
394                        // right right case
395                        self.rot_left().unwrap();
396                    } else {
397                        // right.rank() = -1, right left case
398                        self.go_right().unwrap();
399                        self.rot_right().unwrap();
400                        let res = self.rot_up();
401                        assert!(res == Ok(Side::Right));
402                    }
403                }
404
405                rd => panic!("illegal rank difference: {}", rd),
406            }
407
408            // current node has been balanced. now go up a node,
409            // and check if we need to continue rebalancing.
410            let res = self.walker.go_up(); // ranks may be incorrect, so go up with the inner walker
411            let changed = self.inner_mut().rebuild_ranks();
412            let rd = self.inner().rank_diff();
413            if !changed && -1 <= rd && rd <= 1 {
414                // tree is now balanced correctly
415                break;
416            }
417            if res.is_err() {
418                // reached root
419                break;
420            }
421        }
422    }
423
424    /// Deletes a node and returns it with the box.
425    /// The walker reorganizes the current subtree in order to delete the current node,
426    /// and then rebalances. During rebalancing it may only go up the tree.
427    fn delete_boxed(&mut self) -> Option<Box<BasicNode<D, T>>> {
428        // the delete implementation is modified from `BasicTree`,
429        // in order for rebalancing to be done properly.
430        let mut node = self.walker.take_subtree().into_node_boxed()?;
431        if node.right.is_empty() {
432            self.walker.put_subtree(node.left).unwrap();
433            node.left = BasicTree::Empty;
434            self.rebalance();
435        } else {
436            // find the next node and move it to the current position
437            let mut walker = node.right.walker();
438            while walker.go_left().is_ok() {}
439            let res = walker.go_up();
440            assert_eq!(res, Ok(Side::Left));
441
442            let mut boxed_replacement_node = walker.take_subtree().into_node_boxed().unwrap();
443            assert!(boxed_replacement_node.left.is_empty());
444            walker.put_subtree(boxed_replacement_node.right).unwrap();
445            AVLWalker { walker }.rebalance(); // rebalance here
446
447            boxed_replacement_node.left = node.left;
448            node.left = BasicTree::Empty;
449            boxed_replacement_node.right = node.right;
450            node.right = BasicTree::Empty;
451            boxed_replacement_node.rebuild();
452            self.walker
453                .put_subtree(BasicTree::from_boxed_node(boxed_replacement_node))
454                .unwrap();
455            self.rebalance(); // rebalance here
456        }
457        Some(node)
458    }
459}
460
461impl<'a, D: Data> ModifiableWalker<D> for AVLWalker<'a, D> {
462    /// Inserts the value into the tree at the current empty position.
463    /// If the current position is not empty, return [`None`].
464    /// When the function returns, the walker will be at a position which is an ancestor of the
465    /// newly inserted node.
466    fn insert(&mut self, val: D::Value) -> Option<()> {
467        self.walker
468            .insert_with_alg_data(val, 1 /* rank of a node with no sons */)?;
469        self.rebalance();
470        Some(())
471    }
472
473    
474    /// The walker reorganizes the current subtree in order to delete the current node,
475    /// and then rebalances. During rebalancing it may only go up the tree.
476    fn delete(&mut self) -> Option<D::Value> {
477        Some(self.delete_boxed()?.node_value)
478    }
479}
480
481impl<'a, D: Data> SplittableWalker<D> for AVLWalker<'a, D> {
482    type T = AVLTree<D>;
483
484    /// Will only do anything if the current position is empty.
485    /// If it is empty, it will split the tree: the elements
486    /// to the left will remain, and the elements to the right
487    /// will be put in the new output tree.
488    /// The walker will be at the root after this operation, if it succeeds.
489    ///
490    ///```
491    /// use grove::{SomeTree, avl::AVLTree};
492    /// use grove::example_data::StdNum;
493    ///
494    /// let mut tree: AVLTree<StdNum> = (17..88).collect();
495    /// let mut tree2 = tree.slice(7..7).split_right().unwrap();
496    ///
497    /// assert_eq!(tree.iter().cloned().collect::<Vec<_>>(), (17..24).collect::<Vec<_>>());
498    /// assert_eq!(tree2.iter().cloned().collect::<Vec<_>>(), (24..88).collect::<Vec<_>>());
499    /// # tree.assert_correctness();
500    ///```
501    fn split_right(&mut self) -> Option<Self::T> {
502        if !self.is_empty() {
503            return None;
504        }
505        let mut left = AVLTree::new();
506        let mut right = AVLTree::new();
507
508        // ranks may be incorrect, so go up with the inner walker
509        while let Ok(side) = self.walker.go_up() {
510            let mut node = self.walker.take_subtree().into_node_boxed().unwrap();
511            match side {
512                Side::Left => {
513                    assert!(node.left.is_empty());
514                    let auxiliary_right = AVLTree { tree: node.right };
515                    node.right = BasicTree::Empty;
516                    right.concatenate_boxed_middle_right(node, auxiliary_right);
517                }
518                Side::Right => {
519                    assert!(node.right.is_empty());
520                    let auxiliary_left = AVLTree { tree: node.left };
521                    node.left = BasicTree::Empty;
522                    left.concatenate_boxed_middle_left(auxiliary_left, node);
523                }
524            }
525        }
526
527        // the `self` tree is empty by this point.
528        self.walker.put_subtree(left.tree).unwrap();
529        Some(right)
530    }
531
532    /// Will only do anything if the current position is empty.
533    /// If it is empty, it will split the tree: the elements
534    /// to the left will remain, and the elements to the right
535    /// will be put in the new output tree.
536    /// The walker will be at the root after this operation, if it succeeds.
537    ///
538    ///```
539    /// use grove::{SomeTree, avl::AVLTree};
540    /// use grove::example_data::StdNum;
541    ///
542    /// let mut tree: AVLTree<StdNum> = (17..88).collect();
543    /// let mut tree2 = tree.slice(7..7).split_left().unwrap();
544    ///
545    /// assert_eq!(tree2.iter().cloned().collect::<Vec<_>>(), (17..24).collect::<Vec<_>>());
546    /// assert_eq!(tree.iter().cloned().collect::<Vec<_>>(), (24..88).collect::<Vec<_>>());
547    /// # tree.assert_correctness();
548    ///```
549    fn split_left(&mut self) -> Option<Self::T> {
550        let mut right = self.split_right()?;
551        std::mem::swap(&mut right.tree, self.inner_mut());
552        Some(right)
553    }
554}
555
556impl<D: Data> AVLTree<D> {
557    /// Concatenates the trees together, in place, with a given value for the middle.
558    /// Complexity: `O(log n)`. More precisely, `O(dr)` where `dr` is the difference of ranks between the two trees.
559    ///```
560    /// use grove::{SomeTree, avl::AVLTree};
561    /// use grove::example_data::StdNum;
562    ///
563    /// let mut tree: AVLTree<StdNum> = (17..=89).collect();
564    /// let tree2: AVLTree<StdNum> = (13..=25).collect();
565    /// tree.concatenate_middle_right(5, tree2);
566    ///
567    /// assert_eq!(tree.iter().cloned().collect::<Vec<_>>(), (17..=89).chain(5..=5).chain(13..=25).collect::<Vec<_>>());
568    /// # tree.assert_correctness();
569    ///```
570    pub fn concatenate_middle_right(&mut self, mid: D::Value, right: AVLTree<D>) {
571        let node = BasicNode::new_alg(mid, 0 /* dummy value */);
572        self.concatenate_boxed_middle_right(Box::new(node), right);
573    }
574
575    fn concatenate_boxed_middle_right(
576        &mut self,
577        mut mid: Box<BasicNode<D, T>>,
578        mut right: AVLTree<D>,
579    ) {
580        if self.rank() < right.rank() {
581            std::mem::swap(self, &mut right);
582            self.concatenate_boxed_middle_left(right, mid);
583            return;
584        }
585        let mut walker = self.walker();
586        while walker.rank() > right.rank() {
587            walker.go_right().unwrap();
588        }
589        mid.alg_data = 0;
590        mid.left = walker.walker.take_subtree();
591        mid.right = right.tree;
592        mid.rebuild();
593        walker.walker.put_subtree(BasicTree::from_boxed_node(mid)).unwrap();
594        walker.rebalance();
595    }
596
597    /// Concatenates the trees together, in place, with a given value for the middle.
598    /// Complexity: `O(log n)`. More precisely, `O(dr)` where `dr` is the difference of ranks between the two trees.
599    ///```
600    /// use grove::{SomeTree, avl::AVLTree};
601    /// use grove::example_data::StdNum;
602    ///
603    /// let tree1: AVLTree<StdNum> = (17..=89).collect();
604    /// let mut tree2: AVLTree<StdNum> = (13..=25).collect();
605    /// tree2.concatenate_middle_left(tree1, 5);
606    ///
607    /// assert_eq!(tree2.iter().cloned().collect::<Vec<_>>(), (17..=89).chain(5..=5).chain(13..=25).collect::<Vec<_>>());
608    /// # tree2.assert_correctness();
609    ///```
610    pub fn concatenate_middle_left(&mut self, left: AVLTree<D>, mid: D::Value) {
611        let node = BasicNode::new_alg(mid, 0 /* dummy value */);
612        self.concatenate_boxed_middle_left(left, Box::new(node));
613    }
614
615    fn concatenate_boxed_middle_left(
616        &mut self,
617        mut left: AVLTree<D>,
618        mut mid: Box<BasicNode<D, T>>,
619    ) {
620        if self.rank() < left.rank() {
621            std::mem::swap(self, &mut left);
622            self.concatenate_boxed_middle_right(mid, left);
623            return;
624        }
625        let mut walker = self.walker();
626        while walker.rank() > left.rank() {
627            walker.go_left().unwrap();
628        }
629        mid.alg_data = 0;
630        mid.right = walker.walker.take_subtree();
631        mid.left = left.tree;
632        mid.rebuild();
633        walker.walker.put_subtree(BasicTree::from_boxed_node(mid)).unwrap();
634        walker.rebalance();
635    }
636}
637
638impl<D: Data> ConcatenableTree<D> for AVLTree<D> {
639    /// Concatenates the trees together, in place.
640    /// Complexity: `O(log n)`.
641    ///```
642    /// use grove::{SomeTree, ConcatenableTree, avl::AVLTree};
643    /// use grove::example_data::StdNum;
644    ///
645    /// let mut tree: AVLTree<StdNum> = (17..=89).collect();
646    /// let tree2: AVLTree<StdNum> = (13..=25).collect();
647    /// tree.concatenate_right(tree2);
648    ///
649    /// assert_eq!(tree.iter().cloned().collect::<Vec<_>>(), (17..=89).chain(13..=25).collect::<Vec<_>>());
650    /// # tree.assert_correctness();
651    ///```
652    fn concatenate_right(&mut self, mut right: Self) {
653        if !right.is_empty() {
654            let mut walker = right.search(locators::LeftEdgeOf(..));
655            walker.go_up().unwrap();
656            let mid = walker.delete_boxed().unwrap();
657            drop(walker);
658            self.concatenate_boxed_middle_right(mid, right);
659        }
660    }
661}
662
663/// Concatenates the trees together, in place, with a given value for the middle.
664/// Complexity: `O(log n)`. More precisely, `O(dr)` where `dr` is the difference of ranks between the two trees.
665///```
666/// use grove::{SomeTree, avl::AVLTree, avl::concatenate_with_middle};
667/// use grove::example_data::StdNum;
668///
669/// let tree1: AVLTree<StdNum> = (17..=89).collect();
670/// let tree2: AVLTree<StdNum> = (13..=25).collect();
671/// let mut tree3 = concatenate_with_middle(tree1, 5, tree2);
672///
673/// assert_eq!(tree3.iter().cloned().collect::<Vec<_>>(), (17..=89).chain(5..=5).chain(13..=25).collect::<Vec<_>>());
674/// # tree3.assert_correctness();
675///```
676pub fn concatenate_with_middle<D: Data>(
677    mut left: AVLTree<D>,
678    mid: D::Value,
679    right: AVLTree<D>,
680) -> AVLTree<D> {
681    left.concatenate_middle_right(mid, right);
682    left
683}
grove/trees/avl.rs

grove/trees/
avl.rs
