kevy_map/raw_entry.rs
1//! Single-probe entry API à la hashbrown's `RawEntryMut`.
2//!
3//! Motivation: the existing read/insert APIs (`get`, `get_mut`, `insert`,
4//! `remove`) each cost one full probe. Common Store patterns —
5//! "look up, check expiry, conditionally remove, otherwise return the
6//! borrow" — currently do **two** probes because the borrow returned by
7//! `get` cannot survive a subsequent `remove` (mutable-borrow conflict
8//! with the immutable borrow on the value). The raw-entry API folds the
9//! read and the conditional remove into a single probe by consuming the
10//! `RawOccupiedEntryMut` on `remove(self)`, which releases the borrow at
11//! the call site.
12//!
13//! Design notes:
14//!
15//! * `raw_entry_mut` takes `&mut self` (any subsequent insert/remove
16//! needs exclusive access; matching `Borrow<Q>` lookup matches the
17//! shape of `get_mut`).
18//! * The `RawOccupiedEntryMut` stores a borrowed `&'a mut KevyMap<K, V>`
19//! plus the slot index already located by the probe. `get` /
20//! `get_mut` / `into_mut` reuse that slot; `remove(self)` consumes the
21//! entry so the map borrow is freed and a `set_meta(DELETED)` write
22//! can proceed.
23//! * The `RawVacantEntryMut` stores `&'a mut KevyMap<K, V>` and **does
24//! not** cache the probe's `insert_at`: any subsequent mutation
25//! (notably `maybe_grow` inside `insert`) can invalidate that slot.
26//! `insert(self, k, v)` therefore re-runs `insert` from scratch
27//! (one extra probe in the absent-key insert path; the API-additive
28//! commit is purely about unblocking the *read or remove* fast path —
29//! downstream cumulative attacks may later push the cached probe
30//! through, once the grow-invalidation issue is solved with an
31//! explicit `reserve` API).
32//!
33//! No `unsafe` is added by this file: every memory touch is delegated
34//! to the existing `pub(crate)` helpers in `map.rs` / `map_keyed.rs`.
35
36use std::borrow::Borrow;
37use std::ptr;
38
39use kevy_hash::KevyHash;
40
41use crate::map::{DELETED, KevyMap, ProbeOutcome};
42
43/// Result of [`KevyMap::raw_entry_mut`].
44///
45/// Mirrors `hashbrown::hash_map::RawEntryMut`: an `Occupied` arm grants
46/// read / mutate / consume access to the existing entry; a `Vacant` arm
47/// can be filled with `insert(k, v)`. The defining property — and the
48/// reason this API exists distinct from `get_mut` — is that
49/// [`RawOccupiedEntryMut::remove`] consumes `self`, which releases the
50/// outstanding borrow on the map and lets the caller perform the
51/// deletion within the same borrow scope.
52pub enum RawEntryMut<'a, K, V> {
53 /// The key was present; gives read / mutate / consume access.
54 Occupied(RawOccupiedEntryMut<'a, K, V>),
55 /// The key was absent; `insert(k, v)` writes a new entry.
56 Vacant(RawVacantEntryMut<'a, K, V>),
57}
58
59/// Handle to an existing entry, returned by [`RawEntryMut::Occupied`].
60pub struct RawOccupiedEntryMut<'a, K, V> {
61 map: &'a mut KevyMap<K, V>,
62 /// Slot index inside `map.slots_ptr` for the located entry.
63 slot: usize,
64}
65
66/// Handle to an absent entry, returned by [`RawEntryMut::Vacant`].
67pub struct RawVacantEntryMut<'a, K, V> {
68 map: &'a mut KevyMap<K, V>,
69}
70
71impl<K, V> KevyMap<K, V> {
72 /// Look up `key`; return an [`RawEntryMut`] giving single-probe
73 /// access to the located (or vacant) slot.
74 ///
75 /// One full probe. The returned handle borrows `self` mutably; the
76 /// borrow is released only when the handle is dropped (or consumed
77 /// via [`RawOccupiedEntryMut::remove`] / [`RawOccupiedEntryMut::into_mut`]).
78 pub fn raw_entry_mut<Q>(&mut self, key: &Q) -> RawEntryMut<'_, K, V>
79 where
80 K: Borrow<Q> + KevyHash + Eq,
81 Q: KevyHash + Eq + ?Sized,
82 {
83 match self.probe_by_borrow(key) {
84 ProbeOutcome::Found(slot) => RawEntryMut::Occupied(RawOccupiedEntryMut {
85 map: self,
86 slot,
87 }),
88 ProbeOutcome::NotFound { .. } => {
89 RawEntryMut::Vacant(RawVacantEntryMut { map: self })
90 }
91 }
92 }
93}
94
95impl<'a, K, V> RawOccupiedEntryMut<'a, K, V> {
96 /// Shared access to the stored value.
97 #[inline]
98 pub fn get(&self) -> &V {
99 // SAFETY: `slot` came from `probe_by_borrow::Found`, which only
100 // returns indices into full slots.
101 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_ref() };
102 &kv.1
103 }
104
105 /// Mutable access to the stored value (borrow tied to `&mut self`).
106 #[inline]
107 pub fn get_mut(&mut self) -> &mut V {
108 // SAFETY: see [`get`].
109 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_mut() };
110 &mut kv.1
111 }
112
113 /// Mutable access to the stored value with the outer map's lifetime,
114 /// consuming the handle. The map borrow returned outlives `self`,
115 /// which is exactly the shape `get`+`get_mut` cannot provide.
116 #[inline]
117 pub fn into_mut(self) -> &'a mut V {
118 // SAFETY: see [`get`]. The returned borrow's lifetime `'a` is
119 // tied to the borrow we were constructed with, which is the
120 // caller's `&mut KevyMap<K, V>` borrow.
121 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_mut() };
122 &mut kv.1
123 }
124
125 /// Shared access to the stored key.
126 #[inline]
127 pub fn key(&self) -> &K {
128 // SAFETY: see [`get`].
129 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_ref() };
130 &kv.0
131 }
132
133 /// Remove the entry; returns the previous value. Consumes `self`,
134 /// releasing the map borrow so the caller can immediately re-probe
135 /// or mutate something else.
136 ///
137 /// This is the load-bearing method — it is what makes the API
138 /// strictly more expressive than `get_mut`+`remove`.
139 pub fn remove(self) -> V {
140 self.map.set_meta(self.slot, DELETED);
141 self.map.occupied -= 1;
142 self.map.deleted += 1;
143 // SAFETY: slot was full; we just marked it DELETED so it won't
144 // be re-read as occupied. `ptr::read` moves the (K, V) out;
145 // dropping `k` here is correct because we don't return it.
146 let (_k, v) = unsafe {
147 ptr::read(self.map.slots_ptr.as_ptr().add(self.slot) as *const (K, V))
148 };
149 v
150 }
151}
152
153impl<'a, K, V> RawVacantEntryMut<'a, K, V>
154where
155 K: KevyHash + Eq,
156{
157 /// Insert `(key, value)` and return a mutable borrow of the freshly
158 /// inserted value. The borrow is tied to the original map borrow.
159 ///
160 /// Note: this performs a second probe after a possible grow, because
161 /// the slot located by the first probe (the one that produced
162 /// `Vacant`) can be invalidated by `maybe_grow`. The added probe is
163 /// acceptable for the live_entry pattern (the *read* fast path is
164 /// the one we were chasing; the absent-key insert path is the cold
165 /// side, and `live_entry` itself never takes it).
166 pub fn insert(self, key: K, value: V) -> &'a mut V {
167 // Grow if needed (matches `KevyMap::insert`'s preamble).
168 self.map.maybe_grow();
169 // Re-probe by reference, write into the slot, bump occupancy.
170 let hash = key.kevy_hash();
171 let outcome = self.map.probe_by_borrow(&key);
172 let slot = match outcome {
173 ProbeOutcome::NotFound { insert_at, via_tombstone } => {
174 self.map.set_meta(insert_at, crate::map::h2(hash));
175 // SAFETY: insert_at < cap ⇒ slot pointer in-bounds; we
176 // write (K, V) into a previously uninitialised slot.
177 unsafe {
178 (*self.map.slots_ptr.as_ptr().add(insert_at)).write((key, value));
179 }
180 self.map.occupied += 1;
181 if via_tombstone {
182 self.map.deleted -= 1;
183 }
184 insert_at
185 }
186 ProbeOutcome::Found(_) => {
187 // Cannot happen: we held the only mutable borrow on the
188 // map between `raw_entry_mut` and here, and the first
189 // probe said Vacant. Treat as logic bug rather than
190 // overwriting (overwriting would silently change the
191 // documented contract).
192 unreachable!("raw vacant insert observed an existing key");
193 }
194 };
195 // SAFETY: slot is the one we just initialised.
196 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(slot)).assume_init_mut() };
197 &mut kv.1
198 }
199}