rkyv/collections/swiss_table/
index_map.rs

1//! An archived index map implementation based on Google's high-performance
2//! SwissTable hash map.
3
4use core::{
5    borrow::Borrow,
6    fmt,
7    hash::{Hash, Hasher},
8    iter::FusedIterator,
9    marker::PhantomData,
10    slice::{from_raw_parts, from_raw_parts_mut},
11};
12
13use munge::munge;
14use rancor::{Fallible, Source};
15
16use crate::{
17    collections::{
18        swiss_table::{ArchivedHashTable, HashTableResolver},
19        util::{Entry, EntryAdapter, EntryResolver},
20    },
21    hash::{hash_value, FxHasher64},
22    primitive::{ArchivedUsize, FixedUsize},
23    seal::Seal,
24    ser::{Allocator, Writer, WriterExt as _},
25    Place, Portable, RelPtr, Serialize,
26};
27
28/// An archived `IndexMap`.
29#[derive(Portable)]
30#[cfg_attr(
31    feature = "bytecheck",
32    derive(bytecheck::CheckBytes),
33    bytecheck(verify)
34)]
35#[rkyv(crate)]
36#[repr(C)]
37pub struct ArchivedIndexMap<K, V, H = FxHasher64> {
38    table: ArchivedHashTable<ArchivedUsize>,
39    entries: RelPtr<Entry<K, V>>,
40    _phantom: PhantomData<H>,
41}
42
43impl<K, V, H> ArchivedIndexMap<K, V, H> {
44    fn entries(&self) -> &[Entry<K, V>] {
45        unsafe { from_raw_parts(self.entries.as_ptr(), self.len()) }
46    }
47
48    fn entries_seal(this: Seal<'_, Self>) -> Seal<'_, [Entry<K, V>]> {
49        let len = this.len();
50        munge!(let Self { entries, .. } = this);
51        let slice =
52            unsafe { from_raw_parts_mut(RelPtr::as_mut_ptr(entries), len) };
53        Seal::new(slice)
54    }
55
56    /// Returns `true` if the map contains no elements.
57    pub const fn is_empty(&self) -> bool {
58        self.len() == 0
59    }
60
61    unsafe fn raw_iter(&self) -> RawIter<'_, K, V> {
62        unsafe { RawIter::new(self.entries.as_ptr().cast(), self.len()) }
63    }
64
65    /// Returns an iterator over the key-value pairs of the map in order
66    pub fn iter(&self) -> Iter<'_, K, V> {
67        Iter {
68            inner: unsafe { self.raw_iter() },
69        }
70    }
71
72    /// Returns an iterator over the keys of the map in order
73    pub fn keys(&self) -> Keys<'_, K, V> {
74        Keys {
75            inner: unsafe { self.raw_iter() },
76        }
77    }
78
79    /// Gets the number of items in the index map.
80    pub const fn len(&self) -> usize {
81        self.table.len()
82    }
83
84    /// Returns an iterator over the values of the map in order.
85    pub fn values(&self) -> Values<'_, K, V> {
86        Values {
87            inner: unsafe { self.raw_iter() },
88        }
89    }
90}
91
92impl<K, V, H: Hasher + Default> ArchivedIndexMap<K, V, H> {
93    /// Gets the index, key, and value corresponding to the supplied key using
94    /// the given comparison function.
95    pub fn get_full_with<Q, C>(
96        &self,
97        key: &Q,
98        cmp: C,
99    ) -> Option<(usize, &K, &V)>
100    where
101        Q: Hash + Eq + ?Sized,
102        C: Fn(&Q, &K) -> bool,
103    {
104        let index = self.get_index_of_with(key, cmp)?;
105        let entries = self.entries();
106        let entry = &entries[index];
107        Some((index, &entry.key, &entry.value))
108    }
109
110    /// Gets the index, key, and value corresponding to the supplied key.
111    pub fn get_full<Q>(&self, key: &Q) -> Option<(usize, &K, &V)>
112    where
113        K: Borrow<Q>,
114        Q: Hash + Eq + ?Sized,
115    {
116        self.get_full_with(key, |q, k| q == k.borrow())
117    }
118
119    /// Returns the key-value pair corresponding to the supplied key using the
120    /// given comparison function.
121    pub fn get_key_value_with<Q, C>(&self, key: &Q, cmp: C) -> Option<(&K, &V)>
122    where
123        Q: Hash + Eq + ?Sized,
124        C: Fn(&Q, &K) -> bool,
125    {
126        let (_, k, v) = self.get_full_with(key, cmp)?;
127        Some((k, v))
128    }
129
130    /// Returns the key-value pair corresponding to the supplied key.
131    pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, &V)>
132    where
133        K: Borrow<Q>,
134        Q: Hash + Eq + ?Sized,
135    {
136        let (_, k, v) = self.get_full(key)?;
137        Some((k, v))
138    }
139
140    /// Returns a reference to the value corresponding to the supplied key using
141    /// the given comparison function.
142    pub fn get_with<Q, C>(&self, key: &Q, cmp: C) -> Option<&V>
143    where
144        Q: Hash + Eq + ?Sized,
145        C: Fn(&Q, &K) -> bool,
146    {
147        Some(self.get_full_with(key, cmp)?.2)
148    }
149
150    /// Returns a reference to the value corresponding to the supplied key.
151    pub fn get<Q>(&self, key: &Q) -> Option<&V>
152    where
153        K: Borrow<Q>,
154        Q: Hash + Eq + ?Sized,
155    {
156        Some(self.get_full(key)?.2)
157    }
158
159    /// Gets the mutable index, key, and value corresponding to the supplied key
160    /// using the given comparison function.
161    pub fn get_full_seal_with<'a, Q, C>(
162        this: Seal<'a, Self>,
163        key: &Q,
164        cmp: C,
165    ) -> Option<(usize, &'a K, Seal<'a, V>)>
166    where
167        Q: Hash + Eq + ?Sized,
168        C: Fn(&Q, &K) -> bool,
169    {
170        let index = this.get_index_of_with(key, cmp)?;
171        let entry = Seal::index(Self::entries_seal(this), index);
172        munge!(let Entry { key, value } = entry);
173        Some((index, key.unseal_ref(), value))
174    }
175
176    /// Gets the mutable index, key, and value corresponding to the supplied
177    /// key.
178    pub fn get_full_seal<'a, Q>(
179        this: Seal<'a, Self>,
180        key: &Q,
181    ) -> Option<(usize, &'a K, Seal<'a, V>)>
182    where
183        K: Borrow<Q>,
184        Q: Hash + Eq + ?Sized,
185    {
186        Self::get_full_seal_with(this, key, |q, k| q == k.borrow())
187    }
188
189    /// Returns the mutable key-value pair corresponding to the supplied key
190    /// using the given comparison function.
191    pub fn get_key_value_seal_with<'a, Q, C>(
192        this: Seal<'a, Self>,
193        key: &Q,
194        cmp: C,
195    ) -> Option<(&'a K, Seal<'a, V>)>
196    where
197        K: Borrow<Q>,
198        Q: Hash + Eq + ?Sized,
199        C: Fn(&Q, &K) -> bool,
200    {
201        let (_, k, v) = Self::get_full_seal_with(this, key, cmp)?;
202        Some((k, v))
203    }
204
205    /// Returns the mutable key-value pair corresponding to the supplied key.
206    pub fn get_key_value_seal<'a, Q>(
207        this: Seal<'a, Self>,
208        key: &Q,
209    ) -> Option<(&'a K, Seal<'a, V>)>
210    where
211        K: Borrow<Q>,
212        Q: Hash + Eq + ?Sized,
213    {
214        let (_, k, v) = Self::get_full_seal(this, key)?;
215        Some((k, v))
216    }
217
218    /// Returns a mutable reference to the value corresponding to the supplied
219    /// key using the given comparison function.
220    pub fn get_seal_with<'a, Q, C>(
221        this: Seal<'a, Self>,
222        key: &Q,
223        cmp: C,
224    ) -> Option<Seal<'a, V>>
225    where
226        K: Borrow<Q>,
227        Q: Hash + Eq + ?Sized,
228        C: Fn(&Q, &K) -> bool,
229    {
230        Some(Self::get_full_seal_with(this, key, cmp)?.2)
231    }
232
233    /// Returns a mutable reference to the value corresponding to the supplied
234    /// key.
235    pub fn get_seal<'a, Q>(this: Seal<'a, Self>, key: &Q) -> Option<Seal<'a, V>>
236    where
237        K: Borrow<Q>,
238        Q: Hash + Eq + ?Sized,
239    {
240        Some(Self::get_full_seal(this, key)?.2)
241    }
242
243    /// Returns whether a key is present in the hash map.
244    pub fn contains_key<Q>(&self, key: &Q) -> bool
245    where
246        K: Borrow<Q>,
247        Q: Hash + Eq + ?Sized,
248    {
249        self.get(key).is_some()
250    }
251
252    /// Gets a key-value pair by index.
253    pub fn get_index(&self, index: usize) -> Option<(&K, &V)> {
254        if index < self.len() {
255            let entry = &self.entries()[index];
256            Some((&entry.key, &entry.value))
257        } else {
258            None
259        }
260    }
261
262    /// Gets the index of a key if it exists in the map using the given
263    /// comparison function.
264    pub fn get_index_of_with<Q, C>(&self, key: &Q, cmp: C) -> Option<usize>
265    where
266        Q: Hash + Eq + ?Sized,
267        C: Fn(&Q, &K) -> bool,
268    {
269        let entries = self.entries();
270        let index = self.table.get_with(hash_value::<Q, H>(key), |i| {
271            cmp(key, &entries[i.to_native() as usize].key)
272        })?;
273        Some(index.to_native() as usize)
274    }
275
276    /// Gets the index of a key if it exists in the map.
277    pub fn get_index_of<Q>(&self, key: &Q) -> Option<usize>
278    where
279        K: Borrow<Q>,
280        Q: Hash + Eq + ?Sized,
281    {
282        self.get_index_of_with(key, |q, k| q == k.borrow())
283    }
284
285    /// Resolves an archived index map from a given length and parameters.
286    pub fn resolve_from_len(
287        len: usize,
288        load_factor: (usize, usize),
289        resolver: IndexMapResolver,
290        out: Place<Self>,
291    ) {
292        munge!(let ArchivedIndexMap { table, entries, _phantom: _ } = out);
293        ArchivedHashTable::resolve_from_len(
294            len,
295            load_factor,
296            resolver.table_resolver,
297            table,
298        );
299        RelPtr::emplace(resolver.entries_pos as usize, entries);
300    }
301
302    /// Serializes an iterator of key-value pairs as an index map.
303    pub fn serialize_from_iter<I, BKU, BVU, KU, VU, S>(
304        iter: I,
305        load_factor: (usize, usize),
306        serializer: &mut S,
307    ) -> Result<IndexMapResolver, S::Error>
308    where
309        I: Clone + ExactSizeIterator<Item = (BKU, BVU)>,
310        BKU: Borrow<KU>,
311        BVU: Borrow<VU>,
312        KU: Serialize<S, Archived = K> + Hash + Eq,
313        VU: Serialize<S, Archived = V>,
314        S: Fallible + Writer + Allocator + ?Sized,
315        S::Error: Source,
316    {
317        use crate::util::SerVec;
318
319        // Serialize hash table
320        let table_resolver =
321            ArchivedHashTable::<ArchivedUsize>::serialize_from_iter(
322                0..iter.len(),
323                iter.clone()
324                    .map(|(key, _)| hash_value::<KU, H>(key.borrow())),
325                load_factor,
326                serializer,
327            )?;
328
329        // Serialize entries
330        SerVec::with_capacity(
331            serializer,
332            iter.len(),
333            |resolvers, serializer| {
334                for (key, value) in iter.clone() {
335                    resolvers.push(EntryResolver {
336                        key: key.borrow().serialize(serializer)?,
337                        value: value.borrow().serialize(serializer)?,
338                    });
339                }
340
341                let entries_pos = serializer.align_for::<Entry<K, V>>()?;
342                for ((key, value), resolver) in
343                    iter.clone().zip(resolvers.drain())
344                {
345                    unsafe {
346                        serializer.resolve_aligned(
347                            &EntryAdapter::new(key, value),
348                            resolver,
349                        )?;
350                    }
351                }
352
353                Ok(IndexMapResolver {
354                    table_resolver,
355                    entries_pos: entries_pos as FixedUsize,
356                })
357            },
358        )?
359    }
360}
361
362impl<K, V, H> fmt::Debug for ArchivedIndexMap<K, V, H>
363where
364    K: fmt::Debug,
365    V: fmt::Debug,
366{
367    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
368        f.debug_map().entries(self.iter()).finish()
369    }
370}
371
372impl<K, V, H> PartialEq for ArchivedIndexMap<K, V, H>
373where
374    K: PartialEq,
375    V: PartialEq,
376{
377    fn eq(&self, other: &Self) -> bool {
378        self.iter().eq(other.iter())
379    }
380}
381
382impl<K: Eq, V: Eq, H> Eq for ArchivedIndexMap<K, V, H> {}
383
384struct RawIter<'a, K, V> {
385    current: *const Entry<K, V>,
386    remaining: usize,
387    _phantom: PhantomData<(&'a K, &'a V)>,
388}
389
390impl<K, V> RawIter<'_, K, V> {
391    fn new(pairs: *const Entry<K, V>, len: usize) -> Self {
392        Self {
393            current: pairs,
394            remaining: len,
395            _phantom: PhantomData,
396        }
397    }
398}
399
400impl<'a, K, V> Iterator for RawIter<'a, K, V> {
401    type Item = (&'a K, &'a V);
402
403    fn next(&mut self) -> Option<Self::Item> {
404        unsafe {
405            if self.remaining == 0 {
406                None
407            } else {
408                let result = self.current;
409                self.current = self.current.add(1);
410                self.remaining -= 1;
411                let entry = &*result;
412                Some((&entry.key, &entry.value))
413            }
414        }
415    }
416
417    fn size_hint(&self) -> (usize, Option<usize>) {
418        (self.remaining, Some(self.remaining))
419    }
420}
421
422impl<K, V> ExactSizeIterator for RawIter<'_, K, V> {}
423impl<K, V> FusedIterator for RawIter<'_, K, V> {}
424
425/// An iterator over the key-value pairs of an index map.
426#[repr(transparent)]
427pub struct Iter<'a, K, V> {
428    inner: RawIter<'a, K, V>,
429}
430
431impl<'a, K, V> Iterator for Iter<'a, K, V> {
432    type Item = (&'a K, &'a V);
433
434    fn next(&mut self) -> Option<Self::Item> {
435        self.inner.next()
436    }
437
438    fn size_hint(&self) -> (usize, Option<usize>) {
439        self.inner.size_hint()
440    }
441}
442
443impl<K, V> ExactSizeIterator for Iter<'_, K, V> {}
444impl<K, V> FusedIterator for Iter<'_, K, V> {}
445
446/// An iterator over the keys of an index map.
447#[repr(transparent)]
448pub struct Keys<'a, K, V> {
449    inner: RawIter<'a, K, V>,
450}
451
452impl<'a, K, V> Iterator for Keys<'a, K, V> {
453    type Item = &'a K;
454
455    fn next(&mut self) -> Option<Self::Item> {
456        self.inner.next().map(|(k, _)| k)
457    }
458
459    fn size_hint(&self) -> (usize, Option<usize>) {
460        self.inner.size_hint()
461    }
462}
463
464impl<K, V> ExactSizeIterator for Keys<'_, K, V> {}
465impl<K, V> FusedIterator for Keys<'_, K, V> {}
466
467/// An iterator over the values of an index map.
468#[repr(transparent)]
469pub struct Values<'a, K, V> {
470    inner: RawIter<'a, K, V>,
471}
472
473impl<'a, K, V> Iterator for Values<'a, K, V> {
474    type Item = &'a V;
475
476    fn next(&mut self) -> Option<Self::Item> {
477        self.inner.next().map(|(_, v)| v)
478    }
479
480    fn size_hint(&self) -> (usize, Option<usize>) {
481        self.inner.size_hint()
482    }
483}
484
485impl<K, V> ExactSizeIterator for Values<'_, K, V> {}
486impl<K, V> FusedIterator for Values<'_, K, V> {}
487
488/// The resolver for an `IndexMap`.
489pub struct IndexMapResolver {
490    table_resolver: HashTableResolver,
491    entries_pos: FixedUsize,
492}
493
494#[cfg(feature = "bytecheck")]
495mod verify {
496    use bytecheck::{CheckBytes, Verify};
497    use rancor::{Fallible, Source};
498
499    use super::ArchivedIndexMap;
500    use crate::{
501        collections::util::Entry,
502        validation::{ArchiveContext, ArchiveContextExt},
503    };
504
505    unsafe impl<C, K, V, H> Verify<C> for ArchivedIndexMap<K, V, H>
506    where
507        C: Fallible + ArchiveContext + ?Sized,
508        C::Error: Source,
509        K: CheckBytes<C>,
510        V: CheckBytes<C>,
511    {
512        fn verify(
513            &self,
514            context: &mut C,
515        ) -> Result<(), <C as Fallible>::Error> {
516            let ptr = core::ptr::slice_from_raw_parts(
517                self.entries.as_ptr_wrapping(),
518                self.table.len(),
519            );
520
521            context.in_subtree(ptr, |context| {
522                // SAFETY: `in_subtree` has checked that `ptr` is aligned and
523                // points to enough bytes to represent its slice.
524                unsafe { <[Entry<K, V>]>::check_bytes(ptr, context) }
525            })
526        }
527    }
528}
rkyv/collections/swiss_table/index_map.rs

rkyv/collections/swiss_table/
index_map.rs
