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