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