osom_lib_btree/btree/operations/
insert.rs

1#![allow(clippy::needless_return)]
2use core::ops::DerefMut;
3
4use osom_lib_primitives::length::Length;
5
6use crate::btree::{BTree, BTreeConfig, node_ptr::BTreeNodePtr};
7use crate::errors::BTreeError;
8
9use super::helpers::{self, SearchKeyResult};
10
11impl<TKey, TValue, TConfig> BTree<TKey, TValue, TConfig>
12where
13    TKey: Ord,
14    TConfig: BTreeConfig,
15{
16    /// Inserts a `(key, adder())` pair into the tree if it doesn't exist, or updates the value
17    /// with `update(value)` call.
18    ///
19    /// # Notes
20    ///
21    /// * If the key already exists, then this method discards `key`.
22    /// * This method guarantees that either `adder` or `updater` will be called, but not both.
23    ///
24    /// # Errors
25    ///
26    /// See [`BTreeError`] for details.
27    pub fn try_insert_or_update(
28        &mut self,
29        key: TKey,
30        adder: impl FnOnce() -> TValue,
31        updater: impl FnOnce(&mut TValue),
32    ) -> Result<&mut TValue, BTreeError> {
33        unsafe {
34            if self.data.is_null() {
35                let root = BTreeNodePtr::allocate_leaf_node(self.config.deref_mut())?;
36                root.insert_at(0, key, adder());
37                root.len_ptr().write(root.len() + 1);
38                let new_len = self.total_len.as_u32().unchecked_add(1);
39                self.total_len = Length::new_unchecked(new_len);
40                self.data = root;
41                return Ok(root.values_ptr().add(0).as_mut_unchecked());
42            }
43
44            if self.data.is_full() {
45                let old_root = self.data;
46                let new_root = BTreeNodePtr::allocate_internal_node(self.config.deref_mut())?;
47                new_root.children_ptr().write(old_root);
48                let (promoted_key, promoted_value, right) = old_root.split(self.config.deref_mut())?;
49                new_root.insert_at(0, promoted_key, promoted_value);
50                new_root.insert_child_at(1, right);
51                new_root.len_ptr().write(new_root.len() + 1);
52                right.parent_ptr().write(new_root);
53                old_root.parent_ptr().write(new_root);
54                self.data = new_root;
55            }
56
57            self.try_insert_or_update_non_full(self.data, key, adder, updater)
58        }
59    }
60
61    /// Inserts a `(key, value)` pair into the tree.
62    ///
63    /// Returns `None` if the key was not already present, or the previous value if it was.
64    ///
65    /// # Notes
66    ///
67    /// This method discards `key` if it already exists.
68    ///
69    /// # Errors
70    ///
71    /// See [`BTreeError`] for details.
72    pub fn try_insert(&mut self, key: TKey, value: TValue) -> Result<Option<TValue>, BTreeError> {
73        let value_ptr = &raw const value;
74        let adder = || unsafe { value_ptr.read() };
75
76        let mut result: Option<TValue> = None;
77        let updater = |current: &mut TValue| {
78            let value = unsafe { value_ptr.read() };
79            let old_value = core::mem::replace(current, value);
80            result = Some(old_value);
81        };
82
83        let _ = self.try_insert_or_update(key, adder, updater)?;
84        core::mem::forget(value);
85        Ok(result)
86    }
87
88    unsafe fn try_insert_or_update_non_full(
89        &mut self,
90        node: BTreeNodePtr<TKey, TValue, TConfig>,
91        key: TKey,
92        adder: impl FnOnce() -> TValue,
93        updater: impl FnOnce(&mut TValue),
94    ) -> Result<&mut TValue, BTreeError> {
95        unsafe {
96            match helpers::search_key(node, &key) {
97                SearchKeyResult::ExactMatch { index } => {
98                    let value_mut = node.values_ptr().add(index).as_mut_unchecked();
99                    updater(value_mut);
100                    return Ok(value_mut);
101                }
102                SearchKeyResult::InsertionIndex { mut index } => {
103                    if node.is_leaf() {
104                        if self.total_len == Length::MAX_LENGTH {
105                            return Err(BTreeError::TreeSizeOutOfRange);
106                        }
107                        node.insert_at(index, key, adder());
108                        node.len_ptr().write(node.len() + 1);
109                        let new_len = self.total_len.as_u32().unchecked_add(1);
110                        self.total_len = Length::new_unchecked(new_len);
111                        return Ok(node.values_ptr().add(index).as_mut_unchecked());
112                    }
113
114                    let mut child = node.children_ptr().add(index).read();
115                    if child.is_full() {
116                        let (promoted_key, promoted_value, right) = child.split(self.config.deref_mut())?;
117                        let equals_promoted_key = key == promoted_key;
118                        node.insert_at(index, promoted_key, promoted_value);
119                        node.insert_child_at(index + 1, right);
120                        node.len_ptr().write(node.len() + 1);
121                        right.parent_ptr().write(node);
122                        if equals_promoted_key {
123                            let value_mut = node.values_ptr().add(index).as_mut_unchecked();
124                            updater(value_mut);
125                            return Ok(value_mut);
126                        }
127                        if &key > node.keys_ptr().add(index).as_ref_unchecked() {
128                            index += 1;
129                        }
130                        child = node.children_ptr().add(index).read();
131                    }
132
133                    return self.try_insert_or_update_non_full(child, key, adder, updater);
134                }
135            }
136        }
137    }
138}
osom_lib_btree/btree/operations/insert.rs

osom_lib_btree/btree/operations/
insert.rs
