remove.rs - source
Description: Source of the Rust file `library/alloc/src/collections/btree/remove.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
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151

    use super::map::MIN_LEN;
    use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef};
    use super::unwrap_unchecked;

    impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
        /// Removes a key-value pair from the tree, and returns that pair, as well as
        /// the leaf edge corresponding to that former pair. It's possible this empties
        /// a root node that is internal, which the caller should pop from the map
        /// holding the tree. The caller should also decrement the map's length.
        pub fn remove_kv_tracking<F: FnOnce()>(
            self,
            handle_emptied_internal_root: F,
        ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
            match self.force() {
                Leaf(node) => node.remove_leaf_kv(handle_emptied_internal_root),
                Internal(node) => node.remove_internal_kv(handle_emptied_internal_root),
            }
        }
    }

    impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
        fn remove_leaf_kv<F: FnOnce()>(
            self,
            handle_emptied_internal_root: F,
        ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
            let (old_kv, mut pos) = self.remove();
            let len = pos.reborrow().into_node().len();
            if len < MIN_LEN {
                let idx = pos.idx();
                // We have to temporarily forget the child type, because there is no
                // distinct node type for the immediate parents of a leaf.
                let new_pos = match pos.into_node().forget_type().choose_parent_kv() {
                    Ok(Left(left_parent_kv)) => {
                        debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
                        if left_parent_kv.can_merge() {
                            left_parent_kv.merge(Some(Right(idx)))
                        } else {
                            debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
                            left_parent_kv.steal_left(idx)
                        }
                    }
                    Ok(Right(right_parent_kv)) => {
                        debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
                        if right_parent_kv.can_merge() {
                            right_parent_kv.merge(Some(Left(idx)))
                        } else {
                            debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
                            right_parent_kv.steal_right(idx)
                        }
                    }
                    Err(pos) => unsafe { Handle::new_edge(pos, idx) },
                };
                // SAFETY: `new_pos` is the leaf we started from or a sibling.
                pos = unsafe { new_pos.cast_to_leaf_unchecked() };

                // Only if we merged, the parent (if any) has shrunk, but skipping
                // the following step otherwise does not pay off in benchmarks.
                //
                // SAFETY: We won't destroy or rearrange the leaf where `pos` is at
                // by handling its parent recursively; at worst we will destroy or
                // rearrange the parent through the grandparent, thus change the
                // link to the parent inside the leaf.
                if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
                    parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root);
                }
            }
            (old_kv, pos)
        }
    }

    impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
        fn remove_internal_kv<F: FnOnce()>(
            self,
            handle_emptied_internal_root: F,
        ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
            // Remove an adjacent KV from its leaf and then put it back in place of
            // the element we were asked to remove. Prefer the left adjacent KV,
            // for the reasons listed in `choose_parent_kv`.
            let left_leaf_kv = self.left_edge().descend().last_leaf_edge().left_kv();
            let left_leaf_kv = unsafe { unwrap_unchecked(left_leaf_kv.ok()) };
            let (left_kv, left_hole) = left_leaf_kv.remove_leaf_kv(handle_emptied_internal_root);

            // The internal node may have been stolen from or merged. Go back right
            // to find where the original KV ended up.
            let mut internal = unsafe { unwrap_unchecked(left_hole.next_kv().ok()) };
            let old_kv = internal.replace_kv(left_kv.0, left_kv.1);
            let pos = internal.next_leaf_edge();
            (old_kv, pos)
        }
    }

    impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
        /// Stocks up a possibly underfull internal node and its ancestors,
        /// until it reaches an ancestor that has elements to spare or is the root.
        fn handle_shrunk_node_recursively<F: FnOnce()>(mut self, handle_emptied_internal_root: F) {
            loop {
                self = match self.len() {
                    0 => {
                        // An empty node must be the root, because length is only
                        // reduced by one, and non-root underfull nodes are stocked up,
                        // so non-root nodes never have fewer than MIN_LEN - 1 elements.
                        debug_assert!(self.ascend().is_err());
                        handle_emptied_internal_root();
                        return;
                    }
                    1..MIN_LEN => {
                        if let Some(parent) = self.handle_underfull_node_locally() {
                            parent
                        } else {
                            return;
                        }
                    }
                    _ => return,
                }
            }
        }

        /// Stocks up an underfull internal node, possibly at the cost of shrinking
        /// its parent instead, which is then returned.
        fn handle_underfull_node_locally(
            self,
        ) -> Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>> {
            match self.forget_type().choose_parent_kv() {
                Ok(Left(left_parent_kv)) => {
                    debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
                    if left_parent_kv.can_merge() {
                        let pos = left_parent_kv.merge(None);
                        let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
                        Some(parent_edge.into_node())
                    } else {
                        debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
                        left_parent_kv.steal_left(0);
                        None
                    }
                }
                Ok(Right(right_parent_kv)) => {
                    debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
                    if right_parent_kv.can_merge() {
                        let pos = right_parent_kv.merge(None);
                        let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
                        Some(parent_edge.into_node())
                    } else {
                        debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
                        right_parent_kv.steal_right(0);
                        None
                    }
                }
                Err(_) => None,
            }
        }
    }