library/alloc/src/collections/btree/append.rs

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