Skip to main content

forest/cid_collections/
hash_map.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use super::{CidV1DagCborBlake2b256, MaybeCompactedCid, Uncompactable};
5use cid::Cid;
6#[cfg(doc)]
7use std::collections::HashMap;
8use std::collections::hash_map::{
9    Entry as StdEntry, IntoIter as StdIntoIter, OccupiedEntry as StdOccupiedEntry,
10    VacantEntry as StdVacantEntry,
11};
12
13/// A space-optimised hash map of [`Cid`]s, matching the API for [`std::collections::HashMap`].
14///
15/// We accept the implementation complexity of per-compaction-method `HashMap`s for
16/// the space savings, which are constant per-variant, rather than constant per-item.
17///
18/// This is dramatic for large maps!
19/// Using, e.g [`SmallCidNonEmptyVec`](super::SmallCidNonEmptyVec) will cost
20/// 25% more per-CID in the median case (32 B vs 40 B)
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct CidHashMap<V> {
23    compact: ahash::HashMap<CidV1DagCborBlake2b256, V>,
24    uncompact: ahash::HashMap<Uncompactable, V>,
25}
26
27impl<V> CidHashMap<V> {
28    /// Creates an empty `HashMap`.
29    ///
30    /// See also [`HashMap::new`].
31    pub fn new() -> Self {
32        Self::default()
33    }
34    /// Returns the number of elements in the map.
35    ///
36    /// See also [`HashMap::len`].
37    pub fn len(&self) -> usize {
38        let Self { compact, uncompact } = self;
39        compact.len() + uncompact.len()
40    }
41    /// Reflective of reserved capacity of this map.
42    pub fn total_capacity(&self) -> usize {
43        let Self { compact, uncompact } = self;
44        compact.capacity() + uncompact.capacity()
45    }
46    /// Returns `true` if the map contains a value for the specified key.
47    ///
48    /// See also [`HashMap::contains_key`].
49    pub fn contains_key(&self, key: &Cid) -> bool {
50        match MaybeCompactedCid::from(*key) {
51            MaybeCompactedCid::Compact(c) => self.compact.contains_key(&c),
52            MaybeCompactedCid::Uncompactable(u) => self.uncompact.contains_key(&u),
53        }
54    }
55    /// Returns a reference to the value corresponding to the key.
56    ///
57    /// See also [`HashMap::get`].
58    pub fn get(&self, key: &Cid) -> Option<&V> {
59        match MaybeCompactedCid::from(*key) {
60            MaybeCompactedCid::Compact(c) => self.compact.get(&c),
61            MaybeCompactedCid::Uncompactable(u) => self.uncompact.get(&u),
62        }
63    }
64    /// Inserts a key-value pair into the map.
65    ///
66    /// If the map did not have this key present, [`None`] is returned.
67    ///
68    /// If the map did have this key present, the value is updated, and the old
69    /// value is returned.
70    ///
71    /// See also [`HashMap::insert`].
72    pub fn insert(&mut self, key: Cid, value: V) -> Option<V> {
73        match MaybeCompactedCid::from(key) {
74            MaybeCompactedCid::Compact(c) => self.compact.insert(c, value),
75            MaybeCompactedCid::Uncompactable(u) => self.uncompact.insert(u, value),
76        }
77    }
78    /// Removes a key from the map, returning the value at the key if the key
79    /// was previously in the map.
80    ///
81    /// See also [`HashMap::remove`].
82    pub fn remove(&mut self, key: &Cid) -> Option<V> {
83        match MaybeCompactedCid::from(*key) {
84            MaybeCompactedCid::Compact(c) => self.compact.remove(&c),
85            MaybeCompactedCid::Uncompactable(u) => self.uncompact.remove(&u),
86        }
87    }
88    /// Shrinks the capacity of the internal map as much as possible.
89    pub fn shrink_to_fit(&mut self) {
90        self.compact.shrink_to_fit();
91        self.uncompact.shrink_to_fit();
92    }
93}
94
95///////////////
96// Entry API //
97///////////////
98
99impl<V> CidHashMap<V> {
100    /// Gets the given key's corresponding entry in the map for in-place manipulation.
101    ///
102    /// See also [`HashMap::entry`].
103    #[allow(dead_code)]
104    pub fn entry(&mut self, key: Cid) -> Entry<'_, V> {
105        match MaybeCompactedCid::from(key) {
106            MaybeCompactedCid::Compact(c) => match self.compact.entry(c) {
107                StdEntry::Occupied(o) => Entry::Occupied(OccupiedEntry {
108                    inner: OccupiedEntryInner::Compact(o),
109                }),
110                StdEntry::Vacant(v) => Entry::Vacant(VacantEntry {
111                    inner: VacantEntryInner::Compact(v),
112                }),
113            },
114            MaybeCompactedCid::Uncompactable(u) => match self.uncompact.entry(u) {
115                StdEntry::Occupied(o) => Entry::Occupied(OccupiedEntry {
116                    inner: OccupiedEntryInner::Uncompact(o),
117                }),
118                StdEntry::Vacant(v) => Entry::Vacant(VacantEntry {
119                    inner: VacantEntryInner::Uncompact(v),
120                }),
121            },
122        }
123    }
124}
125
126/// A view into a single entry in a map, which may either be vacant or occupied.
127///
128/// This `enum` is constructed using [`CidHashMap::entry`].
129#[allow(dead_code)]
130#[derive(Debug)]
131pub enum Entry<'a, V: 'a> {
132    /// An occupied entry.
133    Occupied(OccupiedEntry<'a, V>),
134    /// A vacant entry.
135    Vacant(VacantEntry<'a, V>),
136}
137
138/// A view into an occupied entry in a `HashMap`.
139/// It is part of the [`Entry`] enum.
140///
141/// See also [`std::collections::hash_map::OccupiedEntry`].
142#[allow(dead_code)]
143#[derive(Debug)]
144pub struct OccupiedEntry<'a, V> {
145    inner: OccupiedEntryInner<'a, V>,
146}
147
148impl<V> OccupiedEntry<'_, V> {
149    /// Gets a reference to the value in the entry.
150    ///
151    /// See also [`std::collections::hash_map::OccupiedEntry::get`].
152    #[allow(dead_code)]
153    pub fn get(&self) -> &V {
154        match &self.inner {
155            OccupiedEntryInner::Compact(c) => c.get(),
156            OccupiedEntryInner::Uncompact(u) => u.get(),
157        }
158    }
159}
160
161/// Hides compaction from users.
162#[allow(dead_code)]
163#[derive(Debug)]
164enum OccupiedEntryInner<'a, V> {
165    Compact(StdOccupiedEntry<'a, CidV1DagCborBlake2b256, V>),
166    Uncompact(StdOccupiedEntry<'a, Uncompactable, V>),
167}
168
169/// A view into a vacant entry in a `HashMap`.
170/// It is part of the [`Entry`] enum.
171///
172/// See also [`std::collections::hash_map::VacantEntry`].
173#[allow(dead_code)]
174#[derive(Debug)]
175pub struct VacantEntry<'a, V> {
176    inner: VacantEntryInner<'a, V>,
177}
178
179impl<'a, V> VacantEntry<'a, V> {
180    /// Sets the value of the entry with the `VacantEntry`'s key,
181    /// and returns a mutable reference to it.
182    ///
183    /// See also [`std::collections::hash_map::VacantEntry::insert`].
184    #[allow(dead_code)]
185    pub fn insert(self, value: V) -> &'a mut V {
186        match self.inner {
187            VacantEntryInner::Compact(c) => c.insert(value),
188            VacantEntryInner::Uncompact(u) => u.insert(value),
189        }
190    }
191}
192
193/// Hides compaction from users.
194#[allow(dead_code)]
195#[derive(Debug)]
196enum VacantEntryInner<'a, V> {
197    Compact(StdVacantEntry<'a, CidV1DagCborBlake2b256, V>),
198    Uncompact(StdVacantEntry<'a, Uncompactable, V>),
199}
200
201////////////////////
202// Collection Ops //
203////////////////////
204
205impl<V> Default for CidHashMap<V> {
206    fn default() -> Self {
207        Self {
208            compact: Default::default(),
209            uncompact: Default::default(),
210        }
211    }
212}
213
214impl<V> Extend<(Cid, V)> for CidHashMap<V> {
215    fn extend<T: IntoIterator<Item = (Cid, V)>>(&mut self, iter: T) {
216        for (cid, v) in iter {
217            match MaybeCompactedCid::from(cid) {
218                MaybeCompactedCid::Compact(compact) => {
219                    self.compact.insert(compact, v);
220                }
221                MaybeCompactedCid::Uncompactable(uncompact) => {
222                    self.uncompact.insert(uncompact, v);
223                }
224            };
225        }
226    }
227}
228
229impl<V> FromIterator<(Cid, V)> for CidHashMap<V> {
230    fn from_iter<T: IntoIterator<Item = (Cid, V)>>(iter: T) -> Self {
231        let mut this = Self::new();
232        this.extend(iter);
233        this
234    }
235}
236
237pub struct IntoIter<V> {
238    compact: StdIntoIter<CidV1DagCborBlake2b256, V>,
239    uncompact: StdIntoIter<Uncompactable, V>,
240}
241
242impl<V> Iterator for IntoIter<V> {
243    type Item = (Cid, V);
244
245    fn next(&mut self) -> Option<Self::Item> {
246        self.compact
247            .next()
248            .map(|(k, v)| (MaybeCompactedCid::Compact(k).into(), v))
249            .or_else(|| {
250                self.uncompact
251                    .next()
252                    .map(|(k, v)| (MaybeCompactedCid::Uncompactable(k).into(), v))
253            })
254    }
255    fn size_hint(&self) -> (usize, Option<usize>) {
256        join_size_hints(self.compact.size_hint(), self.uncompact.size_hint())
257    }
258}
259
260fn join_size_hints(
261    left: (usize, Option<usize>),
262    right: (usize, Option<usize>),
263) -> (usize, Option<usize>) {
264    let (l_lower, l_upper) = left;
265    let (r_lower, r_upper) = right;
266    let lower = l_lower.saturating_add(r_lower);
267    let upper = match (l_upper, r_upper) {
268        (Some(l), Some(r)) => l.checked_add(r),
269        _ => None,
270    };
271    (lower, upper)
272}
273
274impl<V> IntoIterator for CidHashMap<V> {
275    type Item = (Cid, V);
276
277    type IntoIter = IntoIter<V>;
278
279    fn into_iter(self) -> Self::IntoIter {
280        let Self { compact, uncompact } = self;
281        IntoIter {
282            compact: compact.into_iter(),
283            uncompact: uncompact.into_iter(),
284        }
285    }
286}
287
288//////////
289// Keys //
290//////////
291
292#[cfg(test)]
293use std::collections::hash_map::Keys as StdKeys;
294
295#[cfg(test)]
296impl<V> CidHashMap<V> {
297    /// An iterator visiting all keys in arbitrary order.
298    ///
299    /// In a notable departure from [`HashMap::keys`], the element type is [`Cid`], not [`&Cid`].
300    ///
301    pub fn keys(&self) -> Keys<'_, V> {
302        let Self { compact, uncompact } = self;
303        Keys {
304            compact: compact.keys(),
305            uncompact: uncompact.keys(),
306        }
307    }
308}
309
310/// An iterator over the keys of a `HashMap`.
311///
312/// See [`CidHashMap::keys`].
313#[cfg(test)]
314pub struct Keys<'a, V> {
315    compact: StdKeys<'a, CidV1DagCborBlake2b256, V>,
316    uncompact: StdKeys<'a, Uncompactable, V>,
317}
318
319#[cfg(test)]
320impl<V> Iterator for Keys<'_, V> {
321    type Item = Cid;
322
323    fn next(&mut self) -> Option<Self::Item> {
324        self.compact
325            .next()
326            .copied()
327            .map(MaybeCompactedCid::Compact)
328            .map(Into::into)
329            .or_else(|| {
330                self.uncompact
331                    .next()
332                    .copied()
333                    .map(MaybeCompactedCid::Uncompactable)
334                    .map(Into::into)
335            })
336    }
337
338    fn size_hint(&self) -> (usize, Option<usize>) {
339        join_size_hints(self.compact.size_hint(), self.uncompact.size_hint())
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346
347    #[derive(derive_quickcheck_arbitrary::Arbitrary, Clone, Debug)]
348    enum Operation {
349        ContainsKey(MaybeCompactedCid),
350        Get(MaybeCompactedCid),
351        Insert(MaybeCompactedCid, u8),
352        Entry { key: MaybeCompactedCid, value: u8 },
353    }
354
355    #[quickcheck_macros::quickcheck]
356    fn operations(operations: Vec<Operation>) {
357        use Operation as Op;
358
359        let mut subject = CidHashMap::default();
360        let mut reference = ahash::HashMap::default();
361        for operation in operations {
362            match operation {
363                Op::ContainsKey(key) => {
364                    let key = key.into();
365                    assert_eq!(subject.contains_key(&key), reference.contains_key(&key))
366                }
367                Op::Get(key) => {
368                    let key = key.into();
369                    assert_eq!(subject.get(&key), reference.get(&key))
370                }
371                Op::Insert(key, val) => {
372                    let key = key.into();
373                    assert_eq!(subject.insert(key, val), reference.insert(key, val))
374                }
375                Op::Entry { key, value } => {
376                    let key = key.into();
377                    match (subject.entry(key), reference.entry(key)) {
378                        (Entry::Occupied(subj), StdEntry::Occupied(refr)) => {
379                            assert_eq!(subj.get(), refr.get())
380                        }
381                        (Entry::Vacant(subj), StdEntry::Vacant(refr)) => {
382                            assert_eq!(subj.insert(value), refr.insert(value))
383                        }
384                        (subj, refr) => panic!("{subj:?}, {refr:?}"),
385                    }
386                }
387            }
388        }
389        assert_eq!(reference, ahash::HashMap::from_iter(subject));
390    }
391
392    #[quickcheck_macros::quickcheck]
393    fn collect(pairs: Vec<(Cid, u8)>) {
394        let refr = ahash::HashMap::from_iter(pairs.clone());
395        let via_subject = ahash::HashMap::from_iter(CidHashMap::from_iter(pairs));
396        assert_eq!(refr, via_subject);
397    }
398}