forest/cid_collections/
hash_set.rs1use 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 fn insert(&mut self, cid: Cid) -> anyhow::Result<bool>;
18}
19
20#[derive(Default, Clone, Debug, PartialEq, Eq, derive_more::Deref)]
24pub struct CidHashSet {
25 inner: CidHashMap<()>,
26}
27
28impl CidHashSet {
29 pub fn new() -> Self {
33 Self::default()
34 }
35
36 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
52impl 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
70pub struct FileBackedCidHashSet {
73 db: parity_db::Db,
74 _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), })
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}