augmented_rbtree/entry.rs
1//! The [`Entry`] API for in-place manipulation of tree entries.
2//!
3//! This mirrors the `Entry` API found in [`alloc::collections::BTreeMap`].
4
5use core::ptr::NonNull;
6
7use crate::{
8 alloc_proxy::proxy::{Allocator, Layout, handle_alloc_error},
9 augmented_rbtree::{OutOfMemoryError, internal_details::AugmentedRBTreeInt},
10 node::Node,
11 policy::internal_details::TreePolicy,
12};
13
14/// A view into a single entry in a tree, which may either be vacant or occupied.
15///
16/// This `enum` is constructed from the `AugmentedRBTreeInt::entry` method on
17/// [`AugmentedRBTree`](crate::AugmentedRBTree).
18#[derive(Debug)]
19pub enum Entry<'a, K, V, S, A: Allocator, P: TreePolicy<K = K, V = V, S = S>>
20where
21 K: Ord,
22{
23 /// An occupied entry.
24 Occupied(OccupiedEntry<'a, K, V, S, A, P>),
25 /// A vacant entry.
26 Vacant(VacantEntry<'a, K, V, S, A, P>),
27}
28
29/// A view into an occupied entry in an `AugmentedRBTree`.
30///
31/// The entry holds a raw pointer to the existing node, so no key clone is needed.
32#[derive(Debug)]
33pub struct OccupiedEntry<'a, K, V, S, A, P>
34where
35 K: Ord,
36 P: TreePolicy<K = K, V = V, S = S>,
37 A: Allocator,
38{
39 tree: &'a mut AugmentedRBTreeInt<K, V, S, A, P>,
40 // Raw pointer to the existing node — valid for 'a because the tree is mutably borrowed.
41 node: NonNull<Node<K, V, S>>,
42}
43
44/// A view into a vacant entry in an `AugmentedRBTree`.
45#[derive(Debug)]
46pub struct VacantEntry<'a, K, V, S, A, P>
47where
48 P: TreePolicy<K = K, V = V, S = S>,
49 K: Ord,
50 A: Allocator,
51{
52 tree: &'a mut AugmentedRBTreeInt<K, V, S, A, P>,
53 key: K,
54}
55
56impl<'a, K, V, S, A: Allocator, P: TreePolicy<K = K, V = V, S = S>> Entry<'a, K, V, S, A, P>
57where
58 K: Ord,
59{
60 pub(crate) fn new(tree: &'a mut AugmentedRBTreeInt<K, V, S, A, P>, key: K) -> Self {
61 // Safety: the tree is mutably borrowed for 'a, so the node pointer is stable.
62 if let Some(node) = tree.layout.find_node(&key) {
63 // Key exists — store the node pointer, discard the key (no clone needed)
64 drop(key);
65 Entry::Occupied(OccupiedEntry {
66 tree,
67 node: node.ptr,
68 })
69 } else {
70 Entry::Vacant(VacantEntry { tree, key })
71 }
72 }
73
74 /// Ensures a value is in the entry by inserting the default if empty, and returns
75 /// a mutable reference to the value in the entry.
76 ///
77 /// # Examples
78 ///
79 /// ```
80 /// # use augmented_rbtree::{AugmentedRBTree, augmentations::SubtreeSize};
81 /// let mut tree = AugmentedRBTree::<&str, u32, SubtreeSize>::new();
82 /// tree.entry("hello").or_insert(3);
83 /// assert_eq!(tree.get(&"hello"), Some(&3));
84 ///
85 /// *tree.entry("hello").or_insert(10) += 2;
86 /// assert_eq!(tree.get(&"hello"), Some(&5));
87 /// ```
88 pub fn or_insert(self, default: V) -> &'a mut V {
89 match self {
90 Entry::Occupied(e) => e.into_mut(),
91 Entry::Vacant(e) => e.insert(default),
92 }
93 }
94
95 /// Ensures a value is in the entry by inserting the result of the default function if empty,
96 /// and returns a mutable reference to the value in the entry.
97 ///
98 /// # Examples
99 ///
100 /// ```
101 /// # use augmented_rbtree::{AugmentedRBTree, augmentations::SubtreeSize};
102 /// let mut tree = AugmentedRBTree::<&str, Vec<i32>, SubtreeSize>::new();
103 /// tree.entry("hello").or_insert_with(Vec::new).push(1);
104 /// assert_eq!(tree.get(&"hello"), Some(&vec![1]));
105 /// ```
106 pub fn or_insert_with(self, default: impl FnOnce() -> V) -> &'a mut V {
107 match self {
108 Entry::Occupied(e) => e.into_mut(),
109 Entry::Vacant(e) => e.insert(default()),
110 }
111 }
112
113 /// Ensures a value is in the entry by inserting the default value if empty, and returns a
114 /// mutable reference to the value in the entry.
115 ///
116 /// # Examples
117 ///
118 /// ```
119 /// # use augmented_rbtree::{AugmentedRBTree, augmentations::SubtreeSize};
120 /// let mut tree = AugmentedRBTree::<&str, u32, SubtreeSize>::new();
121 /// tree.entry("hello").or_default();
122 /// assert_eq!(tree.get(&"hello"), Some(&0));
123 /// ```
124 pub fn or_default(self) -> &'a mut V
125 where
126 V: Default,
127 {
128 self.or_insert_with(V::default)
129 }
130
131 /// Returns a reference to this entry's key.
132 ///
133 /// # Examples
134 ///
135 /// ```
136 /// # use augmented_rbtree::{AugmentedRBTree, Unit};
137 /// let mut tree = AugmentedRBTree::<&str, u32, Unit>::new();
138 /// assert_eq!(tree.entry("hello").key(), &"hello");
139 /// ```
140 pub fn key(&self) -> &K {
141 match self {
142 Entry::Occupied(e) => e.key(),
143 Entry::Vacant(e) => &e.key,
144 }
145 }
146
147 /// Provides in-place mutable access to an occupied entry before any potential inserts.
148 ///
149 /// # Examples
150 ///
151 /// ```
152 /// # use augmented_rbtree::{AugmentedRBTree, Unit};
153 /// let mut tree = AugmentedRBTree::<&str, u32, Unit>::new();
154 /// tree.entry("hello").and_modify(|e| *e += 1).or_insert(42);
155 /// assert_eq!(tree.get(&"hello"), Some(&42));
156 ///
157 /// tree.entry("hello").and_modify(|e| *e += 1).or_insert(42);
158 /// assert_eq!(tree.get(&"hello"), Some(&43));
159 /// ```
160 #[must_use]
161 pub fn and_modify(self, f: impl FnOnce(&mut V)) -> Self {
162 match self {
163 Entry::Occupied(mut e) => {
164 f(e.get_mut());
165 Entry::Occupied(e)
166 }
167 Entry::Vacant(e) => Entry::Vacant(e),
168 }
169 }
170}
171
172impl<'a, K, V, S, A, P> OccupiedEntry<'a, K, V, S, A, P>
173where
174 K: Ord,
175 P: TreePolicy<K = K, V = V, S = S>,
176 A: Allocator,
177{
178 /// Returns a reference to the key of this entry.
179 #[must_use]
180 pub fn key(&self) -> &K {
181 // Safety: node is valid for 'a (tree is mutably borrowed).
182 unsafe { &(*self.node.as_ptr()).key }
183 }
184
185 /// Gets a reference to the value in the entry.
186 #[must_use]
187 pub fn get(&self) -> &V {
188 unsafe { &(*self.node.as_ptr()).value }
189 }
190
191 /// Gets a mutable reference to the value in the entry.
192 pub fn get_mut(&mut self) -> &mut V {
193 unsafe { &mut (*self.node.as_ptr()).value }
194 }
195
196 /// Converts the `OccupiedEntry` into a mutable reference to the value in the entry with a
197 /// lifetime bound to the tree itself.
198 #[must_use]
199 pub fn into_mut(self) -> &'a mut V {
200 unsafe { &mut (*self.node.as_ptr()).value }
201 }
202
203 /// Sets the value of the entry, and returns the old value.
204 pub fn insert(&mut self, value: V) -> V {
205 unsafe { core::mem::replace(&mut (*self.node.as_ptr()).value, value) }
206 }
207
208 /// Takes the value out of the entry, and returns it.
209 ///
210 /// # Panics
211 ///
212 /// Panics if the tree is corrupted (entry exists but key cannot be found on removal).
213 #[must_use]
214 pub fn remove(self) -> V
215 where
216 K: Clone,
217 {
218 let key = self.key().clone();
219 self.tree
220 .remove(&key)
221 .expect("occupied entry must have a value")
222 }
223}
224
225impl<'a, K, V, S, A, P> VacantEntry<'a, K, V, S, A, P>
226where
227 K: Ord,
228 P: TreePolicy<K = K, V = V, S = S>,
229 A: Allocator,
230{
231 /// Gets a reference to the key that would be used when inserting a value through the `VacantEntry`.
232 pub fn key(&self) -> &K {
233 &self.key
234 }
235
236 /// Take ownership of the key.
237 pub fn into_key(self) -> K {
238 self.key
239 }
240
241 /// Try to insert a value with the `VacantEntry`'s key, and returns a mutable reference
242 /// to it.
243 ///
244 /// # Errors
245 ///
246 /// Returns an `OutOfMemoryError` if the allocation fails.
247 pub fn try_insert(self, value: V) -> Result<&'a mut V, OutOfMemoryError>
248 where
249 P: TreePolicy<K = K, V = V, S = S>,
250 {
251 // We own the key — insert it, then retrieve a pointer to the newly inserted node.
252 let node = self.tree.layout.try_insert_node_get_ref(self.key, value)?;
253 Ok(unsafe { &mut (*node.ptr.as_ptr()).value })
254 }
255
256 /// Set the value of the entry with the `VacantEntry`'s key, and returns a mutable reference
257 /// to it.
258 pub fn insert(self, value: V) -> &'a mut V {
259 self.try_insert(value)
260 .unwrap_or_else(|_| handle_alloc_error(Layout::new::<Node<K, V, S>>()))
261 }
262}