append.rs - source
Description: Source of the Rust file `library/alloc/src/collections/btree/append.rs`.
Keywords: rust, rustlang, rust-lang

      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124

    use super::map::MIN_LEN;
    use super::merge_iter::MergeIterInner;
    use super::node::{self, ForceResult::*, Root};
    use core::iter::FusedIterator;

    impl<K, V> Root<K, V> {
        /// Appends all key-value pairs from the union of two ascending iterators,
        /// incrementing a `length` variable along the way. The latter makes it
        /// easier for the caller to avoid a leak when a drop handler panicks.
        ///
        /// If both iterators produce the same key, this method drops the pair from
        /// the left iterator and appends the pair from the right iterator.
        ///
        /// If you want the tree to end up in a strictly ascending order, like for
        /// a `BTreeMap`, both iterators should produce keys in strictly ascending
        /// order, each greater than all keys in the tree, including any keys
        /// already in the tree upon entry.
        pub fn append_from_sorted_iters<I>(&mut self, left: I, right: I, length: &mut usize)
        where
            K: Ord,
            I: Iterator<Item = (K, V)> + FusedIterator,
        {
            // We prepare to merge `left` and `right` into a sorted sequence in linear time.
            let iter = MergeIter(MergeIterInner::new(left, right));

            // Meanwhile, we build a tree from the sorted sequence in linear time.
            self.bulk_push(iter, length)
        }

        /// Pushes all key-value pairs to the end of the tree, incrementing a
        /// `length` variable along the way. The latter makes it easier for the
        /// caller to avoid a leak when the iterator panicks.
        fn bulk_push<I>(&mut self, iter: I, length: &mut usize)
        where
            I: Iterator<Item = (K, V)>,
        {
            let mut cur_node = self.node_as_mut().last_leaf_edge().into_node();
            // Iterate through all key-value pairs, pushing them into nodes at the right level.
            for (key, value) in iter {
                // Try to push key-value pair into the current leaf node.
                if cur_node.len() < node::CAPACITY {
                    cur_node.push(key, value);
                } else {
                    // No space left, go up and push there.
                    let mut open_node;
                    let mut test_node = cur_node.forget_type();
                    loop {
                        match test_node.ascend() {
                            Ok(parent) => {
                                let parent = parent.into_node();
                                if parent.len() < node::CAPACITY {
                                    // Found a node with space left, push here.
                                    open_node = parent;
                                    break;
                                } else {
                                    // Go up again.
                                    test_node = parent.forget_type();
                                }
                            }
                            Err(_) => {
                                // We are at the top, create a new root node and push there.
                                open_node = self.push_internal_level();
                                break;
                            }
                        }
                    }

                    // Push key-value pair and new right subtree.
                    let tree_height = open_node.height() - 1;
                    let mut right_tree = Root::new_leaf();
                    for _ in 0..tree_height {
                        right_tree.push_internal_level();
                    }
                    open_node.push(key, value, right_tree);

                    // Go down to the right-most leaf again.
                    cur_node = open_node.forget_type().last_leaf_edge().into_node();
                }

                // Increment length every iteration, to make sure the map drops
                // the appended elements even if advancing the iterator panicks.
                *length += 1;
            }
            self.fix_right_edge();
        }

        fn fix_right_edge(&mut self) {
            // Handle underfull nodes, start from the top.
            let mut cur_node = self.node_as_mut();
            while let Internal(internal) = cur_node.force() {
                // Check if right-most child is underfull.
                let mut last_edge = internal.last_edge();
                let right_child_len = last_edge.reborrow().descend().len();
                if right_child_len < MIN_LEN {
                    // We need to steal.
                    let mut last_kv = match last_edge.left_kv() {
                        Ok(left) => left,
                        Err(_) => unreachable!(),
                    };
                    last_kv.bulk_steal_left(MIN_LEN - right_child_len);
                    last_edge = last_kv.right_edge();
                }

                // Go further down.
                cur_node = last_edge.descend();
            }
        }
    }

    // An iterator for merging two sorted sequences into one
    struct MergeIter<K, V, I: Iterator<Item = (K, V)>>(MergeIterInner<I>);

    impl<K: Ord, V, I> Iterator for MergeIter<K, V, I>
    where
        I: Iterator<Item = (K, V)> + FusedIterator,
    {
        type Item = (K, V);

        /// If two keys are equal, returns the key-value pair from the right source.
        fn next(&mut self) -> Option<(K, V)> {
            let (a_next, b_next) = self.0.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0));
            b_next.or(a_next)
        }
    }