allocated_btree/btree/
entry.rs

1use core::mem::ManuallyDrop;
2use core::ops::Add;
3use core::ops::Mul;
4use core::ptr::NonNull;
5
6use allocated::AllocResult;
7use allocator_api2::alloc::Allocator;
8use generic_array::ArrayLength;
9use typenum::Prod;
10use typenum::Sum;
11use typenum::U1;
12use typenum::U2;
13
14use super::node::InsertOption;
15use super::node::Node;
16use super::node::NodeEntry;
17use super::node::OccupiedNodeEntry;
18use super::node::VacantNodeEntry;
19use super::AllocatedBTreeMap;
20
21/// A view into a single entry in a map, which may either be vacant or occupied.
22///
23/// This enum is constructed from the [`entry`](super::wrapper::NaiveBTreeMap::entry) method on
24/// [`NaiveBTreeMap`](super::wrapper::NaiveBTreeMap).
25#[allow(clippy::module_name_repetitions)]
26pub enum Entry<'a, 's, A: Allocator, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength>
27where
28    U2: Mul<B>,
29    Prod<U2, B>: ArrayLength,
30    U1: Add<Prod<U2, B>>,
31    Sum<U1, Prod<U2, B>>: ArrayLength,
32{
33    /// A vacant entry.
34    Vacant(VacantEntry<'a, 's, A, K, V, B>),
35    /// An occupied entry.
36    Occupied(OccupiedEntry<'a, 's, A, K, V, B>),
37}
38
39/// A view into a vacant entry in a `NaiveBTreeMap`.
40/// It is part of the [`Entry`] enum.
41#[allow(clippy::module_name_repetitions)]
42pub struct VacantEntry<
43    'a,
44    's,
45    A: Allocator,
46    K: core::cmp::PartialOrd + core::fmt::Debug,
47    V,
48    B: ArrayLength,
49> where
50    U2: Mul<B>,
51    Prod<U2, B>: ArrayLength,
52    U1: Add<Prod<U2, B>>,
53    Sum<U1, Prod<U2, B>>: ArrayLength,
54{
55    alloc: &'a A,
56    inner: VacantNodeEntry<'s, K, K, V, B, &'s mut Node<K, V, B>>,
57    map: NonNull<AllocatedBTreeMap<K, V, B>>,
58}
59
60/// A view into an occupied entry in a `NaiveBTreeMap`.
61/// It is part of the [`Entry`] enum.
62#[allow(clippy::module_name_repetitions)]
63pub struct OccupiedEntry<
64    'a,
65    's,
66    A: Allocator,
67    K: core::cmp::PartialOrd + core::fmt::Debug,
68    V,
69    B: ArrayLength,
70> where
71    U2: Mul<B>,
72    Prod<U2, B>: ArrayLength,
73    U1: Add<Prod<U2, B>>,
74    Sum<U1, Prod<U2, B>>: ArrayLength,
75{
76    alloc: &'a A,
77    inner: OccupiedNodeEntry<'s, K, V, B, &'s mut Node<K, V, B>>,
78    map: NonNull<AllocatedBTreeMap<K, V, B>>,
79}
80
81impl<'a, 's, A: Allocator, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength>
82    VacantEntry<'a, 's, A, K, V, B>
83where
84    U2: Mul<B>,
85    Prod<U2, B>: ArrayLength,
86    U1: Add<Prod<U2, B>>,
87    Sum<U1, Prod<U2, B>>: ArrayLength,
88{
89    /// Gets a reference to the key that would be used when inserting a value through the `VacantEntry`.
90    pub fn key(&self) -> &K {
91        self.inner.key()
92    }
93
94    /// Takes ownership of the key.
95    pub fn into_key(self) -> K {
96        self.inner.into_key()
97    }
98
99    /// Returns a reference to the key that would be immediately below (predecessor of)
100    /// the vacant entry's key in the tree's sorted order.
101    /// Returns `None` if the vacant entry would be the minimum key.
102    #[must_use]
103    pub fn key_below(&self) -> Option<&K> {
104        self.inner.key_below()
105    }
106
107    /// Returns a reference to the key that would be immediately above (successor of)
108    /// the vacant entry's key in the tree's sorted order.
109    /// Returns `None` if the vacant entry would be the maximum key.
110    #[must_use]
111    pub fn key_above(&self) -> Option<&K> {
112        self.inner.key_above()
113    }
114
115    /// Returns a reference to the value that would be immediately below (predecessor of)
116    /// the vacant entry's key in the tree's sorted order.
117    /// Returns `None` if the vacant entry would be the minimum key.
118    #[must_use]
119    pub fn value_below(&self) -> Option<&V> {
120        self.inner.value_below()
121    }
122
123    /// Returns a reference to the value that would be immediately above (successor of)
124    /// the vacant entry's key in the tree's sorted order.
125    /// Returns `None` if the vacant entry would be the maximum key.
126    #[must_use]
127    pub fn value_above(&self) -> Option<&V> {
128        self.inner.value_above()
129    }
130
131    /// Returns a reference to the key-value pair that would be immediately below (predecessor of)
132    /// the vacant entry's key in the tree's sorted order.
133    /// Returns `None` if the vacant entry would be the minimum key.
134    #[must_use]
135    pub fn below(&self) -> Option<(&K, &V)> {
136        self.inner.below()
137    }
138
139    /// Returns a reference to the key-value pair that would be immediately above (successor of)
140    /// the vacant entry's key in the tree's sorted order.
141    /// Returns `None` if the vacant entry would be the maximum key.
142    #[must_use]
143    pub fn above(&self) -> Option<(&K, &V)> {
144        self.inner.above()
145    }
146
147    /// Sets the value of the entry with the `VacantEntry`'s key, and returns `None`.
148    ///
149    /// # Errors
150    ///
151    /// Will return `Err` if the allocation fails.
152    pub fn insert(mut self, value: V) -> AllocResult<Option<V>> {
153        // Safety: `alloc` allocated `inner` per `new` requirements
154        let io = unsafe { self.inner.insert_in(self.alloc, value)? };
155        match io {
156            InsertOption::NewKey(_) => {
157                // Safety: we own self, and this is the only reference created.
158                let map = unsafe { self.map.as_mut() };
159                map.n += 1;
160                Ok(None)
161            }
162            InsertOption::Split((k, v, n)) => {
163                // Safety: we own self, and this is the only reference created.
164                let map = unsafe { self.map.as_mut() };
165
166                let root = Node::root(self.alloc, k, v, map.root.take().unwrap(), n.into_inner())?;
167                map.root.replace(root.into_inner());
168                map.n += 1;
169
170                Ok(None)
171            }
172        }
173    }
174}
175
176impl<'a, 's, A: Allocator, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength>
177    OccupiedEntry<'a, 's, A, K, V, B>
178where
179    U2: Mul<B>,
180    Prod<U2, B>: ArrayLength,
181    U1: Add<Prod<U2, B>>,
182    Sum<U1, Prod<U2, B>>: ArrayLength,
183{
184    /// # Safety
185    ///
186    /// `alloc` MUST be the allocator that allocated `map`
187    pub(crate) unsafe fn new(
188        alloc: &'a A,
189        inner: OccupiedNodeEntry<'s, K, V, B, &'s mut Node<K, V, B>>,
190        map: NonNull<AllocatedBTreeMap<K, V, B>>,
191    ) -> Self {
192        OccupiedEntry { alloc, inner, map }
193    }
194
195    /// Gets a reference to the key in the entry.
196    #[must_use]
197    pub fn key(&self) -> &K {
198        self.inner.key()
199    }
200
201    /// Returns a reference to the key that is immediately below (predecessor of)
202    /// this occupied entry's key in the tree's sorted order.
203    /// Returns `None` if this entry is the minimum key.
204    #[must_use]
205    pub fn key_below(&self) -> Option<&K> {
206        self.inner.key_below()
207    }
208
209    /// Returns a reference to the key that is immediately above (successor of)
210    /// this occupied entry's key in the tree's sorted order.
211    /// Returns `None` if this entry is the maximum key.
212    #[must_use]
213    pub fn key_above(&self) -> Option<&K> {
214        self.inner.key_above()
215    }
216
217    /// Gets a reference to the value in the entry.
218    #[must_use]
219    pub fn get(&self) -> &V {
220        self.inner.get()
221    }
222
223    /// Gets a mutable reference to the value in the entry.
224    pub fn get_mut(&mut self) -> &mut V {
225        self.inner.get_mut()
226    }
227
228    /// Converts the entry into a mutable reference to its value.
229    #[must_use]
230    pub fn into_mut(self) -> &'s mut V {
231        self.inner.into_mut()
232    }
233
234    /// Sets the value of the entry with the `OccupiedEntry`'s key, and returns the entry's old value.
235    pub fn insert(&mut self, value: V) -> V {
236        self.inner.insert(value)
237    }
238
239    /// Takes the value of the entry out of the map, and returns it.
240    #[must_use]
241    pub fn remove(self) -> V {
242        self.remove_entry().1
243    }
244
245    /// Takes the key-value pair out of the map, and returns it.
246    #[must_use]
247    pub fn remove_entry(mut self) -> (K, V) {
248        // Safety: requirements match `new` requirements
249        let (entry, check_underflow) = unsafe { self.inner.remove(self.alloc) };
250        // Safety: we own `self`, and this is the only reference created.
251        let map: &mut AllocatedBTreeMap<K, V, B> = unsafe { self.map.as_mut() };
252        map.n -= 1;
253
254        if check_underflow {
255            let root = map.root.as_mut().unwrap();
256            // Only reduce tree height if root has 0 keys AND has a child (interior node)
257            // If root is a leaf with 0 keys, just leave it as an empty tree
258            if root.n_keys() == 0 && root.child_at(0).is_some() {
259                let mut old_root = map.root.take().unwrap();
260                let new_root = old_root.take_child_at_in(self.alloc, 0);
261                map.root.replace(ManuallyDrop::new(new_root));
262            }
263        }
264
265        entry
266    }
267}
268
269impl<'a, 's, A: Allocator, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength>
270    Entry<'a, 's, A, K, V, B>
271where
272    U2: Mul<B>,
273    Prod<U2, B>: ArrayLength,
274    U1: Add<Prod<U2, B>>,
275    Sum<U1, Prod<U2, B>>: ArrayLength,
276{
277    pub(super) fn new(
278        alloc: &'a A,
279        inner: NodeEntry<'s, K, K, V, B, &'s mut Node<K, V, B>>,
280        map: NonNull<AllocatedBTreeMap<K, V, B>>,
281    ) -> Self {
282        match inner {
283            NodeEntry::Vacant(inner) => Entry::Vacant(VacantEntry { alloc, inner, map }),
284            NodeEntry::Occupied(inner) => Entry::Occupied(OccupiedEntry { alloc, inner, map }),
285        }
286    }
287
288    /// Provides in-place mutable access to an occupied entry before any potential inserts into the map.
289    #[must_use]
290    pub fn and_modify<F>(self, f: F) -> Self
291    where
292        F: FnOnce(&mut V),
293    {
294        match self {
295            Self::Occupied(mut o) => {
296                f(o.get_mut());
297                Self::Occupied(o)
298            }
299            Self::Vacant(v) => Self::Vacant(v),
300        }
301    }
302
303    /// Returns a reference to this entry's key.
304    #[inline]
305    pub fn key(&self) -> &K {
306        match self {
307            Self::Occupied(o) => o.key(),
308            Self::Vacant(v) => v.key(),
309        }
310    }
311
312    /// Ensures a value is in the entry by inserting the provided value if empty,
313    /// and returns the old value if the entry was occupied.
314    ///
315    /// # Errors
316    ///
317    /// Will return `Err` if the allocation fails.
318    pub fn insert(self, v: V) -> AllocResult<Option<V>> {
319        match self {
320            Self::Occupied(mut o) => Ok(Some(o.insert(v))),
321            Self::Vacant(o) => {
322                o.insert(v)?;
323                Ok(None)
324            }
325        }
326    }
327
328    /// Returns the `OccupiedEntry` if this entry is occupied, panics otherwise.
329    pub fn unwrap_occupied(self) -> OccupiedEntry<'a, 's, A, K, V, B> {
330        match self {
331            Self::Occupied(o) => o,
332            Self::Vacant(_) => panic!("Expected Occupied(_)"),
333        }
334    }
335}
allocated_btree/btree/entry.rs

allocated_btree/btree/
entry.rs
