embed_collections/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    pub fn peek_backward(&self) -> Option<(&'a K, &'a V)> {
123        match self {
124            Entry::Occupied(ent) => ent.peek_backward(),
125            Entry::Vacant(ent) => ent.peek_backward(),
126        }
127    }
128
129    /// Peak the next OccupiedEntry
130    #[inline(always)]
131    pub fn peek_forward(&self) -> Option<(&'a K, &'a V)> {
132        match self {
133            Entry::Occupied(ent) => ent.peek_forward(),
134            Entry::Vacant(ent) => ent.peek_forward(),
135        }
136    }
137}
138
139impl<'a, K: Ord + Clone + Sized + fmt::Debug, V: Sized + fmt::Debug> fmt::Debug
140    for OccupiedEntry<'a, K, V>
141{
142    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
143        f.debug_struct("OccupiedEntry").field("key", self.key()).field("value", self.get()).finish()
144    }
145}
146
147impl<'a, K: Ord + Clone + Sized + fmt::Debug, V: Sized + fmt::Debug> fmt::Debug
148    for VacantEntry<'a, K, V>
149{
150    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151        f.debug_struct("VacantEntry").field("key", &self.key).finish()
152    }
153}
154
155impl<'a, K: Ord + Clone + Sized + fmt::Debug, V: Sized + fmt::Debug> fmt::Debug
156    for Entry<'a, K, V>
157{
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        match self {
160            Entry::Occupied(ent) => f.debug_tuple("Occupied").field(ent).finish(),
161            Entry::Vacant(ent) => f.debug_tuple("Vacant").field(ent).finish(),
162        }
163    }
164}
165
166impl<'a, K: Ord + Clone + Sized, V: Sized> OccupiedEntry<'a, K, V> {
167    /// Get a reference to the key
168    #[inline]
169    pub fn key(&self) -> &K {
170        unsafe {
171            let key_ptr = self.leaf.key_ptr(self.idx);
172            (*key_ptr).assume_init_ref()
173        }
174    }
175
176    /// Remove the key-value pair from the tree and return the value
177    #[inline(always)]
178    pub fn remove(self) -> V {
179        self.remove_entry().1
180    }
181
182    /// Remove the key-value pair from the tree and return the key and value
183    #[inline]
184    pub fn remove_entry(self) -> (K, V) {
185        self._remove_entry(true)
186    }
187
188    /// Remove the key-value pair from the tree and return the key and value
189    #[inline(always)]
190    pub(crate) fn _remove_entry(mut self, merge: bool) -> (K, V) {
191        let (key, val) = self.leaf.remove_pair_no_borrow(self.idx);
192        self.tree.len -= 1;
193        // Check for underflow and handle merge
194        let new_count = self.leaf.key_count();
195        let min_count = LeafNode::<K, V>::cap() >> 1;
196        if new_count < min_count && self.tree.root_is_inter() {
197            // The cache should already contain the path from the entry lookup
198            self.tree.handle_leaf_underflow(self.leaf, merge);
199        }
200        (key, val)
201    }
202
203    /// Get a reference to the value
204    #[inline]
205    pub fn get(&self) -> &V {
206        unsafe {
207            let val_ptr = self.leaf.value_ptr(self.idx);
208            (*val_ptr).assume_init_ref()
209        }
210    }
211
212    /// Get a mutable reference to the value
213    #[inline]
214    pub fn get_mut(&mut self) -> &mut V {
215        unsafe {
216            let val_ptr = self.leaf.value_ptr(self.idx);
217            (*val_ptr).assume_init_mut()
218        }
219    }
220
221    /// Convert the OccupiedEntry into a mutable reference bounded by
222    /// the tree's lifetime
223    #[inline]
224    pub fn into_mut(self) -> &'a mut V {
225        unsafe {
226            let val_ptr = self.leaf.value_ptr(self.idx);
227            (*val_ptr).assume_init_mut()
228        }
229    }
230
231    /// replace a value into the tree and return the old value
232    #[inline]
233    pub fn insert(&mut self, value: V) -> V {
234        self.leaf.replace(self.idx, value)
235    }
236
237    /// Peak previous OccupiedEntry
238    #[inline(always)]
239    pub fn peek_backward(&self) -> Option<(&'a K, &'a V)> {
240        let mut cursor = IterBackward { back_leaf: self.leaf.clone(), back_idx: self.idx };
241        unsafe {
242            if let Some((k, v)) = cursor.prev_pair() {
243                return Some((&*k, &*v));
244            }
245        }
246        None
247    }
248
249    /// Peak the next OccupiedEntry
250    #[inline(always)]
251    pub fn peek_forward(&self) -> Option<(&'a K, &'a V)> {
252        let mut cursor = IterForward { front_leaf: self.leaf.clone(), idx: self.idx + 1 };
253        unsafe {
254            if let Some((k, v)) = cursor.next_pair() {
255                return Some((&*k, &*v));
256            }
257        }
258        None
259    }
260
261    /// Move to previous OccupiedEntry
262    ///
263    /// When reaching the front, return the original entry in Err()
264    #[inline]
265    pub fn move_backward(self) -> Result<Self, Self> {
266        if self.idx > 0 {
267            Ok(Self { tree: self.tree, leaf: self.leaf, idx: self.idx - 1 })
268        } else if let Some(leaf) = self.leaf.get_left_node() {
269            if let Some(info) = self.tree._get_info().as_mut() {
270                info.move_left();
271            }
272            let count = leaf.key_count();
273            debug_assert!(count > 0);
274            Ok(Self { tree: self.tree, leaf, idx: count - 1 })
275        } else {
276            Err(self)
277        }
278    }
279
280    /// Move to next OccupiedEntry
281    ///
282    /// When reaching the end, return the original entry in Err()
283    #[inline]
284    pub fn move_forward(self) -> Result<Self, Self> {
285        let next_idx = self.idx + 1;
286        if self.leaf.key_count() > next_idx {
287            Ok(Self { tree: self.tree, leaf: self.leaf, idx: next_idx })
288        } else if let Some(right) = self.leaf.get_right_node() {
289            if let Some(info) = self.tree._get_info().as_mut() {
290                info.move_right();
291            }
292            debug_assert!(right.key_count() > 0);
293            Ok(Self { tree: self.tree, leaf: right, idx: 0 })
294        } else {
295            Err(self)
296        }
297    }
298
299    /// Try to alter the key of this entry
300    ///
301    /// On successful returns  Ok() ;
302    /// If key is not in strict order among the neighbors, return  Err() .
303    #[inline]
304    pub fn alter_key(&mut self, k: K) -> Result<(), ()> {
305        if let Some((_k, _v)) = self.peek_backward()
306            && _k >= &k
307        {
308            return Err(());
309        }
310        if let Some((_k, _v)) = self.peek_forward()
311            && _k <= &k
312        {
313            return Err(());
314        }
315        unsafe {
316            let k_ref = (*self.leaf.key_ptr(self.idx)).assume_init_mut();
317            if self.idx == 0 && self.tree._get_info().is_some() {
318                // We need to keep the PathCache intact, use peek rather than move_to_ancenstor
319                // it's allowed to move the entry or remove afterwards
320                self.tree.update_ancestor_sep_key::<false>(k.clone());
321            }
322            *k_ref = k;
323            Ok(())
324        }
325    }
326
327    #[cfg(test)]
328    pub(crate) fn validate_cache_path(&self) {
329        let k = self.leaf.get_keys()[self.idx as usize].clone();
330        if let Some(info) = self.tree._get_info().as_mut() {
331            info.fix_center();
332            let backup = info.to_vec();
333            let mut _info = TreeInfo::new(info.leaf_count(), info.inter_count());
334            let _leaf = self
335                .tree
336                .search_leaf_with(|inter| inter.find_leaf_with_cache(&mut _info, &k))
337                .unwrap();
338            assert_eq!(self.leaf, _leaf);
339            assert_eq!(backup, _info.to_vec());
340        } else {
341            return;
342        }
343    }
344}
345
346impl<'a, K: Ord + Clone + Sized, V: Sized> VacantEntry<'a, K, V> {
347    /// Get a reference to the key
348    #[inline]
349    pub fn key(&self) -> &K {
350        &self.key
351    }
352
353    /// Take ownership of the key
354    #[inline]
355    pub fn into_key(self) -> K {
356        self.key
357    }
358
359    /// Insert a value into the tree
360    #[inline]
361    pub fn insert(self, value: V) -> &'a mut V {
362        let (key, tree, idx) = (self.key, self.tree, self.idx);
363        if tree.root.is_none() {
364            return tree.init_empty(key, value);
365        }
366        tree.len += 1;
367        // Get the leaf node where we should insert
368        let mut leaf = self.leaf.expect("VacantEntry should have a node when root is not None");
369        let count = leaf.key_count();
370        // Check if leaf has space
371        let value_p = if count < LeafNode::<K, V>::cap() {
372            leaf.insert_no_split_with_idx(idx, key, value)
373        } else {
374            // Leaf is full, need to split
375            tree.insert_with_split(key, value, leaf, idx)
376            // NOTE: the PathCache might be a different path with the one inserted,
377            // because borrowing on inter node might happen, and the cache is consumed during
378            // propagate_split moves upwards.
379            // It's too complex to provide returning OccupiedEntry because the subsequence operation
380            // relies on a correct PathCache
381        };
382        unsafe { &mut *value_p }
383    }
384
385    /// Peak previous OccupiedEntry
386    #[inline(always)]
387    pub fn peek_backward(&self) -> Option<(&'a K, &'a V)> {
388        if let Some(leaf) = self.leaf.as_ref() {
389            // The key of previous pos is always smaller than self.key ;
390            // the key at current idx (if exists) must larger than self.key.
391            let mut cursor = IterBackward { back_leaf: leaf.clone(), back_idx: self.idx };
392            unsafe {
393                if let Some((k, v)) = cursor.prev_pair() {
394                    return Some((&*k, &*v));
395                }
396            }
397        }
398        None
399    }
400
401    /// Peak the next OccupiedEntry
402    #[inline(always)]
403    pub fn peek_forward(&self) -> Option<(&'a K, &'a V)> {
404        if let Some(leaf) = self.leaf.as_ref() {
405            unsafe {
406                if let Some((k, v)) = leaf.get_raw_pair(self.idx) {
407                    // get_raw_pair will validate idx
408                    return Some((&*k, &*v));
409                }
410                if let Some(right) = leaf.get_right_node()
411                    && let Some((k, v)) = right.get_raw_pair(0)
412                {
413                    return Some((&*k, &*v));
414                }
415            }
416        }
417        None
418    }
419
420    /// Move to previous OccupiedEntry
421    ///
422    /// When reaching the front return the original entry in Err()
423    #[inline]
424    pub fn move_backward(self) -> Result<OccupiedEntry<'a, K, V>, Self> {
425        if let Some(leaf) = self.leaf.as_ref() {
426            // The key of previous pos is always smaller than self.key ;
427            // the key at current idx (if exists) must larger than self.key.
428            // It's possible the leaf.idx may be == leaf.key_count(), but it's the same.
429            if self.idx > 0 {
430                return Ok(OccupiedEntry {
431                    tree: self.tree,
432                    leaf: leaf.clone(),
433                    idx: self.idx - 1,
434                });
435            }
436            if let Some(left) = leaf.get_left_node() {
437                let count = left.key_count();
438                debug_assert!(count > 0);
439                if let Some(info) = self.tree._get_info().as_mut() {
440                    info.move_left();
441                }
442                return Ok(OccupiedEntry { tree: self.tree, leaf: left, idx: count - 1 });
443            }
444        }
445        Err(self)
446    }
447
448    /// Move to next OccupiedEntry
449    ///
450    /// When reaching the end, return the original entry in Err()
451    #[inline]
452    pub fn move_forward(self) -> Result<OccupiedEntry<'a, K, V>, Self> {
453        if let Some(leaf) = self.leaf.as_ref() {
454            if leaf.key_count() > self.idx {
455                // the key at current idx (if exists) must larger than self.key, no need to move
456                return Ok(OccupiedEntry { tree: self.tree, leaf: leaf.clone(), idx: self.idx });
457            } else if let Some(right) = leaf.get_right_node() {
458                debug_assert!(right.key_count() > 0);
459                if let Some(info) = self.tree._get_info().as_mut() {
460                    info.move_right();
461                }
462                return Ok(OccupiedEntry { tree: self.tree, leaf: right, idx: 0 });
463            }
464        }
465        Err(self)
466    }
467}
embed_collections/btree/entry.rs

embed_collections/btree/
entry.rs
