rc_hashmap/
rc_hash_map.rs

1use crate::tokens::{Count, RcCount, Token};
2// Keepalive handled via direct Rc strong-count inc/dec per entry.
3use crate::counted_hash_map::{CountedHandle, CountedHashMap, PutResult};
4use crate::handle_hash_map::InsertError;
5use crate::hash::DefaultHashBuilder;
6use core::cell::UnsafeCell;
7use core::hash::{Hash, Hasher};
8use core::marker::PhantomData;
9use core::mem::ManuallyDrop;
10use std::ptr::NonNull;
11use std::rc::Rc;
12
13// Stored value wrapper that holds a keepalive token from `Inner`'s RcCount
14// to keep the allocation alive. The token is returned when the last Ref
15// for this entry is dropped and the entry is removed.
16struct RcVal<K, V, S> {
17    value: V,
18    keepalive_token: Token<'static, RcCount<Inner<K, V, S>>>,
19}
20
21struct Inner<K, V, S> {
22    map: UnsafeCell<CountedHashMap<K, RcVal<K, V, S>, S>>, // interior mutability via UnsafeCell
23    keepalive: RcCount<Inner<K, V, S>>,
24}
25
26pub struct RcHashMap<K, V, S = DefaultHashBuilder> {
27    inner: Rc<Inner<K, V, S>>,
28}
29
30impl<K, V> RcHashMap<K, V>
31where
32    K: Eq + core::hash::Hash + 'static,
33    V: 'static,
34{
35    pub fn new() -> Self {
36        Self {
37            inner: Rc::new_cyclic(|weak| Inner {
38                map: UnsafeCell::new(CountedHashMap::new()),
39                keepalive: RcCount::from_weak(weak),
40            }),
41        }
42    }
43}
44
45impl<K, V> Default for RcHashMap<K, V>
46where
47    K: Eq + core::hash::Hash + 'static,
48    V: 'static,
49{
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55impl<K, V, S> RcHashMap<K, V, S>
56where
57    K: Eq + core::hash::Hash + 'static,
58    V: 'static,
59    S: core::hash::BuildHasher + Clone + Default + 'static,
60{
61    // Internal helpers to access the inner map via UnsafeCell in one place.
62    fn map(&self) -> &CountedHashMap<K, RcVal<K, V, S>, S> {
63        unsafe { &*self.inner.map.get() }
64    }
65    fn map_mut(&mut self) -> &mut CountedHashMap<K, RcVal<K, V, S>, S> {
66        unsafe { &mut *self.inner.map.get() }
67    }
68    #[allow(clippy::type_complexity)]
69    fn map_and_rccount_mut(
70        &mut self,
71    ) -> (
72        &mut CountedHashMap<K, RcVal<K, V, S>, S>,
73        &RcCount<Inner<K, V, S>>,
74    ) {
75        let m = unsafe { &mut *self.inner.map.get() };
76        let rc = &self.inner.keepalive;
77        (m, rc)
78    }
79    pub fn with_hasher(hasher: S) -> Self {
80        Self {
81            inner: Rc::new_cyclic(|weak| Inner {
82                map: UnsafeCell::new(CountedHashMap::with_hasher(hasher)),
83                keepalive: RcCount::from_weak(weak),
84            }),
85        }
86    }
87
88    pub fn len(&self) -> usize {
89        self.map().len()
90    }
91    pub fn is_empty(&self) -> bool {
92        self.map().is_empty()
93    }
94
95    pub fn contains_key<Q>(&self, q: &Q) -> bool
96    where
97        K: core::borrow::Borrow<Q>,
98        Q: ?Sized + core::hash::Hash + Eq,
99    {
100        self.map().contains_key(q)
101    }
102
103    pub fn insert(&mut self, key: K, value: V) -> Result<Ref<K, V, S>, InsertError> {
104        let (map, keepalive) = self.map_and_rccount_mut();
105        let res = map.insert_with(key, || RcVal {
106            value,
107            keepalive_token: keepalive.get(),
108        });
109        match res {
110            Ok(ch) => Ok(Ref::new(NonNull::from(self.inner.as_ref()), ch)),
111            Err(e) => Err(e),
112        }
113    }
114
115    pub fn find<Q>(&self, q: &Q) -> Option<Ref<K, V, S>>
116    where
117        K: core::borrow::Borrow<Q>,
118        Q: ?Sized + core::hash::Hash + Eq,
119    {
120        self.map()
121            .find(q)
122            .map(|ch| Ref::new(NonNull::from(self.inner.as_ref()), ch))
123    }
124
125    pub fn iter(&self) -> Iter<'_, K, V, S> {
126        let owner_ptr = NonNull::from(self.inner.as_ref());
127        let inner = self.map().iter_raw();
128        Iter { owner_ptr, inner }
129    }
130
131    pub fn iter_mut(&mut self) -> IterMut<'_, K, V, S> {
132        let owner_ptr = NonNull::from(self.inner.as_ref());
133        let inner = self.map_mut().iter_mut_raw();
134        IterMut { owner_ptr, inner }
135    }
136}
137
138/// A reference to an entry inside RcHashMap. Clone increments per-entry count;
139/// dropping decrements and removes the entry when it reaches zero.
140pub struct Ref<K, V, S = DefaultHashBuilder>
141where
142    K: Eq + core::hash::Hash + 'static,
143    V: 'static,
144    S: core::hash::BuildHasher + Clone + Default + 'static,
145{
146    owner_ptr: NonNull<Inner<K, V, S>>,
147    handle: ManuallyDrop<CountedHandle<'static>>,
148    _nosend: PhantomData<*mut ()>,
149}
150
151/// Owner-mismatch error for Ref accessors.
152#[derive(Copy, Clone, Debug, Eq, PartialEq)]
153pub struct WrongMap;
154
155impl<K, V, S> Ref<K, V, S>
156where
157    K: Eq + core::hash::Hash,
158    S: core::hash::BuildHasher + Clone + Default,
159{
160    fn new(owner_ptr: NonNull<Inner<K, V, S>>, handle: CountedHandle<'static>) -> Self {
161        Self {
162            owner_ptr,
163            handle: ManuallyDrop::new(handle),
164            _nosend: PhantomData,
165        }
166    }
167
168    #[inline]
169    fn check_owner<'a>(&'a self, map: &'a RcHashMap<K, V, S>) -> Result<(), WrongMap> {
170        // Safety: owner_ptr is created from Rc::as_ref; compare raw pointers for identity.
171        let ptr = NonNull::from(map.inner.as_ref());
172        if ptr == self.owner_ptr {
173            Ok(())
174        } else {
175            Err(WrongMap)
176        }
177    }
178
179    /// Borrow the entry's key, validating owner identity.
180    pub fn key<'a>(&'a self, map: &'a RcHashMap<K, V, S>) -> Result<&'a K, WrongMap> {
181        self.check_owner(map)?;
182        self.handle.key_ref(map.map()).ok_or(WrongMap)
183    }
184
185    /// Borrow the entry's value, validating owner identity.
186    pub fn value<'a>(&'a self, map: &'a RcHashMap<K, V, S>) -> Result<&'a V, WrongMap> {
187        self.check_owner(map)?;
188        self.handle
189            .value_ref(map.map())
190            .map(|rcv| &rcv.value)
191            .ok_or(WrongMap)
192    }
193
194    /// Mutably borrow the entry's value, validating owner identity.
195    pub fn value_mut<'a>(&'a self, map: &'a mut RcHashMap<K, V, S>) -> Result<&'a mut V, WrongMap> {
196        if NonNull::from(map.inner.as_ref()) != self.owner_ptr {
197            return Err(WrongMap);
198        }
199        // SAFETY: owner validated and we have &mut map, so exclusive access for 'a
200        self.check_owner(map)?; // ensure owner match
201        self.handle
202            .value_mut(map.map_mut())
203            .map(|rcv| &mut rcv.value)
204            .ok_or(WrongMap)
205    }
206}
207
208impl<K, V, S> Clone for Ref<K, V, S>
209where
210    K: Eq + core::hash::Hash + 'static,
211    V: 'static,
212    S: core::hash::BuildHasher + Clone + Default + 'static,
213{
214    fn clone(&self) -> Self {
215        // Increment per-entry count via counted handle API.
216        let inner = unsafe { self.owner_ptr.as_ref() };
217        let handle = unsafe { &*inner.map.get() }.get(&self.handle);
218        Ref::new(self.owner_ptr, handle)
219    }
220}
221
222impl<K, V, S> Drop for Ref<K, V, S>
223where
224    K: Eq + core::hash::Hash + 'static,
225    V: 'static,
226    S: core::hash::BuildHasher + Clone + Default + 'static,
227{
228    fn drop(&mut self) {
229        let inner = unsafe { &mut *(self.owner_ptr.as_ptr()) };
230        // Move out the handle without running its destructor.
231        let ch = unsafe { ManuallyDrop::take(&mut self.handle) };
232        let res = unsafe { &mut *inner.map.get() }.put(ch);
233        match res {
234            PutResult::Live => {}
235            PutResult::Removed { key, value } => {
236                // Drop user data first while keepalive still holds Inner alive via strong count
237                let RcVal {
238                    value: user_value,
239                    keepalive_token,
240                } = value;
241                drop(key);
242                drop(user_value);
243                // Return the keepalive token to decrement the strong count.
244                inner.keepalive.put(keepalive_token);
245            }
246        }
247    }
248}
249
250impl<K, V, S> PartialEq for Ref<K, V, S>
251where
252    K: Eq + core::hash::Hash,
253    S: core::hash::BuildHasher + Clone + Default,
254{
255    fn eq(&self, other: &Self) -> bool {
256        self.owner_ptr == other.owner_ptr && self.handle.handle == other.handle.handle
257    }
258}
259
260impl<K, V, S> Eq for Ref<K, V, S>
261where
262    K: Eq + core::hash::Hash,
263    S: core::hash::BuildHasher + Clone + Default,
264{
265}
266
267impl<K, V, S> Hash for Ref<K, V, S>
268where
269    K: Eq + core::hash::Hash,
270    S: core::hash::BuildHasher + Clone + Default,
271{
272    fn hash<H: Hasher>(&self, state: &mut H) {
273        (self.owner_ptr.as_ptr() as usize).hash(state);
274        self.handle.handle.hash(state);
275    }
276}
277/// Placeholder for future mutable iterator item (see design docs).
278pub struct ItemMut<'a, K, V, S = DefaultHashBuilder>
279where
280    K: Eq + core::hash::Hash + 'static,
281    V: 'static,
282    S: core::hash::BuildHasher + Clone + Default + 'static,
283{
284    r: Ref<K, V, S>,
285    k: &'a K,
286    v: &'a mut V,
287}
288impl<'a, K, V, S> ItemMut<'a, K, V, S>
289where
290    K: Eq + core::hash::Hash + 'static,
291    V: 'static,
292    S: core::hash::BuildHasher + Clone + Default + 'static,
293{
294    pub fn r#ref(&self) -> &Ref<K, V, S> {
295        &self.r
296    }
297    pub fn key(&self) -> &K {
298        self.k
299    }
300    pub fn value_mut(&mut self) -> &mut V {
301        self.v
302    }
303}
304
305/// Immutable iterator for RcHashMap yielding `Ref`.
306pub struct Iter<'a, K, V, S = DefaultHashBuilder>
307where
308    K: Eq + core::hash::Hash + 'static,
309    V: 'static,
310    S: core::hash::BuildHasher + Clone + Default + 'static,
311{
312    owner_ptr: NonNull<Inner<K, V, S>>,
313    inner: crate::counted_hash_map::Iter<'a, K, RcVal<K, V, S>, S>,
314}
315
316impl<'a, K, V, S> Iterator for Iter<'a, K, V, S>
317where
318    K: Eq + core::hash::Hash + 'static,
319    V: 'static,
320    S: core::hash::BuildHasher + Clone + Default + 'static,
321{
322    type Item = Ref<K, V, S>;
323    fn next(&mut self) -> Option<Self::Item> {
324        self.inner
325            .next()
326            .map(|(ch, _k, _rv)| Ref::new(self.owner_ptr, ch))
327    }
328}
329
330/// Mutable iterator for RcHashMap yielding ItemMut.
331pub struct IterMut<'a, K, V, S = DefaultHashBuilder>
332where
333    K: Eq + core::hash::Hash + 'static,
334    V: 'static,
335    S: core::hash::BuildHasher + Clone + Default + 'static,
336{
337    owner_ptr: NonNull<Inner<K, V, S>>,
338    inner: crate::counted_hash_map::IterMut<'a, K, RcVal<K, V, S>, S>,
339}
340
341impl<'a, K, V, S> Iterator for IterMut<'a, K, V, S>
342where
343    K: Eq + core::hash::Hash + 'static,
344    V: 'static,
345    S: core::hash::BuildHasher + Clone + Default + 'static,
346{
347    type Item = ItemMut<'a, K, V, S>;
348    fn next(&mut self) -> Option<Self::Item> {
349        self.inner.next().map(|(ch, k, rv)| {
350            let r = Ref::new(self.owner_ptr, ch);
351            ItemMut {
352                r,
353                k,
354                v: &mut rv.value,
355            }
356        })
357    }
358}
rc_hashmap/rc_hash_map.rs

rc_hashmap/
rc_hash_map.rs
