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.
52#[derive(Debug)]
53pub enum RawEntryMut<'a, K, V> {
54 /// The key was present; gives read / mutate / consume access.
55 Occupied(RawOccupiedEntryMut<'a, K, V>),
56 /// The key was absent; `insert(k, v)` writes a new entry.
57 Vacant(RawVacantEntryMut<'a, K, V>),
58}
59
60/// Handle to an existing entry, returned by [`RawEntryMut::Occupied`].
61#[derive(Debug)]
62pub struct RawOccupiedEntryMut<'a, K, V> {
63 map: &'a mut KevyMap<K, V>,
64 /// Slot index inside `map.slots_ptr` for the located entry.
65 slot: usize,
66}
67
68/// Handle to an absent entry, returned by [`RawEntryMut::Vacant`].
69#[derive(Debug)]
70pub struct RawVacantEntryMut<'a, K, V> {
71 map: &'a mut KevyMap<K, V>,
72}
73
74impl<K, V> KevyMap<K, V> {
75 /// Look up `key`; return an [`RawEntryMut`] giving single-probe
76 /// access to the located (or vacant) slot.
77 ///
78 /// One full probe. The returned handle borrows `self` mutably; the
79 /// borrow is released only when the handle is dropped (or consumed
80 /// via [`RawOccupiedEntryMut::remove`] / [`RawOccupiedEntryMut::into_mut`]).
81 pub fn raw_entry_mut<Q>(&mut self, key: &Q) -> RawEntryMut<'_, K, V>
82 where
83 K: Borrow<Q> + KevyHash + Eq,
84 Q: KevyHash + Eq + ?Sized,
85 {
86 match self.probe_by_borrow(key) {
87 ProbeOutcome::Found(slot) => {
88 RawEntryMut::Occupied(RawOccupiedEntryMut { map: self, slot })
89 }
90 ProbeOutcome::NotFound { .. } => RawEntryMut::Vacant(RawVacantEntryMut { map: self }),
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) =
147 unsafe { ptr::read(self.map.slots_ptr.as_ptr().add(self.slot) as *const (K, V)) };
148 v
149 }
150}
151
152impl<'a, K, V> RawVacantEntryMut<'a, K, V>
153where
154 K: KevyHash + Eq,
155{
156 /// Insert `(key, value)` and return a mutable borrow of the freshly
157 /// inserted value. The borrow is tied to the original map borrow.
158 ///
159 /// Note: this performs a second probe after a possible grow, because
160 /// the slot located by the first probe (the one that produced
161 /// `Vacant`) can be invalidated by `maybe_grow`. The added probe is
162 /// acceptable for the live_entry pattern (the *read* fast path is
163 /// the one we were chasing; the absent-key insert path is the cold
164 /// side, and `live_entry` itself never takes it).
165 pub fn insert(self, key: K, value: V) -> &'a mut V {
166 // Grow if needed (matches `KevyMap::insert`'s preamble).
167 self.map.maybe_grow();
168 // Re-probe by reference, write into the slot, bump occupancy.
169 let hash = key.kevy_hash();
170 let outcome = self.map.probe_by_borrow(&key);
171 let slot = match outcome {
172 ProbeOutcome::NotFound { insert_at, via_tombstone } => {
173 self.map.set_meta(insert_at, crate::map::h2(hash));
174 // SAFETY: insert_at < cap ⇒ slot pointer in-bounds; we
175 // write (K, V) into a previously uninitialised slot.
176 unsafe {
177 (*self.map.slots_ptr.as_ptr().add(insert_at)).write((key, value));
178 }
179 self.map.occupied += 1;
180 if via_tombstone {
181 self.map.deleted -= 1;
182 }
183 insert_at
184 }
185 ProbeOutcome::Found(_) => {
186 // Cannot happen: we held the only mutable borrow on the
187 // map between `raw_entry_mut` and here, and the first
188 // probe said Vacant. Treat as logic bug rather than
189 // overwriting (overwriting would silently change the
190 // documented contract).
191 unreachable!("raw vacant insert observed an existing key");
192 }
193 };
194 // SAFETY: slot is the one we just initialised.
195 let kv = unsafe { (*self.map.slots_ptr.as_ptr().add(slot)).assume_init_mut() };
196 &mut kv.1
197 }
198}