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 peak_backward(&self) -> Option<(&'a K, &'a V)> {
123        match self {
124            Entry::Occupied(ent) => ent.peak_backward(),
125            Entry::Vacant(ent) => ent.peak_backward(),
126        }
127    }
128
129    /// Peak the next OccupiedEntry
130    #[inline(always)]
131    pub fn peak_forward(&self) -> Option<(&'a K, &'a V)> {
132        match self {
133            Entry::Occupied(ent) => ent.peak_forward(),
134            Entry::Vacant(ent) => ent.peak_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 peak_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 peak_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            return 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            return 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            return 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            return 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.peak_backward() {
306            if _k >= &k {
307                return Err(());
308            }
309        }
310        if let Some((_k, _v)) = self.peak_forward() {
311            if _k <= &k {
312                return Err(());
313            }
314        }
315        unsafe {
316            let k_ref = (*self.leaf.key_ptr(self.idx)).assume_init_mut();
317            if self.idx == 0 {
318                if self.tree._get_info().is_some() {
319                    // We need to keep the PathCache intact, use peak rather than move_to_ancenstor
320                    // it's allowed to move the entry or remove afterwards
321                    self.tree.update_ancestor_sep_key::<false>(k.clone());
322                }
323            }
324            *k_ref = k;
325            Ok(())
326        }
327    }
328
329    #[cfg(test)]
330    pub(crate) fn validate_cache_path(&self) {
331        let k = self.leaf.get_keys()[self.idx as usize].clone();
332        if let Some(info) = self.tree._get_info().as_mut() {
333            info.fix_center();
334            let backup = info.to_vec();
335            let mut _info = TreeInfo::new(info.leaf_count(), info.inter_count());
336            let _leaf = self
337                .tree
338                .search_leaf_with(|inter| inter.find_leaf_with_cache(&mut _info, &k))
339                .unwrap();
340            assert_eq!(self.leaf, _leaf);
341            assert_eq!(backup, _info.to_vec());
342        } else {
343            return;
344        }
345    }
346}
347
348impl<'a, K: Ord + Clone + Sized, V: Sized> VacantEntry<'a, K, V> {
349    /// Get a reference to the key
350    #[inline]
351    pub fn key(&self) -> &K {
352        &self.key
353    }
354
355    /// Take ownership of the key
356    #[inline]
357    pub fn into_key(self) -> K {
358        self.key
359    }
360
361    /// Insert a value into the tree
362    #[inline]
363    pub fn insert(self, value: V) -> &'a mut V {
364        let (key, tree, idx) = (self.key, self.tree, self.idx);
365        if tree.root.is_none() {
366            return tree.init_empty(key, value);
367        }
368        tree.len += 1;
369        // Get the leaf node where we should insert
370        let mut leaf = self.leaf.expect("VacantEntry should have a node when root is not None");
371        let count = leaf.key_count();
372        // Check if leaf has space
373        let value_p = if count < LeafNode::<K, V>::cap() {
374            leaf.insert_no_split_with_idx(idx, key, value)
375        } else {
376            // Leaf is full, need to split
377            tree.insert_with_split(key, value, leaf, idx)
378            // NOTE: the PathCache might be a different path with the one inserted,
379            // because borrowing on inter node might happen, and the cache is consumed during
380            // propagate_split moves upwards.
381            // It's too complex to provide returning OccupiedEntry because the subsequence operation
382            // relies on a correct PathCache
383        };
384        unsafe { &mut *value_p }
385    }
386
387    /// Peak previous OccupiedEntry
388    #[inline(always)]
389    pub fn peak_backward(&self) -> Option<(&'a K, &'a V)> {
390        if let Some(leaf) = self.leaf.as_ref() {
391            // The key of previous pos is always smaller than self.key ;
392            // the key at current idx (if exists) must larger than self.key.
393            let mut cursor = IterBackward { back_leaf: leaf.clone(), back_idx: self.idx };
394            unsafe {
395                if let Some((k, v)) = cursor.prev_pair() {
396                    return Some((&*k, &*v));
397                }
398            }
399        }
400        None
401    }
402
403    /// Peak the next OccupiedEntry
404    #[inline(always)]
405    pub fn peak_forward(&self) -> Option<(&'a K, &'a V)> {
406        if let Some(leaf) = self.leaf.as_ref() {
407            unsafe {
408                if let Some((k, v)) = leaf.get_raw_pair(self.idx) {
409                    // get_raw_pair will validate idx
410                    return Some((&*k, &*v));
411                }
412                if let Some(right) = leaf.get_right_node() {
413                    if let Some((k, v)) = right.get_raw_pair(0) {
414                        return Some((&*k, &*v));
415                    }
416                }
417            }
418        }
419        None
420    }
421
422    /// Move to previous OccupiedEntry
423    ///
424    /// When reaching the front return the original entry in Err()
425    #[inline]
426    pub fn move_backward(self) -> Result<OccupiedEntry<'a, K, V>, Self> {
427        if let Some(leaf) = self.leaf.as_ref() {
428            // The key of previous pos is always smaller than self.key ;
429            // the key at current idx (if exists) must larger than self.key.
430            // It's possible the leaf.idx may be == leaf.key_count(), but it's the same.
431            if self.idx > 0 {
432                return Ok(OccupiedEntry {
433                    tree: self.tree,
434                    leaf: leaf.clone(),
435                    idx: self.idx - 1,
436                });
437            }
438            if let Some(left) = leaf.get_left_node() {
439                let count = left.key_count();
440                debug_assert!(count > 0);
441                if let Some(info) = self.tree._get_info().as_mut() {
442                    info.move_left();
443                }
444                return Ok(OccupiedEntry { tree: self.tree, leaf: left, idx: count - 1 });
445            }
446        }
447        Err(self)
448    }
449
450    /// Move to next OccupiedEntry
451    ///
452    /// When reaching the end, return the original entry in Err()
453    #[inline]
454    pub fn move_forward(self) -> Result<OccupiedEntry<'a, K, V>, Self> {
455        if let Some(leaf) = self.leaf.as_ref() {
456            if leaf.key_count() > self.idx {
457                // the key at current idx (if exists) must larger than self.key, no need to move
458                return Ok(OccupiedEntry { tree: self.tree, leaf: leaf.clone(), idx: self.idx });
459            } else if let Some(right) = leaf.get_right_node() {
460                debug_assert!(right.key_count() > 0);
461                if let Some(info) = self.tree._get_info().as_mut() {
462                    info.move_right();
463                }
464                return Ok(OccupiedEntry { tree: self.tree, leaf: right, idx: 0 });
465            }
466        }
467        Err(self)
468    }
469}
embed_collections/btree/entry.rs

embed_collections/btree/
entry.rs
