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.get_root_unwrap().is_leaf() {
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        unsafe {
235            let val_ptr = self.leaf.value_ptr(self.idx);
236            let old = (*val_ptr).assume_init_read();
237            (*val_ptr).write(value);
238            old
239        }
240    }
241
242    /// Peak previous OccupiedEntry
243    #[inline(always)]
244    pub fn peak_backward(&self) -> Option<(&'a K, &'a V)> {
245        let mut cursor = IterBackward { back_leaf: self.leaf.clone(), back_idx: self.idx };
246        unsafe {
247            if let Some((k, v)) = cursor.prev_pair() {
248                return Some((&*k, &*v));
249            }
250        }
251        None
252    }
253
254    /// Peak the next OccupiedEntry
255    #[inline(always)]
256    pub fn peak_forward(&self) -> Option<(&'a K, &'a V)> {
257        let mut cursor = IterForward { front_leaf: self.leaf.clone(), idx: self.idx + 1 };
258        unsafe {
259            if let Some((k, v)) = cursor.next_pair() {
260                return Some((&*k, &*v));
261            }
262        }
263        None
264    }
265
266    /// Move to previous OccupiedEntry
267    ///
268    /// When reaching the front, return the original entry in Err()
269    #[inline]
270    pub fn move_backward(self) -> Result<Self, Self> {
271        let mut cursor = IterBackward { back_leaf: self.leaf.clone(), back_idx: self.idx };
272        if let Some((prev_leaf, idx)) = cursor.prev() {
273            self.tree.get_cache().move_left();
274            return Ok(Self { tree: self.tree, leaf: prev_leaf.clone(), idx });
275        }
276        Err(self)
277    }
278
279    /// Move to next OccupiedEntry
280    ///
281    /// When reaching the end, return the original entry in Err()
282    #[inline]
283    pub fn move_forward(self) -> Result<Self, Self> {
284        let mut cursor = IterForward { front_leaf: self.leaf.clone(), idx: self.idx + 1 };
285        if let Some((next_leaf, idx)) = cursor.next() {
286            self.tree.get_cache().move_right();
287            return Ok(Self { tree: self.tree, leaf: next_leaf.clone(), idx });
288        }
289        Err(self)
290    }
291
292    /// Try to alter the key of this entry
293    ///
294    /// On successful returns  Ok() ;
295    /// If key is not in strict order among the neighbors, return  Err() .
296    #[inline]
297    pub fn alter_key(&mut self, k: K) -> Result<(), ()> {
298        if let Some((_k, _v)) = self.peak_backward() {
299            if _k >= &k {
300                return Err(());
301            }
302        }
303        if let Some((_k, _v)) = self.peak_forward() {
304            if _k <= &k {
305                return Err(());
306            }
307        }
308        unsafe {
309            let k_ref = (*self.leaf.key_ptr(self.idx)).assume_init_mut();
310            if self.idx == 0 {
311                self.tree.update_ancestor_sep_key::<false>(k.clone());
312            }
313            *k_ref = k;
314            Ok(())
315        }
316    }
317}
318
319impl<'a, K: Ord + Clone + Sized, V: Sized> VacantEntry<'a, K, V> {
320    /// Get a reference to the key
321    #[inline]
322    pub fn key(&self) -> &K {
323        &self.key
324    }
325
326    /// Take ownership of the key
327    #[inline]
328    pub fn into_key(self) -> K {
329        self.key
330    }
331
332    /// Insert a value into the tree
333    #[inline]
334    pub fn insert(self, value: V) -> &'a mut V {
335        let (key, tree, idx) = (self.key, self.tree, self.idx);
336        if tree.root.is_none() {
337            unsafe {
338                // empty tree
339                let mut leaf = LeafNode::<K, V>::alloc();
340                tree.root = Some(Node::Leaf(leaf.clone()));
341                tree.len = 1;
342                tree.leaf_count += 1;
343                return &mut *leaf.insert_no_split_with_idx(0, key, value);
344            }
345        }
346        tree.len += 1;
347        // Get the leaf node where we should insert
348        let mut leaf = self.leaf.expect("VacantEntry should have a node when root is not None");
349        let count = leaf.key_count();
350        // Check if leaf has space
351        let value_p = if count < LeafNode::<K, V>::cap() {
352            leaf.insert_no_split_with_idx(idx, key, value)
353        } else {
354            // Leaf is full, need to split
355            tree.insert_with_split(key, value, leaf, idx)
356            // NOTE: the PathCache might be a different path with the one inserted,
357            // because borrowing on inter node might happen, and the cache is consumed during
358            // propagate_split moves upwards.
359            // It's too complex to provide returning OccupiedEntry because the subsequence operation
360            // relies on a correct PathCache
361        };
362        unsafe { &mut *value_p }
363    }
364
365    /// Peak previous OccupiedEntry
366    #[inline(always)]
367    pub fn peak_backward(&self) -> Option<(&'a K, &'a V)> {
368        if let Some(leaf) = self.leaf.as_ref() {
369            // The key of previous pos is always smaller than self.key ;
370            // the key at current idx (if exists) must larger than self.key.
371            let mut cursor = IterBackward { back_leaf: leaf.clone(), back_idx: self.idx };
372            unsafe {
373                if let Some((k, v)) = cursor.prev_pair() {
374                    return Some((&*k, &*v));
375                }
376            }
377        }
378        None
379    }
380
381    /// Peak the next OccupiedEntry
382    #[inline(always)]
383    pub fn peak_forward(&self) -> Option<(&'a K, &'a V)> {
384        if let Some(leaf) = self.leaf.as_ref() {
385            unsafe {
386                if let Some((k, v)) = leaf.get_raw_pair(self.idx) {
387                    return Some((&*k, &*v));
388                }
389                if let Some(right) = leaf.get_right_node() {
390                    if let Some((k, v)) = right.get_raw_pair(0) {
391                        return Some((&*k, &*v));
392                    }
393                }
394            }
395        }
396        None
397    }
398
399    /// Move to previous OccupiedEntry
400    ///
401    /// When reaching the front return the original entry in Err()
402    #[inline]
403    pub fn move_backward(self) -> Result<OccupiedEntry<'a, K, V>, Self> {
404        if let Some(leaf) = self.leaf.as_ref() {
405            let mut cursor = IterBackward { back_leaf: leaf.clone(), back_idx: self.idx };
406            // The key of previous pos is always smaller than self.key ;
407            // the key at current idx (if exists) must larger than self.key.
408            // Be aware the leaf.idx may be larger than leaf.key_count(), but IterBackward
409            // does not care.
410            if cursor.prev().is_some() {
411                self.tree.get_cache().move_left();
412                return Ok(OccupiedEntry {
413                    tree: self.tree,
414                    leaf: cursor.back_leaf.clone(),
415                    idx: cursor.back_idx,
416                });
417            }
418        }
419        Err(self)
420    }
421
422    /// Move to next OccupiedEntry
423    ///
424    /// When reaching the end, return the original entry in Err()
425    #[inline]
426    pub fn move_forward(self) -> Result<OccupiedEntry<'a, K, V>, Self> {
427        if let Some(leaf) = self.leaf.as_ref() {
428            if leaf.key_count() > self.idx {
429                return Ok(OccupiedEntry { tree: self.tree, leaf: leaf.clone(), idx: self.idx });
430            } else if let Some(right) = leaf.get_right_node() {
431                debug_assert!(right.key_count() > 0);
432                self.tree.get_cache().move_right();
433                return Ok(OccupiedEntry { tree: self.tree, leaf: right, idx: 0 });
434            }
435        }
436        Err(self)
437    }
438}
embed_collections/btree/entry.rs

embed_collections/btree/
entry.rs
