Skip to main content

osom_lib_btree/btree/operations/
remove.rs

1#![allow(clippy::needless_return)]
2use core::{borrow::Borrow, ops::DerefMut};
3
4use osom_lib_macros::debug_check_or_release_hint;
5use osom_lib_primitives::{kvp::KVP, length::Length};
6
7use crate::btree::{BTree, BTreeConfig, node_ptr::BTreeNodePtr};
8
9use super::helpers;
10
11impl<TKey, TValue, TConfig> BTree<TKey, TValue, TConfig>
12where
13    TKey: Ord,
14    TConfig: BTreeConfig,
15{
16    const MIN_KVP_LEN: usize = (TConfig::CHILDREN_COUNT - 1) / 2;
17
18    /// Removes a key-value pair from the [`BTree`].
19    ///
20    /// Returns the removed key-value pair if the key was present, otherwise `None`.
21    #[inline]
22    pub fn remove<Q>(&mut self, key: &Q) -> Option<KVP<TKey, TValue>>
23    where
24        TKey: Borrow<Q>,
25        Q: Ord + ?Sized,
26    {
27        if self.data.is_null() {
28            return None;
29        }
30
31        self.remove_recursive(self.data, key)
32    }
33
34    fn remove_recursive<Q>(&mut self, node: BTreeNodePtr<TKey, TValue, TConfig>, key: &Q) -> Option<KVP<TKey, TValue>>
35    where
36        TKey: Borrow<Q>,
37        Q: Ord + ?Sized,
38    {
39        unsafe {
40            match helpers::search_key(node, key) {
41                helpers::SearchKeyResult::ExactMatch { index } => {
42                    self.total_len = Length::new_unchecked(self.total_len.as_u32() - 1);
43                    return Some(if node.is_leaf() {
44                        self.remove_leaf_node(node, index)
45                    } else {
46                        self.remove_internal_node(node, index)
47                    });
48                }
49                helpers::SearchKeyResult::InsertionIndex { index } => {
50                    if node.is_leaf() {
51                        return None;
52                    }
53
54                    let child = node.children_ptr().add(index).read();
55                    self.remove_recursive(child, key)
56                }
57            }
58        }
59    }
60
61    fn remove_leaf_node<Q>(&mut self, node: BTreeNodePtr<TKey, TValue, TConfig>, index: usize) -> KVP<TKey, TValue>
62    where
63        TKey: Borrow<Q>,
64        Q: Ord + ?Sized,
65    {
66        unsafe {
67            // Remove the (key, value) pair from the given node.
68            let key = node.keys_ptr().add(index).read();
69            let value = node.values_ptr().add(index).read();
70            node.remove_at(index);
71            let new_len = node.len() - 1;
72            node.len_ptr().write(new_len);
73
74            // We didn't go below the minimum node length, we are done.
75            if new_len > Self::MIN_KVP_LEN {
76                return KVP { key, value };
77            }
78
79            // Otherwise get the parent. If there is no parent, then this is root.
80            // Root is special and it can have arbitrarily small number of children, we are done.
81            let parent = node.parent_ptr().read();
82            if parent.is_null() {
83                return KVP { key, value };
84            }
85
86            // If we have parent, then try to steal items from left or right sibling.
87            let index = Self::get_index_in_parent(node);
88
89            if Self::try_steal_left(node, index) {
90                return KVP { key, value };
91            }
92
93            if Self::try_steal_right(node, index) {
94                return KVP { key, value };
95            }
96
97            // If we couldn't steal then we have to merge with a sibling.
98            if self.try_merge_with_left(node, index) {
99                self.maybe_shrink_root();
100                return KVP { key, value };
101            }
102
103            if self.try_merge_with_right(node, index) {
104                self.maybe_shrink_root();
105                return KVP { key, value };
106            }
107
108            // B Tree invariants ensure that one of the previous operations succeeds.
109            unreachable!("remove_leaf_node: failed to remove (key, value) pair from leaf node");
110        }
111    }
112
113    /// If the root is an internal node with no keys, replace it with its only child.
114    fn maybe_shrink_root(&mut self) {
115        unsafe {
116            let root = self.data;
117            if root.is_null() || root.is_leaf() {
118                return;
119            }
120
121            if root.len() == 0 {
122                let child = root.children_ptr().read();
123                debug_check_or_release_hint!(!child.is_null(), "maybe_shrink_root: root has no children");
124                child.parent_ptr().write(BTreeNodePtr::NULL);
125                self.data = child;
126                root.deallocate(self.config.deref_mut());
127            }
128        }
129    }
130
131    fn remove_internal_node<Q>(
132        &mut self,
133        mut node: BTreeNodePtr<TKey, TValue, TConfig>,
134        index: usize,
135    ) -> KVP<TKey, TValue>
136    where
137        TKey: Borrow<Q>,
138        Q: Ord + ?Sized,
139    {
140        unsafe {
141            // We swap the found value with either the direct predecessor or successor.
142            // That element lives in a child node. We then recursively remove it.
143            // That way we will reach a leaf node ultimately.
144            loop {
145                let child;
146                let child_key;
147                let child_value;
148                let child_index;
149
150                if index > 0 {
151                    child = node.children_ptr().add(index - 1).read();
152                    debug_check_or_release_hint!(!child.is_null(), "remove_internal_node: child is null");
153                    child_index = child.len() - 1;
154                    child_key = child.keys_ptr().add(child_index);
155                    child_value = child.values_ptr().add(child_index);
156                } else {
157                    child = node.children_ptr().add(index + 1).read();
158                    debug_check_or_release_hint!(!child.is_null(), "remove_internal_node: child is null");
159                    child_index = 0;
160                    child_key = child.keys_ptr().add(child_index);
161                    child_value = child.values_ptr().add(child_index);
162                }
163
164                let current_key = node.keys_ptr().add(index);
165                let current_value = node.values_ptr().add(index);
166                current_key.swap(child_key);
167                current_value.swap(child_value);
168
169                if child.is_leaf() {
170                    return self.remove_leaf_node(child, child_index);
171                }
172
173                node = child;
174            }
175        }
176    }
177
178    fn get_index_in_parent(node: BTreeNodePtr<TKey, TValue, TConfig>) -> usize {
179        unsafe {
180            debug_check_or_release_hint!(!node.is_null(), "get_index_in_parent: node is null");
181            let parent = node.parent_ptr().read();
182            debug_check_or_release_hint!(!parent.is_null(), "get_index_in_parent: parent is null");
183            let children = parent.children_ptr();
184            for i in 0..=parent.len() {
185                let child = children.add(i).read();
186                if child == node {
187                    return i;
188                }
189            }
190            panic!("node is not a child of its parent");
191        }
192    }
193
194    fn try_steal_left(node: BTreeNodePtr<TKey, TValue, TConfig>, index_in_parent: usize) -> bool {
195        unsafe {
196            debug_check_or_release_hint!(!node.is_null(), "try_rotate_left: node is null");
197            debug_check_or_release_hint!(node.is_leaf(), "try_rotate_left: node is not a leaf node");
198
199            if index_in_parent == 0 {
200                return false;
201            }
202
203            let parent = node.parent_ptr().read();
204            debug_check_or_release_hint!(!parent.is_null(), "try_rotate_left: parent is null");
205            let children = parent.children_ptr();
206            let left_sibling = children.add(index_in_parent - 1).read();
207            debug_check_or_release_hint!(!left_sibling.is_null(), "try_rotate_left: left sibling is null");
208            debug_check_or_release_hint!(
209                left_sibling.is_leaf(),
210                "try_rotate_left: left sibling is not a leaf node"
211            );
212            if left_sibling.len() <= Self::MIN_KVP_LEN {
213                return false;
214            }
215            let promoted_key = left_sibling.keys_ptr().add(left_sibling.len() - 1).read();
216            let promoted_value = left_sibling.values_ptr().add(left_sibling.len() - 1).read();
217
218            left_sibling.remove_at(left_sibling.len() - 1);
219            let new_len = left_sibling.len() - 1;
220            left_sibling.len_ptr().write(new_len);
221
222            let parent_key = parent.keys_ptr().add(index_in_parent - 1).read();
223            let parent_value = parent.values_ptr().add(index_in_parent - 1).read();
224            parent.keys_ptr().add(index_in_parent - 1).write(promoted_key);
225            parent.values_ptr().add(index_in_parent - 1).write(promoted_value);
226            node.insert_at(0, parent_key, parent_value);
227            node.len_ptr().write(node.len() + 1);
228            return true;
229        }
230    }
231
232    fn try_steal_right(node: BTreeNodePtr<TKey, TValue, TConfig>, index_in_parent: usize) -> bool {
233        unsafe {
234            debug_check_or_release_hint!(!node.is_null(), "try_rotate_right: node is null");
235            debug_check_or_release_hint!(node.is_leaf(), "try_rotate_right: node is not a leaf node");
236
237            let parent = node.parent_ptr().read();
238            debug_check_or_release_hint!(!parent.is_null(), "try_rotate_right: parent is null");
239
240            if index_in_parent == parent.len() {
241                return false;
242            }
243
244            let children = parent.children_ptr();
245            let right_sibling = children.add(index_in_parent + 1).read();
246            debug_check_or_release_hint!(!right_sibling.is_null(), "try_rotate_right: right sibling is null");
247            debug_check_or_release_hint!(
248                right_sibling.is_leaf(),
249                "try_rotate_right: right sibling is not a leaf node"
250            );
251            if right_sibling.len() <= Self::MIN_KVP_LEN {
252                return false;
253            }
254
255            let promoted_key = right_sibling.keys_ptr().add(0).read();
256            let promoted_value = right_sibling.values_ptr().add(0).read();
257
258            right_sibling.remove_at(0);
259            let new_len = right_sibling.len() - 1;
260            right_sibling.len_ptr().write(new_len);
261
262            let parent_key = parent.keys_ptr().add(index_in_parent).read();
263            let parent_value = parent.values_ptr().add(index_in_parent).read();
264            parent.keys_ptr().add(index_in_parent).write(promoted_key);
265            parent.values_ptr().add(index_in_parent).write(promoted_value);
266            node.insert_at(node.len(), parent_key, parent_value);
267            node.len_ptr().write(node.len() + 1);
268            return true;
269        }
270    }
271
272    fn try_merge_with_left(&mut self, node: BTreeNodePtr<TKey, TValue, TConfig>, index_in_parent: usize) -> bool {
273        unsafe {
274            debug_check_or_release_hint!(!node.is_null(), "try_rotate_right: node is null");
275            debug_check_or_release_hint!(node.is_leaf(), "try_rotate_right: node is not a leaf node");
276
277            if index_in_parent == 0 {
278                return false;
279            }
280
281            let parent = node.parent_ptr().read();
282            debug_check_or_release_hint!(!parent.is_null(), "try_rotate_right: parent is null");
283            let left_sibling = parent.children_ptr().add(index_in_parent - 1).read();
284            debug_check_or_release_hint!(!left_sibling.is_null(), "try_rotate_left: left sibling is null");
285            debug_check_or_release_hint!(
286                left_sibling.is_leaf(),
287                "try_rotate_left: left sibling is not a leaf node"
288            );
289
290            // Left sibling and current node have at most TConfig::CHILDREN_COUNT - 2 items together.
291            // So we add separator in parent to left sibling, then move all items from node
292            // to the sibling. And finally we remove separator from parent and remove node,
293            // and set the left sibling as the new child. Node is then deallocated.
294            let left_len = left_sibling.len();
295            let node_len = node.len();
296            let separator_key = parent.keys_ptr().add(index_in_parent - 1).read();
297            let separator_value = parent.values_ptr().add(index_in_parent - 1).read();
298            left_sibling.insert_at(left_len, separator_key, separator_value);
299            left_sibling
300                .keys_ptr()
301                .add(left_len + 1)
302                .copy_from_nonoverlapping(node.keys_ptr(), node_len);
303            left_sibling
304                .values_ptr()
305                .add(left_len + 1)
306                .copy_from_nonoverlapping(node.values_ptr(), node_len);
307            left_sibling.len_ptr().write(left_len + node_len + 1);
308            parent.remove_at(index_in_parent - 1);
309            parent.remove_child_at(index_in_parent);
310            parent.len_ptr().write(parent.len() - 1);
311            node.deallocate(self.config.deref_mut());
312            return true;
313        }
314    }
315
316    fn try_merge_with_right(&mut self, node: BTreeNodePtr<TKey, TValue, TConfig>, index_in_parent: usize) -> bool {
317        unsafe {
318            debug_check_or_release_hint!(!node.is_null(), "try_rotate_right: node is null");
319            debug_check_or_release_hint!(node.is_leaf(), "try_rotate_right: node is not a leaf node");
320            let parent = node.parent_ptr().read();
321            debug_check_or_release_hint!(!parent.is_null(), "try_rotate_right: parent is null");
322
323            if index_in_parent == parent.len() {
324                return false;
325            }
326
327            let right_sibling = parent.children_ptr().add(index_in_parent + 1).read();
328            debug_check_or_release_hint!(!right_sibling.is_null(), "try_rotate_right: right sibling is null");
329            debug_check_or_release_hint!(
330                right_sibling.is_leaf(),
331                "try_rotate_right: right sibling is not a leaf node"
332            );
333
334            // Merge the right sibling into `node`, mirroring `try_merge_with_left`.
335            let node_len = node.len();
336            let right_len = right_sibling.len();
337            let separator_key = parent.keys_ptr().add(index_in_parent).read();
338            let separator_value = parent.values_ptr().add(index_in_parent).read();
339            node.insert_at(node_len, separator_key, separator_value);
340            node.keys_ptr()
341                .add(node_len + 1)
342                .copy_from_nonoverlapping(right_sibling.keys_ptr(), right_len);
343            node.values_ptr()
344                .add(node_len + 1)
345                .copy_from_nonoverlapping(right_sibling.values_ptr(), right_len);
346            node.len_ptr().write(node_len + right_len + 1);
347            parent.remove_at(index_in_parent);
348            parent.remove_child_at(index_in_parent + 1);
349            parent.len_ptr().write(parent.len() - 1);
350            right_sibling.deallocate(self.config.deref_mut());
351            return true;
352        }
353    }
354}