embed_btree/bptree/
entry.rs

1use super::iter::{IterBackward, IterForward};
2use super::*;
3use core::fmt::{self, Debug};
4
5/// Entry for an existing key-value pair in the tree
6pub struct OccupiedEntry<'a, K: Key, V: Value, S: Stats> {
7    pub(crate) inner: S::EntryInner<'a, K, V>,
8    pub(crate) leaf: LeafNode<K, V>,
9}
10
11/// Entry for a vacant key position in the tree
12pub struct VacantEntry<'a, K: Key, V: Value, S: Stats> {
13    pub(crate) inner: S::EntryInner<'a, K, V>,
14    pub(crate) leaf: Option<LeafNode<K, V>>,
15    pub(crate) key: K,
16}
17
18/// Entry into a BTreeMap for in-place manipulation
19pub enum Entry<'a, K: Key, V: Value, S: Stats> {
20    Occupied(OccupiedEntry<'a, K, V, S>),
21    Vacant(VacantEntry<'a, K, V, S>),
22}
23
24pub(crate) trait EntryInner<K: Key, V: Value> {
25    fn get_cache<'a>(&'a mut self) -> impl PathBuffer + 'a;
26
27    #[cfg(test)]
28    fn get_tree(&self) -> &BTreeInner<K, V>;
29
30    fn get_tree_cache(&mut self) -> (&mut BTreeInner<K, V>, impl PathBuffer);
31
32    fn set_idx(&mut self, idx: u8);
33
34    fn get_idx(&self) -> u8;
35}
36
37impl<'a, K: Key, V: Value, S: Stats> Entry<'a, K, V, S> {
38    #[inline]
39    pub fn exists(&self) -> bool {
40        matches!(self, Entry::Occupied(_))
41    }
42
43    /// Ensures a value is in the entry by inserting the default if empty,
44    /// and returns a mutable reference to the value in the entry.
45    #[inline]
46    pub fn or_insert(self, default: V) -> &'a mut V
47    where
48        K: Ord + 'a,
49    {
50        match self {
51            Entry::Occupied(entry) => entry.into_mut(),
52            Entry::Vacant(entry) => entry.insert(default),
53        }
54    }
55
56    /// Ensures a value is in the entry by inserting the result of the default function if empty,
57    /// and returns a mutable reference to the value in the entry.
58    #[inline]
59    pub fn or_insert_with<F>(self, default: F) -> &'a mut V
60    where
61        F: FnOnce() -> V,
62        K: Ord + 'a,
63    {
64        match self {
65            Entry::Occupied(entry) => entry.into_mut(),
66            Entry::Vacant(entry) => entry.insert(default()),
67        }
68    }
69
70    /// Returns a reference to this entry's key.
71    #[inline]
72    pub fn key(&self) -> &K {
73        match self {
74            Entry::Occupied(entry) => entry.key(),
75            Entry::Vacant(entry) => &entry.key,
76        }
77    }
78
79    /// Provides in-place mutable access to an occupied entry before any
80    /// potential inserts into the tree.
81    #[inline]
82    pub fn and_modify<F>(self, f: F) -> Self
83    where
84        F: FnOnce(&mut V),
85    {
86        match self {
87            Entry::Occupied(mut entry) => {
88                f(entry.get_mut());
89                Entry::Occupied(entry)
90            }
91            Entry::Vacant(entry) => Entry::Vacant(entry),
92        }
93    }
94
95    // NOTE: Since rust does not alloc multiple mutable borrow, the moving api should assume ownership
96
97    /// Move to previous OccupiedEntry
98    ///
99    /// When reaching the front, return the original entry in Err()
100    #[inline]
101    pub fn move_backward(self) -> Result<OccupiedEntry<'a, K, V, S>, Self> {
102        match self {
103            Entry::Occupied(ent) => match ent.move_backward() {
104                Ok(_ent) => Ok(_ent),
105                Err(_ent) => Err(Entry::Occupied(_ent)),
106            },
107            Entry::Vacant(ent) => match ent.move_backward() {
108                Ok(_ent) => Ok(_ent),
109                Err(_ent) => Err(Entry::Vacant(_ent)),
110            },
111        }
112    }
113
114    /// Move to next OccupiedEntry
115    ///
116    /// When reaching the end, return the original entry in Err()
117    #[inline]
118    pub fn move_forward(self) -> Result<OccupiedEntry<'a, K, V, S>, Self> {
119        match self {
120            Entry::Occupied(ent) => match ent.move_forward() {
121                Ok(_ent) => Ok(_ent),
122                Err(_ent) => Err(Entry::Occupied(_ent)),
123            },
124            Entry::Vacant(ent) => match ent.move_forward() {
125                Ok(_ent) => Ok(_ent),
126                Err(_ent) => Err(Entry::Vacant(_ent)),
127            },
128        }
129    }
130
131    /// Peak previous OccupiedEntry
132    #[inline(always)]
133    #[allow(clippy::needless_lifetimes)]
134    pub fn peek_backward<'b>(&'b self) -> Option<(&'b K, &'b V)> {
135        match self {
136            Entry::Occupied(ent) => ent.peek_backward(),
137            Entry::Vacant(ent) => ent.peek_backward(),
138        }
139    }
140
141    /// Peak the next OccupiedEntry
142    #[inline(always)]
143    #[allow(clippy::needless_lifetimes)]
144    pub fn peek_forward<'b>(&'b self) -> Option<(&'b K, &'b V)> {
145        match self {
146            Entry::Occupied(ent) => ent.peek_forward(),
147            Entry::Vacant(ent) => ent.peek_forward(),
148        }
149    }
150}
151
152impl<'a, K: Key + Debug, V: Value + Debug, S: Stats> Debug for OccupiedEntry<'a, K, V, S> {
153    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154        f.debug_struct("OccupiedEntry").field("key", self.key()).field("value", self.get()).finish()
155    }
156}
157
158impl<'a, K: Key + Debug, V: Value + Debug, S: Stats> Debug for VacantEntry<'a, K, V, S> {
159    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
160        f.debug_struct("VacantEntry").field("key", &self.key).finish()
161    }
162}
163
164impl<'a, K: Key + Debug, V: Value + Debug, S: Stats> Debug for Entry<'a, K, V, S> {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        match self {
167            Entry::Occupied(ent) => f.debug_tuple("Occupied").field(ent).finish(),
168            Entry::Vacant(ent) => f.debug_tuple("Vacant").field(ent).finish(),
169        }
170    }
171}
172
173impl<'a, K: Key, V: Value, S: Stats> OccupiedEntry<'a, K, V, S> {
174    /// Get a reference to the key
175    #[inline]
176    pub fn key(&self) -> &K {
177        unsafe {
178            let key_ptr = self.leaf.key_ptr(self.inner.get_idx());
179            (*key_ptr).assume_init_ref()
180        }
181    }
182
183    /// Remove the key-value pair from the tree and return the value
184    #[inline(always)]
185    pub fn remove(self) -> V {
186        self.remove_entry().1
187    }
188
189    /// Remove the key-value pair from the tree and return the key and value
190    #[inline]
191    pub fn remove_entry(self) -> (K, V) {
192        self._remove_entry(true)
193    }
194
195    /// Remove the key-value pair from the tree and return the key and value
196    #[inline(always)]
197    pub(crate) fn _remove_entry(mut self, merge: bool) -> (K, V) {
198        let (key, val) = self.leaf.remove_pair_no_borrow(self.inner.get_idx());
199        let (tree, mut cache) = self.inner.get_tree_cache();
200        tree.len -= 1;
201        // Check for underflow and handle merge
202        let new_count = self.leaf.key_count();
203        let min_count = LeafNode::<K, V>::cap() >> 1;
204        if new_count < min_count && tree.root_is_inter() {
205            // The cache should already contain the path from the entry lookup
206            tree.handle_leaf_underflow(&mut cache, self.leaf, merge);
207        }
208        (key, val)
209    }
210
211    /// Get a reference to the value
212    #[inline]
213    pub fn get(&self) -> &V {
214        unsafe {
215            let val_ptr = self.leaf.value_ptr(self.inner.get_idx());
216            (*val_ptr).assume_init_ref()
217        }
218    }
219
220    /// Get a mutable reference to the value
221    #[inline]
222    pub fn get_mut(&mut self) -> &mut V {
223        unsafe {
224            let val_ptr = self.leaf.value_ptr_mut(self.inner.get_idx());
225            (*val_ptr).assume_init_mut()
226        }
227    }
228
229    /// Convert the OccupiedEntry into a mutable reference bounded by
230    /// the tree's lifetime
231    #[inline]
232    pub fn into_mut(mut self) -> &'a mut V {
233        unsafe {
234            let val_ptr = self.leaf.value_ptr_mut(self.inner.get_idx());
235            (*val_ptr).assume_init_mut()
236        }
237    }
238
239    /// replace a value into the tree and return the old value
240    #[inline]
241    pub fn insert(&mut self, value: V) -> V {
242        self.leaf.replace(self.inner.get_idx(), value)
243    }
244
245    /// Peak previous OccupiedEntry
246    #[inline(always)]
247    #[allow(clippy::needless_lifetimes)]
248    pub fn peek_backward<'b>(&'b self) -> Option<(&'b K, &'b V)> {
249        let mut cursor =
250            IterBackward { back_leaf: self.leaf.clone(), back_idx: self.inner.get_idx() };
251        unsafe {
252            if let Some((k, v)) = cursor.prev_pair() {
253                return Some((&*k, &*v));
254            }
255        }
256        None
257    }
258
259    /// Peak the next OccupiedEntry
260    #[inline(always)]
261    #[allow(clippy::needless_lifetimes)]
262    pub fn peek_forward<'b>(&'b self) -> Option<(&'b K, &'b V)> {
263        let mut cursor =
264            IterForward { front_leaf: self.leaf.clone(), idx: self.inner.get_idx() + 1 };
265        unsafe {
266            if let Some((k, v)) = cursor.next_pair() {
267                return Some((&*k, &*v));
268            }
269        }
270        None
271    }
272
273    /// Move to previous OccupiedEntry
274    ///
275    /// When reaching the front, return the original entry in Err()
276    #[inline]
277    pub fn move_backward(mut self) -> Result<Self, Self> {
278        let idx = self.inner.get_idx();
279        if idx > 0 {
280            self.inner.set_idx(idx - 1);
281            Ok(Self { inner: self.inner, leaf: self.leaf })
282        } else if let Some(leaf) = self.leaf.get_left_node() {
283            self.inner.get_cache().move_path_left::<K>();
284            let count = leaf.key_count();
285            debug_assert!(count > 0);
286            self.inner.set_idx(count - 1);
287            Ok(Self { inner: self.inner, leaf })
288        } else {
289            Err(self)
290        }
291    }
292
293    /// Move to next OccupiedEntry
294    ///
295    /// When reaching the end, return the original entry in Err()
296    #[inline]
297    pub fn move_forward(mut self) -> Result<Self, Self> {
298        let next_idx = self.inner.get_idx() + 1;
299        if self.leaf.key_count() > next_idx {
300            self.inner.set_idx(next_idx);
301            Ok(Self { inner: self.inner, leaf: self.leaf })
302        } else if let Some(right) = self.leaf.get_right_node() {
303            self.inner.get_cache().move_path_right::<K>();
304            debug_assert!(right.key_count() > 0);
305            self.inner.set_idx(0);
306            Ok(Self { inner: self.inner, leaf: right })
307        } else {
308            Err(self)
309        }
310    }
311
312    /// Try to alter the key of this entry
313    ///
314    /// On successful returns  Ok() ;
315    /// If key is not in strict order among the neighbors, return  Err() .
316    #[inline]
317    pub fn alter_key(&mut self, k: K) -> Result<(), ()> {
318        if let Some((_k, _v)) = self.peek_backward()
319            && _k >= &k
320        {
321            return Err(());
322        }
323        if let Some((_k, _v)) = self.peek_forward()
324            && _k <= &k
325        {
326            return Err(());
327        }
328        let idx = self.inner.get_idx();
329        unsafe {
330            let k_ref = (*self.leaf.key_ptr_mut(idx)).assume_init_mut();
331            let (tree, mut cache) = self.inner.get_tree_cache();
332            if idx == 0 && tree.root_is_inter() {
333                // We need to keep the PathBuffer intact, use peek rather than move_to_ancestor
334                // it's allowed to move the entry or remove afterwards
335                tree.update_ancestor_sep_key::<false, _>(&mut cache, k.clone());
336            }
337            *k_ref = k;
338            Ok(())
339        }
340    }
341
342    #[cfg(test)]
343    pub(crate) fn validate_cache_path(&mut self) {
344        let k = self.leaf.get_keys()[self.inner.get_idx() as usize].clone();
345        if let Some(root) = self.inner.get_tree().root_as_inter() {
346            let backup = {
347                let mut cache = self.inner.get_cache();
348                cache.fix_path_center::<K>();
349                cache.to_vec()
350            };
351            // new cache, don't mix with original
352            let mut _stats = S::default();
353            let mut stack = S::BufferStack::default();
354            let mut cache = _stats.get_cache(&mut stack, Some(root));
355            let _leaf = self
356                .inner
357                .get_tree()
358                .search_leaf_with(|inter| inter.find_leaf_with_cache::<V, _, _>(&mut cache, &k))
359                .unwrap();
360            assert_eq!(self.leaf, _leaf);
361            assert_eq!(backup, cache.to_vec::<K>());
362        }
363    }
364}
365
366impl<'a, K: Key, V: Value, S: Stats> VacantEntry<'a, K, V, S> {
367    /// Get a reference to the key
368    #[inline]
369    pub fn key(&self) -> &K {
370        &self.key
371    }
372
373    /// Take ownership of the key
374    #[inline]
375    pub fn into_key(self) -> K {
376        self.key
377    }
378
379    /// Insert a value into the tree
380    #[inline]
381    pub fn insert(self, value: V) -> &'a mut V
382    where
383        K: 'a,
384    {
385        let (key, mut inner) = (self.key, self.inner);
386        let idx = inner.get_idx();
387        let (tree, mut cache) = inner.get_tree_cache();
388        if tree.root.is_none() {
389            return unsafe { &mut *tree.init_empty(key, value) };
390        }
391        tree.len += 1;
392        // Get the leaf node where we should insert
393        let mut leaf = self.leaf.expect("VacantEntry should have a node when root is not None");
394        let count = leaf.key_count();
395        // Check if leaf has space
396        let value_p = if count < LeafNode::<K, V>::cap() {
397            leaf.insert_no_split_with_idx(idx, key, value)
398        } else {
399            // Leaf is full, need to split
400            tree.insert_with_split(&mut cache, key, value, leaf, idx)
401            // NOTE: the PathBuffer might be a different path with the one inserted,
402            // because borrowing on inter node might happen, and the cache is consumed during
403            // propagate_split moves upwards.
404            // It's too complex to provide returning OccupiedEntry because the subsequence operation
405            // relies on a correct PathBuffer
406        };
407        unsafe { &mut *value_p }
408    }
409
410    /// Peak previous OccupiedEntry
411    #[inline(always)]
412    #[allow(clippy::needless_lifetimes)]
413    pub fn peek_backward<'b>(&'b self) -> Option<(&'b K, &'b V)> {
414        if let Some(leaf) = self.leaf.as_ref() {
415            // The key of previous pos is always smaller than self.key ;
416            // the key at current idx (if exists) must larger than self.key.
417            let mut cursor =
418                IterBackward { back_leaf: leaf.clone(), back_idx: self.inner.get_idx() };
419            unsafe {
420                if let Some((k, v)) = cursor.prev_pair() {
421                    return Some((&*k, &*v));
422                }
423            }
424        }
425        None
426    }
427
428    /// Peak the next OccupiedEntry
429    #[inline(always)]
430    #[allow(clippy::needless_lifetimes)]
431    pub fn peek_forward<'b>(&'b self) -> Option<(&'b K, &'b V)> {
432        if let Some(leaf) = self.leaf.as_ref() {
433            unsafe {
434                if let Some((k, v)) = leaf.get_raw_pair(self.inner.get_idx()) {
435                    // get_raw_pair will validate idx
436                    return Some((&*k, &*v));
437                }
438                if let Some(right) = leaf.get_right_node()
439                    && let Some((k, v)) = right.get_raw_pair(0)
440                {
441                    return Some((&*k, &*v));
442                }
443            }
444        }
445        None
446    }
447
448    /// Move to previous OccupiedEntry
449    ///
450    /// When reaching the front return the original entry in Err()
451    #[inline]
452    pub fn move_backward(mut self) -> Result<OccupiedEntry<'a, K, V, S>, Self> {
453        if let Some(leaf) = self.leaf.as_ref() {
454            // The key of previous pos is always smaller than self.key ;
455            // the key at current idx (if exists) must larger than self.key.
456            // It's possible the leaf.idx may be == leaf.key_count(), but it's the same.
457            let idx = self.inner.get_idx();
458            if idx > 0 {
459                self.inner.set_idx(idx - 1);
460                return Ok(OccupiedEntry { inner: self.inner, leaf: leaf.clone() });
461            }
462            if let Some(left) = leaf.get_left_node() {
463                let count = left.key_count();
464                debug_assert!(count > 0);
465                self.inner.get_cache().move_path_left::<K>();
466                self.inner.set_idx(count - 1);
467                return Ok(OccupiedEntry { inner: self.inner, leaf: left });
468            }
469        }
470        Err(self)
471    }
472
473    /// Move to next OccupiedEntry
474    ///
475    /// When reaching the end, return the original entry in Err()
476    #[inline]
477    pub fn move_forward(mut self) -> Result<OccupiedEntry<'a, K, V, S>, Self> {
478        if let Some(leaf) = self.leaf.as_ref() {
479            let idx = self.inner.get_idx();
480            if leaf.key_count() > idx {
481                // the key at current idx (if exists) must larger than self.key, no need to move
482                return Ok(OccupiedEntry { inner: self.inner, leaf: leaf.clone() });
483            } else if let Some(right) = leaf.get_right_node() {
484                debug_assert!(right.key_count() > 0);
485                self.inner.get_cache().move_path_right::<K>();
486                self.inner.set_idx(0);
487                return Ok(OccupiedEntry { inner: self.inner, leaf: right });
488            }
489        }
490        Err(self)
491    }
492}
embed_btree/bptree/entry.rs

embed_btree/bptree/
entry.rs
