twizzler/collections/hachage/
map.rs

1use std::{
2    hash::{BuildHasher, Hash},
3    marker::PhantomData,
4};
5
6pub use equivalent::Equivalent;
7
8use crate::{
9    alloc::Allocator,
10    collections::hachage::{raw::*, DefaultHashBuilder},
11    marker::Invariant,
12    object::{Object, ObjectBuilder, TxObject, TypedObject},
13    ptr::{GlobalPtr, Ref, RefMut},
14    Result,
15};
16
17pub(crate) fn make_hasher<Q, V, S>(hash_builder: &S) -> impl Fn(&(Q, V)) -> u64 + '_
18where
19    Q: Hash,
20    S: BuildHasher,
21{
22    move |val| make_hash::<Q, S>(hash_builder, &val.0)
23}
24
25pub(crate) fn make_hash<Q, S>(hash_builder: &S, val: &Q) -> u64
26where
27    Q: Hash + ?Sized,
28    S: BuildHasher,
29{
30    use std::hash::Hasher;
31    let mut state = hash_builder.build_hasher();
32    val.hash(&mut state);
33    state.finish()
34}
35
36pub(crate) fn equivalent_key<Q, K, V>(k: &Q) -> impl Fn(&(K, V)) -> bool + '_
37where
38    Q: Equivalent<K> + ?Sized,
39{
40    move |x| k.equivalent(&x.0)
41}
42
43pub struct PersistentHashMap<
44    K: Invariant,
45    V: Invariant,
46    S = DefaultHashBuilder,
47    A: Allocator = HashTableAlloc,
48> {
49    pub(crate) table: Object<RawTable<(K, V), S, A>>,
50}
51
52impl<K: Invariant, V: Invariant> PersistentHashMap<K, V, DefaultHashBuilder, HashTableAlloc> {
53    pub fn new() -> Result<Self> {
54        let builder = ObjectBuilder::default();
55        Self::with_builder(builder)
56    }
57
58    pub fn new_persist() -> Result<Self> {
59        let builder = ObjectBuilder::default().persist(true);
60        Self::with_builder(builder)
61    }
62
63    pub fn with_builder(
64        builder: ObjectBuilder<RawTable<(K, V), DefaultHashBuilder, HashTableAlloc>>,
65    ) -> Result<Self> {
66        let phm = Self::with_hasher_in(builder, Default::default(), Default::default())?;
67
68        // There's a circular dependency if an RawTable attempts to allocate
69        // before the object so we need to do part of the allocation afterwards
70        // so that an empty table can be made.
71        let mut phm_tx = phm.table.as_tx()?;
72        let mut base = phm_tx.base_mut();
73
74        base.bootstrap(1)?;
75
76        Ok(phm)
77    }
78}
79
80impl<K: Invariant, V: Invariant, S, A: Allocator> From<Object<RawTable<(K, V), S, A>>> for PersistentHashMap<K, V, S, A> {
81    fn from(table: Object<RawTable<(K, V), S, A>>) -> Self {
82        Self { table }
83    }
84}
85
86impl<K: Invariant, V: Invariant, S, A: Allocator> From<GlobalPtr<RawTable<(K, V), S, A>>>
87    for PersistentHashMap<K, V, S, A>
88{
89    fn from(value: GlobalPtr<RawTable<(K, V), S, A>>) -> Self {
90        unsafe {
91            let value = Ref::handle(&value.resolve()).clone();
92            Self {
93                table: Object::from_handle(value).unwrap(),
94            }
95        }
96    }
97}
98
99impl<K: Invariant, V: Invariant, S, A: Allocator> PersistentHashMap<K, V, S, A> {
100    pub fn object(&self) -> &Object<RawTable<(K, V), S, A>> {
101        &self.table
102    }
103
104    pub fn into_object(self) -> Object<RawTable<(K, V), S, A>> {
105        self.table
106    }
107
108    pub fn capacity(&self) -> usize {
109        self.table.base().capacity()
110    }
111
112    #[inline]
113    pub fn allocator(&self) -> &A {
114        self.table.base().allocator()
115    }
116
117    pub fn with_hasher_in(
118        builder: ObjectBuilder<RawTable<(K, V), S, A>>,
119        hasher: S,
120        alloc: A,
121    ) -> Result<Self> {
122        let phm = Self {
123            table: builder.build_inplace(|tx| tx.write(RawTable::with_hasher_in(hasher, alloc)))?,
124        };
125
126        Ok(phm)
127    }
128
129    pub fn len(&self) -> usize {
130        self.table.base().len()
131    }
132
133    pub fn is_empty(&self) -> bool {
134        self.len() == 0
135    }
136
137    pub fn ctx(&self) -> CarryCtx {
138        self.table.base().carry_ctx()
139    }
140
141    pub fn clear(&mut self) -> Result<()> {
142        let mut tx = self.table.as_tx()?;
143        let mut base = tx.base_mut().owned();
144
145        base.clear();
146
147        Ok(())
148    }
149
150    pub fn iter(&self) -> Iter<'_, K, V> {
151        unsafe {
152            Iter {
153                _backing: self.table.base().backing(),
154                inner: self.table.base().iter(),
155                marker: PhantomData,
156            }
157        }
158    }
159
160    pub fn keys(&self) -> Keys<'_, K, V> {
161        Keys { inner: self.iter() }
162    }
163
164    pub fn values(&self) -> Values<'_, K, V> {
165        Values { inner: self.iter() }
166    }
167
168    pub fn iter_mut(&mut self) -> Result<IterMut<'_, K, V>> {
169        let mut tx = self.table.as_tx()?;
170        let mut base = tx.base_mut();
171
172        unsafe {
173            Ok(IterMut {
174                _backing: base.backing_mut().owned(),
175                inner: self.table.base().iter(),
176                marker: PhantomData,
177            })
178        }
179    }
180
181    pub fn values_mut(&mut self) -> Result<ValuesMut<'_, K, V>> {
182        Ok(ValuesMut {
183            inner: self.iter_mut()?,
184        })
185    }
186}
187
188impl<K: Invariant + Eq + Hash, V: Invariant, S: BuildHasher, A: Allocator>
189    PersistentHashMap<K, V, S, A>
190{
191    pub fn get<Q>(&self, k: &Q) -> Option<&V>
192    where
193        Q: Hash + Equivalent<K> + ?Sized,
194    {
195        let ctx = self.ctx();
196
197        match self.get_inner(k, &ctx) {
198            Some((_, v)) => Some(v),
199            None => None,
200        }
201    }
202
203    pub fn get_pair<Q>(&self, k: &Q, ctx: &impl Ctx) -> Option<&(K, V)>
204    where
205        Q: Hash + Equivalent<K> + ?Sized,
206    {
207        self.get_inner(k, ctx)
208    }
209
210    fn get_inner<Q>(&self, k: &Q, ctx: &impl Ctx) -> Option<&(K, V)>
211    where
212        Q: Hash + Equivalent<K> + ?Sized,
213    {
214        let hash = make_hash::<Q, S>(self.hasher(), k);
215        self.table.base().get(hash, equivalent_key(k), ctx)
216    }
217
218    pub fn hasher(&self) -> &S {
219        self.table.base().hasher()
220    }
221
222    fn remove_entry(&mut self, k: &K) -> Option<(K, V)> {
223        let hash = make_hash::<K, S>(self.hasher(), k);
224        let mut tx = self.table.as_tx().ok()?;
225        let mut base = tx.base_mut().owned();
226
227        let ctx = base.carry_ctx_mut(&base);
228
229        base.remove(hash, equivalent_key(k), &ctx)
230    }
231
232    pub fn remove(&mut self, k: &K) -> Option<V> {
233        match self.remove_entry(k) {
234            Some((_, v)) => Some(v),
235            None => None,
236        }
237    }
238}
239
240impl<K: Invariant + Eq + Hash, V: Invariant> PersistentHashMap<K, V> {
241    pub fn reserve(&mut self, additional: usize) -> Result<()> {
242        let mut tx = self.table.as_tx()?;
243        let mut base = tx.base_mut().owned();
244
245        let ctx = base.carry_ctx_mut(&base);
246
247        base.reserve(additional, make_hasher(self.hasher()), &ctx);
248        Ok(())
249    }
250}
251
252pub struct PHMsession<
253    'a,
254    K: Invariant,
255    V: Invariant,
256    S = DefaultHashBuilder,
257    A: Allocator = HashTableAlloc,
258> {
259    tx_base: RefMut<'a, RawTable<(K, V), S, A>>,
260    imm_base: Ref<'a, RawTable<(K, V), S, A>>,
261    ctx: CarryCtxMut<'a>,
262    _tx: TxObject<()>,
263}
264
265impl<K: Invariant + Eq + Hash, V: Invariant, S: BuildHasher> PHMsession<'_, K, V, S> {
266    pub fn insert(&mut self, k: K, v: V) -> Result<Option<V>> {
267        let hash = make_hash::<K, S>(self.imm_base.hasher(), &k);
268
269        match self.tx_base.find_or_find_insert_slot(
270            hash,
271            equivalent_key(&k),
272            make_hasher(self.imm_base.hasher()),
273            &self.ctx,
274        ) {
275            Ok(bucket) => {
276                let mut tx_ref = bucket.as_tx()?;
277                let mut mut_bucket = tx_ref.as_mut();
278                Ok(Some(std::mem::replace(&mut mut_bucket.1, v)))
279            }
280            Err(slot) => unsafe {
281                self.tx_base.insert_in_slot(hash, slot, (k, v), &self.ctx);
282                Ok(None)
283            },
284        }
285    }
286
287    // The mutable reference can only last as long as the write session
288    pub fn get_inner_mut<Q>(&mut self, k: &Q) -> Option<&mut (K, V)>
289    where
290        Q: Hash + Equivalent<K> + ?Sized,
291    {
292        let hash = make_hash::<Q, S>(self.imm_base.hasher(), k);
293
294        self.tx_base.get_mut(hash, equivalent_key(k), &self.ctx)
295    }
296
297    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
298    where
299        Q: Hash + Equivalent<K> + ?Sized,
300    {
301        match self.get_inner_mut(k) {
302            Some((_, v)) => Some(v),
303            None => None,
304        }
305    }
306}
307
308impl<K: Invariant + Eq + Hash, V: Invariant, S: BuildHasher>
309    PersistentHashMap<K, V, S, HashTableAlloc>
310{
311    pub fn write_session(&mut self) -> Result<PHMsession<'_, K, V, S>> {
312        let mut tx = self.table.as_tx()?;
313        let base = tx.base_mut().owned();
314        let imm_base = base.as_ref().owned();
315        let ctx = base.carry_ctx_mut(&base);
316
317        Ok(PHMsession {
318            tx_base: base,
319            imm_base,
320            ctx,
321            _tx: tx.into_unit(),
322        })
323    }
324
325    pub fn insert(&mut self, k: K, v: V) -> Result<Option<V>> {
326        let mut tx = self.table.as_tx()?;
327        let mut base = tx.base_mut();
328
329        let ctx = base.carry_ctx_mut(&base);
330        let hash = make_hash::<K, S>(base.hasher(), &k);
331
332        match base.find_or_find_insert_slot(
333            hash,
334            equivalent_key(&k),
335            make_hasher(self.hasher()),
336            &ctx,
337        ) {
338            Ok(bucket) => {
339                let mut tx_ref = bucket.as_tx()?;
340                let mut mut_bucket = tx_ref.as_mut();
341                Ok(Some(std::mem::replace(&mut mut_bucket.1, v)))
342            }
343            Err(slot) => unsafe {
344                base.insert_in_slot(hash, slot, (k, v), &ctx);
345                Ok(None)
346            },
347        }
348    }
349
350    pub unsafe fn resize(&mut self, capacity: usize) -> Result<()> {
351        let mut tx = self.table.as_tx()?;
352        let mut base = tx.base_mut().owned();
353
354        let mut ctx = base.carry_ctx_mut(&base);
355
356        base.resize(capacity, make_hasher(self.hasher()), &mut ctx)?;
357
358        Ok(())
359    }
360}
361
362pub struct Iter<'a, K: Invariant, V: Invariant> {
363    _backing: Ref<'a, u8>,
364    inner: RawIter<(K, V)>,
365    marker: PhantomData<(&'a K, &'a V)>,
366}
367
368impl<'a, K: Invariant, V: Invariant> Iterator for Iter<'a, K, V> {
369    type Item = (&'a K, &'a V);
370
371    fn next(&mut self) -> Option<(&'a K, &'a V)> {
372        // Avoid `Option::map` because it bloats LLVM IR.
373        match self.inner.next() {
374            Some(x) => unsafe {
375                let r = x.as_ref();
376                Some((&r.0, &r.1))
377            },
378            None => None,
379        }
380    }
381}
382
383pub struct IterMut<'a, K: Invariant, V: Invariant> {
384    _backing: RefMut<'a, u8>,
385    inner: RawIter<(K, V)>,
386    // To ensure invariance with respect to V
387    marker: PhantomData<(&'a K, &'a mut V)>,
388}
389
390impl<'a, K: Invariant, V: Invariant> Iterator for IterMut<'a, K, V> {
391    type Item = (&'a K, &'a mut V);
392
393    fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
394        match self.inner.next() {
395            Some(mut x) => unsafe {
396                let r = x.as_mut();
397                Some((&r.0, &mut r.1))
398            },
399            None => None,
400        }
401    }
402}
403
404pub struct Keys<'a, K: Invariant, V: Invariant> {
405    inner: Iter<'a, K, V>,
406}
407
408impl<'a, K: Invariant, V: Invariant> Iterator for Keys<'a, K, V> {
409    type Item = &'a K;
410
411    fn next(&mut self) -> Option<&'a K> {
412        match self.inner.next() {
413            Some((k, _)) => Some(k),
414            None => None,
415        }
416    }
417}
418
419pub struct Values<'a, K: Invariant, V: Invariant> {
420    inner: Iter<'a, K, V>,
421}
422
423impl<'a, K: Invariant, V: Invariant> Iterator for Values<'a, K, V> {
424    type Item = &'a V;
425
426    fn next(&mut self) -> Option<&'a V> {
427        match self.inner.next() {
428            Some((_, v)) => Some(v),
429            None => None,
430        }
431    }
432}
433
434pub struct ValuesMut<'a, K: Invariant, V: Invariant> {
435    inner: IterMut<'a, K, V>,
436}
437
438impl<'a, K: Invariant, V: Invariant> Iterator for ValuesMut<'a, K, V> {
439    type Item = &'a mut V;
440
441    fn next(&mut self) -> Option<&'a mut V> {
442        match self.inner.next() {
443            Some((_, v)) => Some(v),
444            None => None,
445        }
446    }
447}