rune_alloc/btree/
append.rs

1use core::iter::FusedIterator;
2
3use crate::alloc::{AllocError, Allocator};
4#[cfg(test)]
5use crate::testing::*;
6
7use super::merge_iter::MergeIterInner;
8use super::node::{self, Root};
9
10impl<K, V> Root<K, V> {
11    /// Appends all key-value pairs from the union of two ascending iterators,
12    /// incrementing a `length` variable along the way. The latter makes it
13    /// easier for the caller to avoid a leak when a drop handler panicks.
14    ///
15    /// If both iterators produce the same key, this method drops the pair from
16    /// the left iterator and appends the pair from the right iterator.
17    ///
18    /// If you want the tree to end up in a strictly ascending order, like for
19    /// a `BTreeMap`, both iterators should produce keys in strictly ascending
20    /// order, each greater than all keys in the tree, including any keys
21    /// already in the tree upon entry.
22    pub(crate) fn try_append_from_sorted_iters<I, A: Allocator>(
23        &mut self,
24        left: I,
25        right: I,
26        length: &mut usize,
27        alloc: &A,
28    ) -> Result<(), AllocError>
29    where
30        K: Ord,
31        I: Iterator<Item = (K, V)> + FusedIterator,
32    {
33        // We prepare to merge `left` and `right` into a sorted sequence in linear time.
34        let iter = MergeIter(MergeIterInner::new(left, right));
35
36        // Meanwhile, we build a tree from the sorted sequence in linear time.
37        self.try_bulk_push(iter, length, alloc)
38    }
39
40    /// Pushes all key-value pairs to the end of the tree, incrementing a
41    /// `length` variable along the way. The latter makes it easier for the
42    /// caller to avoid a leak when the iterator panicks.
43    pub(crate) fn try_bulk_push<I, A: Allocator>(
44        &mut self,
45        iter: I,
46        length: &mut usize,
47        alloc: &A,
48    ) -> Result<(), AllocError>
49    where
50        I: Iterator<Item = (K, V)>,
51    {
52        let mut cur_node = self.borrow_mut().last_leaf_edge().into_node();
53        // Iterate through all key-value pairs, pushing them into nodes at the right level.
54        for (key, value) in iter {
55            // Try to push key-value pair into the current leaf node.
56            if cur_node.len() < node::CAPACITY {
57                cur_node.push(key, value);
58            } else {
59                // No space left, go up and push there.
60                let mut open_node;
61                let mut test_node = cur_node.forget_type();
62                loop {
63                    match test_node.ascend() {
64                        Ok(parent) => {
65                            let parent = parent.into_node();
66                            if parent.len() < node::CAPACITY {
67                                // Found a node with space left, push here.
68                                open_node = parent;
69                                break;
70                            } else {
71                                // Go up again.
72                                test_node = parent.forget_type();
73                            }
74                        }
75                        Err(_) => {
76                            // We are at the top, create a new root node and push there.
77                            open_node = self.push_internal_level(alloc)?;
78                            break;
79                        }
80                    }
81                }
82
83                // Push key-value pair and new right subtree.
84                let tree_height = open_node.height() - 1;
85                let mut right_tree = Root::new(alloc)?;
86
87                for _ in 0..tree_height {
88                    right_tree.push_internal_level(alloc)?;
89                }
90
91                open_node.push(key, value, right_tree);
92
93                // Go down to the right-most leaf again.
94                cur_node = open_node.forget_type().last_leaf_edge().into_node();
95            }
96
97            // Increment length every iteration, to make sure the map drops
98            // the appended elements even if advancing the iterator panicks.
99            *length += 1;
100        }
101        self.fix_right_border_of_plentiful();
102        Ok(())
103    }
104
105    #[cfg(test)]
106    pub(crate) fn bulk_push<I, A: Allocator>(&mut self, iter: I, length: &mut usize, alloc: &A)
107    where
108        I: Iterator<Item = (K, V)>,
109    {
110        self.try_bulk_push(iter, length, alloc).abort()
111    }
112}
113
114// An iterator for merging two sorted sequences into one
115struct MergeIter<K, V, I: Iterator<Item = (K, V)>>(MergeIterInner<I>);
116
117impl<K: Ord, V, I> Iterator for MergeIter<K, V, I>
118where
119    I: Iterator<Item = (K, V)> + FusedIterator,
120{
121    type Item = (K, V);
122
123    /// If two keys are equal, returns the key-value pair from the right source.
124    fn next(&mut self) -> Option<(K, V)> {
125        let (a_next, b_next) = self.0.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0));
126        b_next.or(a_next)
127    }
128}