twizzler/collections/hachage/
map.rs

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