embed_btree/btree/
entry.rs

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

embed_btree/btree/
entry.rs
