allocated_btree/compressed/
entry.rs

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

allocated_btree/compressed/
entry.rs
