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    /// Sets the value of the entry with the `VacantEntry`'s key, and returns `None`.
103    ///
104    /// # Errors
105    ///
106    /// Will return `Err` if the allocation fails.
107    pub fn insert(mut self, value: V) -> AllocResult<Option<V>> {
108        let alloc = self.alloc;
109        let io = self.inner.insert_in(alloc, value)?;
110        match io {
111            InsertOption::NewKey(_) => {
112                // SAFETY: Requires map is a valid non-null pointer with exclusive access
113                unsafe { self.map.as_mut() }.n += 1;
114                Ok(None)
115            }
116            InsertOption::Split((k, v, n)) => {
117                // SAFETY: Requires map is a valid non-null pointer with exclusive access
118                let map = unsafe { self.map.as_mut() };
119
120                let root =
121                    InteriorNode::root(alloc, k, v, map.root.take().unwrap(), n.into_inner())?;
122                map.root
123                    .replace(ManuallyDrop::new(Node::Interior(ManuallyDrop::into_inner(
124                        root.into_inner(),
125                    ))));
126                map.n += 1;
127
128                Ok(None)
129            }
130        }
131    }
132}
133
134impl<'a, 's, A: Allocator, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength>
135    OccupiedEntry<'a, 's, A, K, V, B>
136where
137    U2: Mul<B>,
138    Prod<U2, B>: ArrayLength,
139    U1: Add<Prod<U2, B>>,
140    Sum<U1, Prod<U2, B>>: ArrayLength,
141{
142    /// # Safety
143    ///
144    /// `alloc` MUST be the allocator that allocated `map`
145    pub(crate) unsafe fn new(
146        alloc: &'a A,
147        inner: OccupiedNodeEntry<'s, K, V, B, MutNodeRef<'s, K, V, B>>,
148        map: NonNull<AllocatedBTreeMap<K, V, B>>,
149    ) -> Self {
150        OccupiedEntry { alloc, inner, map }
151    }
152
153    /// Gets a reference to the key in the entry.
154    #[must_use]
155    pub fn key(&self) -> &K {
156        self.inner.key()
157    }
158
159    /// Gets a reference to the value in the entry.
160    #[must_use]
161    pub fn get(&self) -> &V {
162        self.inner.get()
163    }
164
165    /// Gets a mutable reference to the value in the entry.
166    pub fn get_mut(&mut self) -> &mut V {
167        self.inner.get_mut()
168    }
169
170    /// Converts the entry into a mutable reference to its value.
171    #[must_use]
172    pub fn into_mut(self) -> &'s mut V {
173        self.inner.into_mut()
174    }
175
176    /// Sets the value of the entry with the `OccupiedEntry`'s key, and returns the entry's old value.
177    pub fn insert(&mut self, value: V) -> V {
178        self.inner.insert(value)
179    }
180
181    /// Takes the value of the entry out of the map, and returns it.
182    #[must_use]
183    pub fn remove(self) -> V {
184        self.remove_entry().1
185    }
186
187    /// Takes the key-value pair out of the map, and returns it.
188    #[must_use]
189    pub fn remove_entry(mut self) -> (K, V) {
190        // SAFETY: `self.alloc` was used to create `self.inner`
191        let (entry, check_underflow) = unsafe { self.inner.remove(self.alloc) };
192
193        // SAFETY: `self` is `mut`
194        unsafe { self.map.as_mut() }.n -= 1;
195
196        if check_underflow {
197            // SAFETY: `self` is `mut`
198            let map: &mut AllocatedBTreeMap<K, V, B> = unsafe { self.map.as_mut() };
199            let root = map.root.as_mut().unwrap();
200            if root.n_keys() == 0 {
201                let old_root = map.root.take().unwrap();
202                let new_root = match ManuallyDrop::into_inner(old_root) {
203                    Node::Interior(mut n) => n.take_child_at_in(self.alloc, 0),
204                    Node::Leaf(n) => Node::Leaf(n),
205                };
206                map.root.replace(ManuallyDrop::new(new_root));
207            }
208        }
209
210        entry
211    }
212}
213
214impl<'a, 's, A: Allocator, K: core::cmp::PartialOrd + Debug, V, B: ArrayLength>
215    Entry<'a, 's, A, K, V, B>
216where
217    U2: Mul<B>,
218    Prod<U2, B>: ArrayLength,
219    U1: Add<Prod<U2, B>>,
220    Sum<U1, Prod<U2, B>>: ArrayLength,
221{
222    pub(super) fn new(
223        alloc: &'a A,
224        inner: NodeEntry<'s, K, K, V, B, MutNodeRef<'s, K, V, B>>,
225        map: NonNull<AllocatedBTreeMap<K, V, B>>,
226    ) -> Self {
227        match inner {
228            NodeEntry::Vacant(inner) => Entry::Vacant(VacantEntry { alloc, inner, map }),
229            NodeEntry::Occupied(inner) => Entry::Occupied(OccupiedEntry { alloc, inner, map }),
230        }
231    }
232
233    /// Provides in-place mutable access to an occupied entry before any potential inserts into the map.
234    #[must_use]
235    pub fn and_modify<F>(self, f: F) -> Self
236    where
237        F: FnOnce(&mut V),
238    {
239        match self {
240            Self::Occupied(mut o) => {
241                f(o.get_mut());
242                Self::Occupied(o)
243            }
244            Self::Vacant(v) => Self::Vacant(v),
245        }
246    }
247
248    /// Returns a reference to this entry's key.
249    #[inline]
250    pub fn key(&self) -> &K {
251        match self {
252            Self::Occupied(o) => o.key(),
253            Self::Vacant(v) => v.key(),
254        }
255    }
256
257    /// Ensures a value is in the entry by inserting the provided value if empty,
258    /// and returns the old value if the entry was occupied.
259    ///
260    /// # Errors
261    ///
262    /// Will return `Err` if the allocation fails.
263    pub fn insert(self, v: V) -> AllocResult<Option<V>> {
264        match self {
265            Self::Occupied(mut o) => Ok(Some(o.insert(v))),
266            Self::Vacant(o) => {
267                o.insert(v)?;
268                Ok(None)
269            }
270        }
271    }
272
273    /// Returns the `OccupiedEntry` if this entry is occupied, panics otherwise.
274    pub fn unwrap_occupied(self) -> OccupiedEntry<'a, 's, A, K, V, B> {
275        match self {
276            Self::Occupied(o) => o,
277            Self::Vacant(_) => panic!("Expected Occupied(_)"),
278        }
279    }
280}
allocated_btree/compressed/entry.rs

allocated_btree/compressed/
entry.rs
