rune_alloc/btree/
remove.rs

1use crate::alloc::Allocator;
2
3use super::map::MIN_LEN;
4use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef};
5
6impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
7    /// Removes a key-value pair from the tree, and returns that pair, as well as
8    /// the leaf edge corresponding to that former pair. It's possible this empties
9    /// a root node that is internal, which the caller should pop from the map
10    /// holding the tree. The caller should also decrement the map's length.
11    pub(crate) fn remove_kv_tracking<F: FnOnce(), A: Allocator>(
12        self,
13        handle_emptied_internal_root: F,
14        alloc: &A,
15    ) -> (
16        (K, V),
17        Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
18    ) {
19        match self.force() {
20            Leaf(node) => node.remove_leaf_kv(handle_emptied_internal_root, alloc),
21            Internal(node) => node.remove_internal_kv(handle_emptied_internal_root, alloc),
22        }
23    }
24}
25
26impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
27    fn remove_leaf_kv<F: FnOnce(), A: Allocator>(
28        self,
29        handle_emptied_internal_root: F,
30        alloc: &A,
31    ) -> (
32        (K, V),
33        Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
34    ) {
35        let (old_kv, mut pos) = self.remove();
36        let len = pos.reborrow().into_node().len();
37        if len < MIN_LEN {
38            let idx = pos.idx();
39            // We have to temporarily forget the child type, because there is no
40            // distinct node type for the immediate parents of a leaf.
41            let new_pos = match pos.into_node().forget_type().choose_parent_kv() {
42                Ok(Left(left_parent_kv)) => {
43                    debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
44                    if left_parent_kv.can_merge() {
45                        left_parent_kv.merge_tracking_child_edge(Right(idx), alloc)
46                    } else {
47                        debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
48                        left_parent_kv.steal_left(idx)
49                    }
50                }
51                Ok(Right(right_parent_kv)) => {
52                    debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
53                    if right_parent_kv.can_merge() {
54                        right_parent_kv.merge_tracking_child_edge(Left(idx), alloc)
55                    } else {
56                        debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
57                        right_parent_kv.steal_right(idx)
58                    }
59                }
60                Err(pos) => unsafe { Handle::new_edge(pos, idx) },
61            };
62            // SAFETY: `new_pos` is the leaf we started from or a sibling.
63            pos = unsafe { new_pos.cast_to_leaf_unchecked() };
64
65            // Only if we merged, the parent (if any) has shrunk, but skipping
66            // the following step otherwise does not pay off in benchmarks.
67            //
68            // SAFETY: We won't destroy or rearrange the leaf where `pos` is at
69            // by handling its parent recursively; at worst we will destroy or
70            // rearrange the parent through the grandparent, thus change the
71            // link to the parent inside the leaf.
72            if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
73                if !parent
74                    .into_node()
75                    .forget_type()
76                    .fix_node_and_affected_ancestors(alloc)
77                {
78                    handle_emptied_internal_root();
79                }
80            }
81        }
82        (old_kv, pos)
83    }
84}
85
86impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
87    fn remove_internal_kv<F: FnOnce(), A: Allocator>(
88        self,
89        handle_emptied_internal_root: F,
90        alloc: &A,
91    ) -> (
92        (K, V),
93        Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
94    ) {
95        // Remove an adjacent KV from its leaf and then put it back in place of
96        // the element we were asked to remove. Prefer the left adjacent KV,
97        // for the reasons listed in `choose_parent_kv`.
98        let left_leaf_kv = self.left_edge().descend().last_leaf_edge().left_kv();
99        let left_leaf_kv = unsafe { left_leaf_kv.ok().unwrap_unchecked() };
100        let (left_kv, left_hole) = left_leaf_kv.remove_leaf_kv(handle_emptied_internal_root, alloc);
101
102        // The internal node may have been stolen from or merged. Go back right
103        // to find where the original KV ended up.
104        let mut internal = unsafe { left_hole.next_kv().ok().unwrap_unchecked() };
105        let old_kv = internal.replace_kv(left_kv.0, left_kv.1);
106        let pos = internal.next_leaf_edge();
107        (old_kv, pos)
108    }
109}