1use crate::merkle::error::Error;
2use crate::{
3 Hash, InternalNodesMap, KeyId, LeafNodesMap, MerkleBlob, Node, NodeHashToDeltaReaderNode,
4 NodeHashToIndex, ParentFirstIterator, TreeIndex, ValueId,
5};
6#[cfg(feature = "py-bindings")]
7use pyo3::{PyResult, Python, pyclass, pymethods};
8use rayon::iter::{IntoParallelIterator, ParallelExtend, ParallelIterator};
9use std::collections::{HashMap, HashSet};
10use std::path::PathBuf;
11
12pub enum DeltaReaderNode {
13 Internal { left: Hash, right: Hash },
14 Leaf { key: KeyId, value: ValueId },
15}
16
17#[cfg_attr(feature = "py-bindings", pyclass)]
18pub struct DeltaFileCache {
19 hash_to_index: NodeHashToIndex,
20 previous_hashes: HashSet<Hash>,
21 merkle_blob: MerkleBlob,
22}
23
24impl DeltaFileCache {
25 pub fn new(path: &PathBuf) -> Result<Self, Error> {
26 let merkle_blob = MerkleBlob::from_path(path)?;
27 let hash_to_index = merkle_blob.get_hashes_indexes(false)?;
28 Ok(Self {
29 hash_to_index,
30 previous_hashes: HashSet::new(),
31 merkle_blob,
32 })
33 }
34
35 pub fn load_previous_hashes(&mut self, path: &PathBuf) -> Result<(), Error> {
36 let blob = crate::zstd_decode_path(path)?;
37 self.previous_hashes = HashSet::new();
38
39 if !blob.is_empty() {
40 for item in ParentFirstIterator::new(&blob, None) {
41 let (_, block) = item?;
42 self.previous_hashes.insert(block.node.hash());
43 }
44 }
45 Ok(())
46 }
47
48 pub fn get_raw_node(&self, index: TreeIndex) -> Result<Node, Error> {
49 self.merkle_blob.get_node(index)
50 }
51
52 pub fn get_hash_at_index(&self, index: TreeIndex) -> Result<Option<Hash>, Error> {
53 self.merkle_blob.get_hash_at_index(index)
54 }
55
56 pub fn seen_previous_hash(&self, hash: Hash) -> bool {
57 self.previous_hashes.contains(&hash)
58 }
59
60 pub fn get_index(&self, hash: Hash) -> Result<TreeIndex, Error> {
61 self.hash_to_index
62 .get(&hash)
63 .copied()
64 .ok_or(Error::HashNotFound(hash))
65 }
66}
67
68#[cfg_attr(feature = "py-bindings", pyclass)]
69pub struct DeltaReader {
70 nodes: NodeHashToDeltaReaderNode,
71}
72
73impl DeltaReader {
74 pub fn new(internal_nodes: InternalNodesMap, leaf_nodes: LeafNodesMap) -> Result<Self, Error> {
75 let mut nodes = NodeHashToDeltaReaderNode::new();
76
77 for (hash, (left, right)) in internal_nodes {
78 nodes.insert(hash, DeltaReaderNode::Internal { left, right });
79 }
80 for (hash, (key, value)) in leaf_nodes {
81 nodes.insert(hash, DeltaReaderNode::Leaf { key, value });
82 }
83
84 Ok(Self { nodes })
85 }
86
87 pub fn get_missing_hashes(&self, root_hash: Hash) -> HashSet<Hash> {
88 let mut missing_hashes: HashSet<Hash> = HashSet::new();
89
90 for node in self.nodes.values() {
91 let DeltaReaderNode::Internal { left, right } = node else {
92 continue;
93 };
94
95 for hash in [left, right] {
96 if !self.nodes.contains_key(hash) {
97 missing_hashes.insert(*hash);
98 }
99 }
100 }
101
102 if !self.nodes.contains_key(&root_hash) {
103 missing_hashes.insert(root_hash);
104 }
105
106 missing_hashes
107 }
108
109 pub fn collect_from_merkle_blob(
110 &mut self,
111 path: &PathBuf,
112 indexes: &Vec<TreeIndex>,
113 ) -> Result<(), Error> {
114 let vector = crate::zstd_decode_path(path)?;
115
116 for (hash, (_index, node)) in crate::get_internal_terminal(&vector, indexes)? {
117 self.nodes.insert(hash, node);
118 }
119
120 Ok(())
121 }
122
123 pub fn collect_and_return_from_merkle_blobs(
124 &mut self,
125 jobs: &Vec<(Hash, PathBuf)>,
126 hashes: &HashSet<Hash>,
127 ) -> Result<Vec<(Hash, NodeHashToIndex)>, Error> {
128 let mut grouped_results = Vec::new();
129 grouped_results.par_extend(jobs.into_par_iter().map(
130 |(hash, path)| -> Result<(Hash, (NodeHashToDeltaReaderNode, NodeHashToIndex)), Error> {
131 Ok((
132 *hash,
133 crate::collect_and_return_from_merkle_blob(path, hashes, |key| {
134 self.nodes.contains_key(key)
135 })?,
136 ))
137 },
138 ));
139
140 let mut results: Vec<(Hash, NodeHashToIndex)> = Vec::new();
141 let mut seen_hashes: HashSet<Hash> = HashSet::new();
142 for result in grouped_results {
143 let (hash, (nodes, node_hash_to_index)) = result?;
144 self.nodes.extend(nodes);
145 let mut filtered = HashMap::new();
146 for (hash, index) in node_hash_to_index {
147 if seen_hashes.insert(hash) {
148 filtered.insert(hash, index);
149 }
150 }
151 results.push((hash, filtered));
152 }
153
154 Ok(results)
155 }
156
157 pub fn collect_from_merkle_blobs(
158 &mut self,
159 jobs: &Vec<(PathBuf, Vec<TreeIndex>)>,
160 ) -> Result<(), Error> {
161 let mut results = Vec::new();
162
163 results.par_extend(jobs.into_par_iter().map(
164 |(path, indexes)| -> Result<HashMap<Hash, (TreeIndex, DeltaReaderNode)>, Error> {
165 let vector = crate::zstd_decode_path(path)?;
166 crate::get_internal_terminal(&vector, indexes)
167 },
168 ));
169
170 for result in results {
171 for (hash, (_index, node)) in result? {
173 self.nodes.insert(hash, node);
174 }
175 }
176
177 Ok(())
178 }
179
180 pub fn create_merkle_blob_and_filter_unused_nodes(
181 &mut self,
182 root_hash: Hash,
183 interested_hashes: &HashSet<Hash>,
184 ) -> Result<MerkleBlob, Error> {
185 let mut all_used_hashes: HashSet<Hash> = HashSet::new();
186 let merkle_blob = MerkleBlob::build_blob_from_node_list(
187 &self.nodes,
188 root_hash,
189 interested_hashes,
190 &mut all_used_hashes,
191 )?;
192
193 self.nodes.retain(|k, _v| all_used_hashes.contains(k));
194
195 Ok(merkle_blob)
196 }
197}
198
199#[cfg(feature = "py-bindings")]
200#[pymethods]
201impl DeltaReader {
202 #[new]
203 pub fn py_init(internal_nodes: InternalNodesMap, leaf_nodes: LeafNodesMap) -> PyResult<Self> {
204 Ok(Self::new(internal_nodes, leaf_nodes)?)
205 }
206
207 #[pyo3(name = "get_missing_hashes")]
208 pub fn py_get_missing_hashes(&self, root_hash: Hash) -> PyResult<HashSet<Hash>> {
209 Ok(self.get_missing_hashes(root_hash))
210 }
211
212 #[pyo3(name = "add_internal_nodes")]
213 pub fn py_add_internal_nodes(&mut self, internal_nodes: InternalNodesMap) {
214 for (hash, (left, right)) in internal_nodes {
215 self.nodes
216 .insert(hash, DeltaReaderNode::Internal { left, right });
217 }
218 }
219
220 #[pyo3(name = "add_leaf_nodes")]
221 pub fn py_add_leaf_nodes(&mut self, leaf_nodes: LeafNodesMap) {
222 for (hash, (key, value)) in leaf_nodes {
223 self.nodes
224 .insert(hash, DeltaReaderNode::Leaf { key, value });
225 }
226 }
227
228 #[allow(clippy::needless_pass_by_value)]
229 #[pyo3(name = "collect_from_merkle_blob")]
230 pub fn py_collect_from_merkle_blob(
231 &mut self,
232 path: PathBuf,
233 indexes: Vec<TreeIndex>,
234 ) -> PyResult<()> {
235 self.collect_from_merkle_blob(&path, &indexes)?;
236
237 Ok(())
238 }
239
240 #[allow(clippy::needless_pass_by_value)]
241 #[pyo3(name = "collect_and_return_from_merkle_blobs")]
242 pub fn py_collect_and_return_from_merkle_blobs(
243 &mut self,
244 py: Python<'_>,
245 jobs: Vec<(Hash, PathBuf)>,
246 hashes: HashSet<Hash>,
247 ) -> PyResult<Vec<(Hash, NodeHashToIndex)>> {
248 let mut extracted_jobs: Vec<(Hash, PathBuf)> = Vec::new();
249 for (hash, path) in jobs {
250 extracted_jobs.push((hash, path));
251 }
252
253 Ok(py.detach(|| self.collect_and_return_from_merkle_blobs(&extracted_jobs, &hashes))?)
254 }
255
256 #[allow(clippy::needless_pass_by_value)]
257 #[pyo3(name = "collect_from_merkle_blobs")]
258 pub fn py_collect_from_merkle_blobs(
259 &mut self,
260 py: Python<'_>,
261 jobs: Vec<(PathBuf, Vec<TreeIndex>)>,
262 ) -> PyResult<()> {
263 let mut pathed_jobs: Vec<(PathBuf, Vec<TreeIndex>)> = Vec::new();
264 for (path, indexes) in jobs {
265 pathed_jobs.push((path, indexes));
266 }
267 py.detach(|| self.collect_from_merkle_blobs(&pathed_jobs))?;
268
269 Ok(())
270 }
271
272 #[allow(clippy::needless_pass_by_value)]
273 #[pyo3(name = "create_merkle_blob_and_filter_unused_nodes")]
274 pub fn py_create_merkle_blob_and_filter_unused_nodes(
275 &mut self,
276 root_hash: Hash,
277 interested_hashes: HashSet<Hash>,
278 ) -> Result<MerkleBlob, Error> {
279 self.create_merkle_blob_and_filter_unused_nodes(root_hash, &interested_hashes)
280 }
281}
282
283#[cfg(feature = "py-bindings")]
284#[pymethods]
285impl DeltaFileCache {
286 #[allow(clippy::needless_pass_by_value)]
287 #[new]
288 fn py_new(path: PathBuf) -> PyResult<Self> {
289 Ok(Self::new(&path)?)
290 }
291
292 #[allow(clippy::needless_pass_by_value)]
293 #[pyo3(name = "load_previous_hashes")]
294 pub fn py_load_previous_hashes(&mut self, path: PathBuf) -> PyResult<()> {
295 Ok(self.load_previous_hashes(&path)?)
296 }
297
298 #[pyo3(name = "get_index")]
299 pub fn py_get_index(&self, hash: Hash) -> PyResult<TreeIndex> {
300 Ok(self.get_index(hash)?)
301 }
302
303 #[pyo3(name = "seen_previous_hash")]
304 pub fn py_seen_previous_hash(&self, hash: Hash) -> bool {
305 self.seen_previous_hash(hash)
306 }
307
308 #[pyo3(name = "get_raw_node")]
309 pub fn py_get_raw_node(&mut self, index: TreeIndex) -> PyResult<Node> {
310 Ok(self.get_raw_node(index)?)
311 }
312
313 #[pyo3(name = "get_hash_at_index")]
314 pub fn py_get_hash_at_index(&self, index: TreeIndex) -> PyResult<Option<Hash>> {
315 Ok(self.get_hash_at_index(index)?)
316 }
317}
318
319#[cfg(test)]
320mod tests {
321 use super::*;
322 use crate::merkle::test_util::traversal_blob;
323 use crate::merkle::test_util::{HASH_ONE, HASH_TWO, HASH_ZERO, generate_hash, generate_kvid};
324 use crate::{InsertLocation, InternalNodesMap, LeafNodesMap};
325 use expect_test::expect;
326 use rstest::rstest;
327 use std::iter::zip;
328 use tempfile;
329
330 fn incomplete_delta_reader() -> DeltaReader {
331 let mut internal_nodes_map = InternalNodesMap::new();
332 let mut leaf_nodes_map = LeafNodesMap::new();
333
334 internal_nodes_map.insert(HASH_ZERO, (HASH_ONE, HASH_TWO));
335 leaf_nodes_map.insert(HASH_ONE, (KeyId(0), ValueId(1)));
336
337 DeltaReader::new(internal_nodes_map, leaf_nodes_map).unwrap()
338 }
339
340 fn complete_delta_reader() -> DeltaReader {
341 let mut delta_reader = incomplete_delta_reader();
342 delta_reader.nodes.insert(
343 HASH_TWO,
344 DeltaReaderNode::Leaf {
345 key: KeyId(2),
346 value: ValueId(3),
347 },
348 );
349
350 delta_reader
351 }
352
353 #[test]
354 fn test_root_hash_missing() {
355 let internal_nodes_map = InternalNodesMap::new();
356 let leaf_nodes_map = LeafNodesMap::new();
357 let delta_reader = DeltaReader::new(internal_nodes_map, leaf_nodes_map).unwrap();
358 let missing = delta_reader.get_missing_hashes(HASH_ZERO);
359 let expected = expect![[r"
360 {
361 Hash(
362 0000000000000000000000000000000000000000000000000000000000000000,
363 ),
364 }
365 "]];
366 expected.assert_debug_eq(&missing);
367 }
368
369 #[test]
370 fn test_delta_reader_get_missing_hashes_one_known_one_unknown() {
371 let delta_reader = incomplete_delta_reader();
372 let missing = delta_reader.get_missing_hashes(HASH_ZERO);
373
374 let expected = expect![[r"
375 {
376 Hash(
377 0202020202020202020202020202020202020202020202020202020202020202,
378 ),
379 }
380 "]];
381 expected.assert_debug_eq(&missing);
382 }
383
384 #[test]
385 fn test_delta_reader_collect_from_merkle_blob_completes_incomplete() {
386 let mut delta_reader = incomplete_delta_reader();
387 let mut merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
388 merkle_blob
389 .insert(KeyId(2), ValueId(3), &HASH_TWO, InsertLocation::AsRoot {})
390 .unwrap();
391
392 let dir_path = tempfile::tempdir().unwrap();
393 let leaf_blob_path = dir_path.path().join("merkle_blob");
394 merkle_blob.to_path(&leaf_blob_path).unwrap();
395 delta_reader
396 .collect_from_merkle_blob(&leaf_blob_path, &vec![TreeIndex(0)])
397 .unwrap();
398
399 let missing = delta_reader.get_missing_hashes(HASH_ZERO);
400
401 #[allow(clippy::needless_raw_string_hashes)]
402 let expected = expect![[r#"
403 {}
404 "#]];
405 expected.assert_debug_eq(&missing);
406 }
407
408 #[rstest]
409 fn test_delta_reader_collect_from_merkle_blob_is_complete(traversal_blob: MerkleBlob) {
410 let mut delta_reader = DeltaReader {
411 nodes: HashMap::new(),
412 };
413 let dir_path = tempfile::tempdir().unwrap();
414 let blob_path = dir_path.path().join("merkle_blob");
415 traversal_blob.to_path(&blob_path).unwrap();
416 delta_reader
417 .collect_from_merkle_blob(&blob_path, &vec![TreeIndex(0)])
418 .unwrap();
419
420 let missing = delta_reader.get_missing_hashes(
421 traversal_blob
422 .get_hash_at_index(TreeIndex(0))
423 .unwrap()
424 .expect("Expected root hash"),
425 );
426
427 #[allow(clippy::needless_raw_string_hashes)]
428 let expected = expect![[r#"
429 {}
430 "#]];
431 expected.assert_debug_eq(&missing);
432 }
433
434 #[test]
435 fn test_delta_reader_collect_from_merkle_blobs() {
436 let mut delta_reader = incomplete_delta_reader();
437 let mut merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
438 merkle_blob
439 .insert(KeyId(2), ValueId(3), &HASH_TWO, InsertLocation::AsRoot {})
440 .unwrap();
441
442 let dir_path = tempfile::tempdir().unwrap();
443 let leaf_blob_path = dir_path.path().join("merkle_blob");
444 merkle_blob.to_path(&leaf_blob_path).unwrap();
445 delta_reader
446 .collect_from_merkle_blobs(&vec![(leaf_blob_path, vec![TreeIndex(0)])])
447 .unwrap();
448
449 let missing = delta_reader.get_missing_hashes(
450 merkle_blob
451 .get_hash_at_index(TreeIndex(0))
452 .unwrap()
453 .expect("Expected root hash"),
454 );
455
456 #[allow(clippy::needless_raw_string_hashes)]
457 let expected = expect![[r#"
458 {}
459 "#]];
460 expected.assert_debug_eq(&missing);
461 }
462
463 #[test]
464 #[should_panic(expected = "integrity check failed while dropping merkle blob: CycleFound")]
465 fn test_delta_reader_create_merkle_blob_incomplete_fails() {
466 let mut delta_reader = incomplete_delta_reader();
467
468 delta_reader
469 .create_merkle_blob_and_filter_unused_nodes(HASH_ZERO, &HashSet::new())
470 .expect_err("incomplete so should fail");
471 }
472
473 #[test]
474 fn test_delta_reader_create_merkle_blob_works() {
475 let mut delta_reader = complete_delta_reader();
476 let interested_hashes: HashSet<Hash> = delta_reader.nodes.keys().copied().collect();
477
478 let complete_blob = delta_reader
479 .create_merkle_blob_and_filter_unused_nodes(HASH_ZERO, &interested_hashes)
480 .unwrap();
481 complete_blob.check_integrity().unwrap();
482 }
483
484 #[rstest]
485 fn test_collect_and_return(traversal_blob: MerkleBlob) {
486 let dir_path = tempfile::tempdir().unwrap();
487 let file_path = dir_path.path().join("blob");
488 traversal_blob.to_path(&file_path).unwrap();
489 let hashes = traversal_blob
490 .get_hashes_indexes(false)
491 .unwrap()
492 .into_keys()
493 .collect::<HashSet<Hash>>();
494 let root_hash = traversal_blob.get_hash(TreeIndex(0)).unwrap();
495
496 let mut delta_reader = DeltaReader {
497 nodes: HashMap::new(),
498 };
499 let mut root_hash_to_node_hash_to_index = delta_reader
500 .collect_and_return_from_merkle_blobs(&vec![(root_hash, file_path)], &hashes)
501 .unwrap();
502
503 let (collected_root_hash, collected_node_hash_to_index) =
504 root_hash_to_node_hash_to_index.pop().unwrap();
505 assert_eq!(root_hash_to_node_hash_to_index.len(), 0);
506
507 assert_eq!(collected_root_hash, root_hash);
508 assert_eq!(
509 collected_node_hash_to_index,
510 traversal_blob.get_hashes_indexes(false).unwrap()
511 );
512 }
513
514 #[rstest]
515 fn test_delta_file_cache() {
516 let num_inserts = 500;
517
518 let mut merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
519 let mut kv_ids: Vec<(KeyId, ValueId)> = Vec::new();
520 let mut hashes: Vec<Hash> = Vec::new();
521
522 let mut previous_merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
523 let mut prev_kv_ids: Vec<(KeyId, ValueId)> = Vec::new();
524 let mut prev_hashes: Vec<Hash> = Vec::new();
525
526 for seed in 1..=num_inserts {
527 let (key, value) = generate_kvid(seed);
528 kv_ids.push((key, value));
529 hashes.push(generate_hash(seed));
530
531 let (key, value) = generate_kvid(num_inserts + seed);
532 prev_kv_ids.push((key, value));
533 prev_hashes.push(generate_hash(num_inserts + seed));
534 }
535
536 merkle_blob
537 .batch_insert(zip(kv_ids, hashes.clone()).collect())
538 .unwrap();
539 merkle_blob.calculate_lazy_hashes().unwrap();
540
541 previous_merkle_blob
542 .batch_insert(zip(prev_kv_ids, prev_hashes.clone()).collect())
543 .unwrap();
544 previous_merkle_blob.calculate_lazy_hashes().unwrap();
545
546 let dir_path = tempfile::tempdir().unwrap();
547 let blob_path = dir_path.path().join("merkle_blob");
548 merkle_blob.to_path(&blob_path).unwrap();
549 let previous_blob_path = dir_path.path().join("previous_merkle_blob");
550 previous_merkle_blob.to_path(&previous_blob_path).unwrap();
551
552 let mut delta_cache_file = DeltaFileCache::new(&blob_path).unwrap();
553 for hash in &hashes {
554 let index = delta_cache_file.get_index(*hash).unwrap();
555 let received_hash = delta_cache_file.get_hash_at_index(index).unwrap();
556 assert_eq!(received_hash, Some(*hash));
557 let node = delta_cache_file.get_raw_node(index).unwrap();
558 assert_eq!(node.hash(), *hash);
559 }
560
561 delta_cache_file
562 .load_previous_hashes(&previous_blob_path)
563 .unwrap();
564 for hash in &prev_hashes {
565 let exists = delta_cache_file.seen_previous_hash(*hash);
566 assert!(exists);
567 }
568 }
569}