rkyv/collections/index_map/
mod.rs

1//! Archived index map implementation.
2//!
3//! During archiving, hashmaps are built into minimal perfect hashmaps using
4//! [compress, hash and displace](http://cmph.sourceforge.net/papers/esa09.pdf).
5
6#[cfg(feature = "validation")]
7pub mod validation;
8
9use crate::{
10    collections::{
11        hash_index::{ArchivedHashIndex, HashBuilder, HashIndexResolver},
12        util::Entry,
13    },
14    out_field, Archived, RelPtr,
15};
16use core::{borrow::Borrow, fmt, hash::Hash, iter::FusedIterator, marker::PhantomData};
17
18/// An archived `IndexMap`.
19#[cfg_attr(feature = "strict", repr(C))]
20pub struct ArchivedIndexMap<K, V> {
21    index: ArchivedHashIndex,
22    pivots: RelPtr<Archived<usize>>,
23    entries: RelPtr<Entry<K, V>>,
24}
25
26impl<K, V> ArchivedIndexMap<K, V> {
27    #[inline]
28    unsafe fn pivot(&self, index: usize) -> usize {
29        from_archived!(*self.pivots.as_ptr().add(index)) as usize
30    }
31
32    #[inline]
33    unsafe fn entry(&self, index: usize) -> &Entry<K, V> {
34        &*self.entries.as_ptr().add(index)
35    }
36
37    #[inline]
38    fn find<Q: ?Sized>(&self, k: &Q) -> Option<usize>
39    where
40        K: Borrow<Q>,
41        Q: Hash + Eq,
42    {
43        self.index.index(k).and_then(|pivot_index| {
44            let index = unsafe { self.pivot(pivot_index) };
45            let entry = unsafe { self.entry(index) };
46            if entry.key.borrow() == k {
47                Some(index)
48            } else {
49                None
50            }
51        })
52    }
53
54    /// Returns whether a key is present in the hash map.
55    #[inline]
56    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
57    where
58        K: Borrow<Q>,
59        Q: Hash + Eq,
60    {
61        self.find(k).is_some()
62    }
63
64    /// Returns the first key-value pair.
65    #[inline]
66    pub fn first(&self) -> Option<(&K, &V)> {
67        if !self.is_empty() {
68            let entry = unsafe { self.entry(0) };
69            Some((&entry.key, &entry.value))
70        } else {
71            None
72        }
73    }
74
75    /// Gets the value associated with the given key.
76    #[inline]
77    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
78    where
79        K: Borrow<Q>,
80        Q: Hash + Eq,
81    {
82        self.find(k)
83            .map(|index| unsafe { &self.entry(index).value })
84    }
85
86    /// Gets the index, key, and value associated with the given key.
87    #[inline]
88    pub fn get_full<Q: ?Sized>(&self, k: &Q) -> Option<(usize, &K, &V)>
89    where
90        K: Borrow<Q>,
91        Q: Hash + Eq,
92    {
93        self.find(k).map(|index| {
94            let entry = unsafe { &self.entry(index) };
95            (index, &entry.key, &entry.value)
96        })
97    }
98
99    /// Gets a key-value pair by index.
100    #[inline]
101    pub fn get_index(&self, index: usize) -> Option<(&K, &V)> {
102        if index < self.len() {
103            let entry = unsafe { &self.entry(index) };
104            Some((&entry.key, &entry.value))
105        } else {
106            None
107        }
108    }
109
110    /// Gets the index of a key if it exists in the map.
111    #[inline]
112    pub fn get_index_of<Q: ?Sized>(&self, key: &Q) -> Option<usize>
113    where
114        K: Borrow<Q>,
115        Q: Hash + Eq,
116    {
117        self.find(key)
118    }
119
120    /// Gets the key-value pair associated with the given key.
121    #[inline]
122    pub fn get_key_value<Q: ?Sized>(&self, k: &Q) -> Option<(&K, &V)>
123    where
124        K: Borrow<Q>,
125        Q: Hash + Eq,
126    {
127        self.find(k).map(|index| {
128            let entry = unsafe { &self.entry(index) };
129            (&entry.key, &entry.value)
130        })
131    }
132
133    /// Gets the hasher for this index map.
134    #[inline]
135    pub fn hasher(&self) -> HashBuilder {
136        self.index.hasher()
137    }
138
139    /// Returns `true` if the map contains no elements.
140    #[inline]
141    pub const fn is_empty(&self) -> bool {
142        self.len() == 0
143    }
144
145    #[inline]
146    fn raw_iter(&self) -> RawIter<K, V> {
147        RawIter::new(self.entries.as_ptr().cast(), self.len())
148    }
149
150    /// Returns an iterator over the key-value pairs of the map in order
151    #[inline]
152    pub fn iter(&self) -> Iter<K, V> {
153        Iter {
154            inner: self.raw_iter(),
155        }
156    }
157
158    /// Returns an iterator over the keys of the map in order
159    #[inline]
160    pub fn keys(&self) -> Keys<K, V> {
161        Keys {
162            inner: self.raw_iter(),
163        }
164    }
165
166    /// Returns the last key-value pair.
167    #[inline]
168    pub fn last(&self) -> Option<(&K, &V)> {
169        if !self.is_empty() {
170            let entry = unsafe { self.entry(self.len() - 1) };
171            Some((&entry.key, &entry.value))
172        } else {
173            None
174        }
175    }
176
177    /// Gets the number of items in the index map.
178    #[inline]
179    pub const fn len(&self) -> usize {
180        self.index.len()
181    }
182
183    /// Returns an iterator over the values of the map in order.
184    #[inline]
185    pub fn values(&self) -> Values<K, V> {
186        Values {
187            inner: self.raw_iter(),
188        }
189    }
190
191    /// Resolves an archived index map from a given length and parameters.
192    ///
193    /// # Safety
194    ///
195    /// - `len` must be the number of elements that were serialized
196    /// - `pos` must be the position of `out` within the archive
197    /// - `resolver` must be the result of serializing a hash map
198    pub unsafe fn resolve_from_len(
199        len: usize,
200        pos: usize,
201        resolver: IndexMapResolver,
202        out: *mut Self,
203    ) {
204        let (fp, fo) = out_field!(out.index);
205        ArchivedHashIndex::resolve_from_len(len, pos + fp, resolver.index_resolver, fo);
206
207        let (fp, fo) = out_field!(out.pivots);
208        RelPtr::emplace(pos + fp, resolver.pivots_pos, fo);
209
210        let (fp, fo) = out_field!(out.entries);
211        RelPtr::emplace(pos + fp, resolver.entries_pos, fo);
212    }
213}
214
215#[cfg(feature = "alloc")]
216const _: () = {
217    use crate::{
218        ser::{ScratchSpace, Serializer},
219        Serialize,
220    };
221
222    impl<K, V> ArchivedIndexMap<K, V> {
223        /// Serializes an iterator of key-value pairs as an index map.
224        ///
225        /// # Safety
226        ///
227        /// - The keys returned by the iterator must be unique
228        /// - The index function must return the index of the given key within the iterator
229        pub unsafe fn serialize_from_iter_index<'a, UK, UV, I, F, S>(
230            iter: I,
231            index: F,
232            serializer: &mut S,
233        ) -> Result<IndexMapResolver, S::Error>
234        where
235            UK: 'a + Serialize<S, Archived = K> + Hash + Eq,
236            UV: 'a + Serialize<S, Archived = V>,
237            I: Clone + ExactSizeIterator<Item = (&'a UK, &'a UV)>,
238            F: Fn(&UK) -> usize,
239            S: Serializer + ScratchSpace + ?Sized,
240        {
241            use crate::ScratchVec;
242
243            let len = iter.len();
244
245            let mut entries = ScratchVec::new(serializer, iter.len())?;
246            entries.set_len(len);
247            let index_resolver =
248                ArchivedHashIndex::build_and_serialize(iter.clone(), serializer, &mut entries)?;
249            let mut entries = entries.assume_init();
250
251            // Serialize entries
252            let mut resolvers = ScratchVec::new(serializer, iter.len())?;
253            for (key, value) in iter.clone() {
254                resolvers.push((key.serialize(serializer)?, value.serialize(serializer)?));
255            }
256
257            let entries_pos = serializer.align_for::<Entry<K, V>>()?;
258            for ((key, value), (key_resolver, value_resolver)) in iter.zip(resolvers.drain(..)) {
259                serializer
260                    .resolve_aligned(&Entry { key, value }, (key_resolver, value_resolver))?;
261            }
262
263            // Serialize pivots
264            let pivots_pos = serializer.align_for::<Archived<usize>>()?;
265            for (key, _) in entries.drain(..) {
266                serializer.resolve_aligned(&index(key), ())?;
267            }
268
269            // Free scratch vecs
270            resolvers.free(serializer)?;
271            entries.free(serializer)?;
272
273            Ok(IndexMapResolver {
274                index_resolver,
275                pivots_pos,
276                entries_pos,
277            })
278        }
279    }
280};
281
282impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for ArchivedIndexMap<K, V> {
283    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284        f.debug_map().entries(self.iter()).finish()
285    }
286}
287
288impl<K: PartialEq, V: PartialEq> PartialEq for ArchivedIndexMap<K, V> {
289    fn eq(&self, other: &Self) -> bool {
290        self.iter().eq(other.iter())
291    }
292}
293
294struct RawIter<'a, K, V> {
295    current: *const Entry<K, V>,
296    remaining: usize,
297    _phantom: PhantomData<(&'a K, &'a V)>,
298}
299
300impl<'a, K, V> RawIter<'a, K, V> {
301    #[inline]
302    fn new(pairs: *const Entry<K, V>, len: usize) -> Self {
303        Self {
304            current: pairs,
305            remaining: len,
306            _phantom: PhantomData,
307        }
308    }
309}
310
311impl<'a, K, V> Iterator for RawIter<'a, K, V> {
312    type Item = (&'a K, &'a V);
313
314    #[inline]
315    fn next(&mut self) -> Option<Self::Item> {
316        unsafe {
317            if self.remaining == 0 {
318                None
319            } else {
320                let result = self.current;
321                self.current = self.current.add(1);
322                self.remaining -= 1;
323                let entry = &*result;
324                Some((&entry.key, &entry.value))
325            }
326        }
327    }
328
329    #[inline]
330    fn size_hint(&self) -> (usize, Option<usize>) {
331        (self.remaining, Some(self.remaining))
332    }
333}
334
335impl<'a, K, V> ExactSizeIterator for RawIter<'a, K, V> {}
336impl<'a, K, V> FusedIterator for RawIter<'a, K, V> {}
337
338/// An iterator over the key-value pairs of an index map.
339#[repr(transparent)]
340pub struct Iter<'a, K, V> {
341    inner: RawIter<'a, K, V>,
342}
343
344impl<'a, K, V> Iterator for Iter<'a, K, V> {
345    type Item = (&'a K, &'a V);
346
347    #[inline]
348    fn next(&mut self) -> Option<Self::Item> {
349        self.inner.next()
350    }
351
352    #[inline]
353    fn size_hint(&self) -> (usize, Option<usize>) {
354        self.inner.size_hint()
355    }
356}
357
358impl<K, V> ExactSizeIterator for Iter<'_, K, V> {}
359impl<K, V> FusedIterator for Iter<'_, K, V> {}
360
361/// An iterator over the keys of an index map.
362#[repr(transparent)]
363pub struct Keys<'a, K, V> {
364    inner: RawIter<'a, K, V>,
365}
366
367impl<'a, K, V> Iterator for Keys<'a, K, V> {
368    type Item = &'a K;
369
370    #[inline]
371    fn next(&mut self) -> Option<Self::Item> {
372        self.inner.next().map(|(k, _)| k)
373    }
374
375    #[inline]
376    fn size_hint(&self) -> (usize, Option<usize>) {
377        self.inner.size_hint()
378    }
379}
380
381impl<K, V> ExactSizeIterator for Keys<'_, K, V> {}
382impl<K, V> FusedIterator for Keys<'_, K, V> {}
383
384/// An iterator over the values of an index map.
385#[repr(transparent)]
386pub struct Values<'a, K, V> {
387    inner: RawIter<'a, K, V>,
388}
389
390impl<'a, K, V> Iterator for Values<'a, K, V> {
391    type Item = &'a V;
392
393    #[inline]
394    fn next(&mut self) -> Option<Self::Item> {
395        self.inner.next().map(|(_, v)| v)
396    }
397
398    #[inline]
399    fn size_hint(&self) -> (usize, Option<usize>) {
400        self.inner.size_hint()
401    }
402}
403
404impl<K, V> ExactSizeIterator for Values<'_, K, V> {}
405impl<K, V> FusedIterator for Values<'_, K, V> {}
406
407// Archive implementations
408
409/// The resolver for an `IndexMap`.
410pub struct IndexMapResolver {
411    index_resolver: HashIndexResolver,
412    pivots_pos: usize,
413    entries_pos: usize,
414}