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 core::borrow::Borrow;
37use core::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) => {
85 RawEntryMut::Occupied(RawOccupiedEntryMut { map: self, slot })
86 }
87 ProbeOutcome::NotFound { .. } => RawEntryMut::Vacant(RawVacantEntryMut { map: self }),
88 }
89 }
90}
91
92impl<'a, K, V> RawOccupiedEntryMut<'a, K, V> {
93 /// Shared access to the stored value.
94 #[inline]
95 pub fn get(&self) -> &V {
96 // SAFETY: `slot` came from `probe_by_borrow::Found`, which only
97 // returns indices into full slots.
98 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_ref() };
99 &kv.1
100 }
101
102 /// Mutable access to the stored value (borrow tied to `&mut self`).
103 #[inline]
104 pub fn get_mut(&mut self) -> &mut V {
105 // SAFETY: see [`get`].
106 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_mut() };
107 &mut kv.1
108 }
109
110 /// Mutable access to the stored value with the outer map's lifetime,
111 /// consuming the handle. The map borrow returned outlives `self`,
112 /// which is exactly the shape `get`+`get_mut` cannot provide.
113 #[inline]
114 pub fn into_mut(self) -> &'a mut V {
115 // SAFETY: see [`get`]. The returned borrow's lifetime `'a` is
116 // tied to the borrow we were constructed with, which is the
117 // caller's `&mut KevyMap<K, V>` borrow.
118 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_mut() };
119 &mut kv.1
120 }
121
122 /// Shared access to the stored key.
123 #[inline]
124 pub fn key(&self) -> &K {
125 // SAFETY: see [`get`].
126 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(self.slot)).assume_init_ref() };
127 &kv.0
128 }
129
130 /// Remove the entry; returns the previous value. Consumes `self`,
131 /// releasing the map borrow so the caller can immediately re-probe
132 /// or mutate something else.
133 ///
134 /// This is the load-bearing method — it is what makes the API
135 /// strictly more expressive than `get_mut`+`remove`.
136 pub fn remove(self) -> V {
137 self.map.set_meta(self.slot, DELETED);
138 self.map.occupied -= 1;
139 self.map.deleted += 1;
140 // SAFETY: slot was full; we just marked it DELETED so it won't
141 // be re-read as occupied. `ptr::read` moves the (K, V) out;
142 // dropping `k` here is correct because we don't return it.
143 let (_k, v) =
144 unsafe { ptr::read(self.map.slots_ptr.as_ptr().add(self.slot) as *const (K, V)) };
145 v
146 }
147}
148
149impl<'a, K, V> RawVacantEntryMut<'a, K, V>
150where
151 K: KevyHash + Eq,
152{
153 /// Insert `(key, value)` and return a mutable borrow of the freshly
154 /// inserted value. The borrow is tied to the original map borrow.
155 ///
156 /// Note: this performs a second probe after a possible grow, because
157 /// the slot located by the first probe (the one that produced
158 /// `Vacant`) can be invalidated by `maybe_grow`. The added probe is
159 /// acceptable for the live_entry pattern (the *read* fast path is
160 /// the one we were chasing; the absent-key insert path is the cold
161 /// side, and `live_entry` itself never takes it).
162 pub fn insert(self, key: K, value: V) -> &'a mut V {
163 // Grow if needed (matches `KevyMap::insert`'s preamble).
164 self.map.maybe_grow();
165 // Re-probe by reference, write into the slot, bump occupancy.
166 let hash = key.kevy_hash();
167 let outcome = self.map.probe_by_borrow(&key);
168 let slot = match outcome {
169 ProbeOutcome::NotFound { insert_at, via_tombstone } => {
170 self.map.set_meta(insert_at, crate::map::h2(hash));
171 // SAFETY: insert_at < cap ⇒ slot pointer in-bounds; we
172 // write (K, V) into a previously uninitialised slot.
173 unsafe {
174 (*self.map.slots_ptr.as_ptr().add(insert_at)).write((key, value));
175 }
176 self.map.occupied += 1;
177 if via_tombstone {
178 self.map.deleted -= 1;
179 }
180 insert_at
181 }
182 ProbeOutcome::Found(_) => {
183 // Cannot happen: we held the only mutable borrow on the
184 // map between `raw_entry_mut` and here, and the first
185 // probe said Vacant. Treat as logic bug rather than
186 // overwriting (overwriting would silently change the
187 // documented contract).
188 unreachable!("raw vacant insert observed an existing key");
189 }
190 };
191 // SAFETY: slot is the one we just initialised.
192 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(slot)).assume_init_mut() };
193 &mut kv.1
194 }
195}