Skip to main content

forest/cid_collections/
hash_set.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use super::*;
5use anyhow::Context as _;
6use bytes::Bytes;
7use cid::Cid;
8
9#[cfg(doc)]
10use std::collections::HashSet;
11use std::{path::Path, sync::LazyLock};
12
13pub trait CidHashSetLike {
14    /// Adds a value to the set.
15    ///
16    /// Returns whether the value was newly inserted.
17    fn insert(&mut self, cid: Cid) -> anyhow::Result<bool>;
18}
19
20/// A hash set implemented as a `HashMap` where the value is `()`.
21///
22/// See also [`HashSet`].
23#[derive(Default, Clone, Debug, PartialEq, Eq, derive_more::Deref)]
24pub struct CidHashSet {
25    inner: CidHashMap<()>,
26}
27
28impl CidHashSet {
29    /// Creates an empty `HashSet`.
30    ///
31    /// See also [`HashSet::new`].
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    /// Adds a value to the set.
37    ///
38    /// Returns whether the value was newly inserted.
39    ///
40    /// See also [`HashSet::insert`].
41    pub fn insert(&mut self, cid: Cid) -> bool {
42        self.inner.insert(cid, ()).is_none()
43    }
44}
45
46impl CidHashSetLike for CidHashSet {
47    fn insert(&mut self, cid: Cid) -> anyhow::Result<bool> {
48        Ok(self.insert(cid))
49    }
50}
51
52////////////////////
53// Collection Ops //
54////////////////////
55
56impl Extend<Cid> for CidHashSet {
57    fn extend<T: IntoIterator<Item = Cid>>(&mut self, iter: T) {
58        self.inner.extend(iter.into_iter().map(|it| (it, ())))
59    }
60}
61
62impl FromIterator<Cid> for CidHashSet {
63    fn from_iter<T: IntoIterator<Item = Cid>>(iter: T) -> Self {
64        let mut this = Self::new();
65        this.extend(iter);
66        this
67    }
68}
69
70/// A file-backed CID hash set.
71/// This is intended to be used for large sets of CIDs that may not fit in memory, such as when tracking seen CIDs during a chain export.
72pub struct FileBackedCidHashSet {
73    db: parity_db::Db,
74    // for dropping the temporary directory when the set is dropped
75    _dir: tempfile::TempDir,
76    lru: hashlink::LruCache<SmallCid, ()>,
77}
78
79impl FileBackedCidHashSet {
80    pub fn new(temp_dir_root: impl AsRef<Path>) -> anyhow::Result<Self> {
81        let dir = tempfile::tempdir_in(temp_dir_root.as_ref()).with_context(|| {
82            format!(
83                "failed to create temp dir in {}",
84                temp_dir_root.as_ref().display(),
85            )
86        })?;
87        let options = parity_db::Options {
88            path: dir.path().to_path_buf(),
89            sync_wal: false,
90            sync_data: false,
91            stats: false,
92            salt: None,
93            columns: vec![
94                parity_db::ColumnOptions {
95                    uniform: true,
96                    append_only: true,
97                    ..Default::default()
98                },
99                parity_db::ColumnOptions {
100                    append_only: true,
101                    ..Default::default()
102                },
103            ],
104            compression_threshold: Default::default(),
105        };
106        let db = parity_db::Db::open_or_create(&options).with_context(|| {
107            format!(
108                "failed to create temp parity-db at {}",
109                options.path.display()
110            )
111        })?;
112        Ok(Self {
113            db,
114            _dir: dir,
115            #[allow(clippy::disallowed_methods)]
116            lru: hashlink::LruCache::new(2 << 19), // ~80MiB for 1M entries
117        })
118    }
119
120    pub fn new_in_temp_dir() -> anyhow::Result<Self> {
121        Self::new(std::env::temp_dir())
122    }
123}
124
125impl CidHashSetLike for FileBackedCidHashSet {
126    fn insert(&mut self, cid: Cid) -> anyhow::Result<bool> {
127        static EMPTY_VALUE: LazyLock<Bytes> = LazyLock::new(|| Bytes::from_static(&[]));
128
129        let small = SmallCid::from(cid);
130        if self.lru.get(&small).is_some() {
131            return Ok(false);
132        }
133
134        let (col, key) = match &small {
135            SmallCid::Inline(c) => (0, c.digest().to_vec()),
136            SmallCid::Indirect(u) => (1, u.to_bytes()),
137        };
138        if self.db.get(col, &key).ok().flatten().is_some() {
139            self.lru.insert(small, ());
140            Ok(false)
141        } else {
142            self.db.commit_changes_bytes([(
143                col,
144                parity_db::Operation::Set(key, EMPTY_VALUE.clone()),
145            )])?;
146            self.lru.insert(small, ());
147            Ok(true)
148        }
149    }
150}
151
152#[cfg(test)]
153impl Default for FileBackedCidHashSet {
154    fn default() -> Self {
155        Self::new_in_temp_dir().expect("failed to create FileBackedCidHashSet")
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use ahash::HashSet;
163
164    #[quickcheck_macros::quickcheck]
165    fn test_cid_hashset(cids: HashSet<Cid>) {
166        let mut set = CidHashSet::default();
167        for cid in cids.iter() {
168            all_asserts::assert_true!(set.insert(*cid), "expected CID to be newly inserted");
169        }
170        for cid in cids.iter() {
171            all_asserts::assert_false!(set.insert(*cid), "expected CID to be present in the set");
172        }
173    }
174
175    #[quickcheck_macros::quickcheck]
176    fn test_file_backed_cid_hashset(cids: HashSet<Cid>) {
177        let mut set = FileBackedCidHashSet::default();
178        let dir = set._dir.path().to_path_buf();
179        for cid in cids.iter() {
180            all_asserts::assert_true!(
181                set.insert(*cid).unwrap(),
182                "expected CID to be newly inserted"
183            );
184        }
185        for cid in cids.iter() {
186            all_asserts::assert_false!(
187                set.insert(*cid).unwrap(),
188                "expected CID to be present in the set"
189            );
190        }
191        drop(set);
192        all_asserts::assert_false!(dir.exists(), "expected temporary directory to be deleted");
193    }
194}