use crate::merkle::error::Error;
use crate::{
Hash, InternalNodesMap, KeyId, LeafNodesMap, MerkleBlob, Node, NodeHashToDeltaReaderNode,
NodeHashToIndex, ParentFirstIterator, TreeIndex, ValueId,
};
#[cfg(feature = "py-bindings")]
use pyo3::{PyResult, Python, pyclass, pymethods};
use rayon::iter::{IntoParallelIterator, ParallelExtend, ParallelIterator};
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
pub enum DeltaReaderNode {
Internal { left: Hash, right: Hash },
Leaf { key: KeyId, value: ValueId },
}
#[cfg_attr(feature = "py-bindings", pyclass)]
pub struct DeltaFileCache {
hash_to_index: NodeHashToIndex,
previous_hashes: HashSet<Hash>,
merkle_blob: MerkleBlob,
}
impl DeltaFileCache {
pub fn new(path: &PathBuf) -> Result<Self, Error> {
let merkle_blob = MerkleBlob::from_path(path)?;
let hash_to_index = merkle_blob.get_hashes_indexes(false)?;
Ok(Self {
hash_to_index,
previous_hashes: HashSet::new(),
merkle_blob,
})
}
pub fn load_previous_hashes(&mut self, path: &PathBuf) -> Result<(), Error> {
let blob = crate::zstd_decode_path(path)?;
self.previous_hashes = HashSet::new();
if !blob.is_empty() {
for item in ParentFirstIterator::new(&blob, None) {
let (_, block) = item?;
self.previous_hashes.insert(block.node.hash());
}
}
Ok(())
}
pub fn get_raw_node(&self, index: TreeIndex) -> Result<Node, Error> {
self.merkle_blob.get_node(index)
}
pub fn get_hash_at_index(&self, index: TreeIndex) -> Result<Option<Hash>, Error> {
self.merkle_blob.get_hash_at_index(index)
}
pub fn seen_previous_hash(&self, hash: Hash) -> bool {
self.previous_hashes.contains(&hash)
}
pub fn get_index(&self, hash: Hash) -> Result<TreeIndex, Error> {
self.hash_to_index
.get(&hash)
.copied()
.ok_or(Error::HashNotFound(hash))
}
}
#[cfg_attr(feature = "py-bindings", pyclass)]
pub struct DeltaReader {
nodes: NodeHashToDeltaReaderNode,
}
impl DeltaReader {
pub fn new(internal_nodes: InternalNodesMap, leaf_nodes: LeafNodesMap) -> Result<Self, Error> {
let mut nodes = NodeHashToDeltaReaderNode::new();
for (hash, (left, right)) in internal_nodes {
nodes.insert(hash, DeltaReaderNode::Internal { left, right });
}
for (hash, (key, value)) in leaf_nodes {
nodes.insert(hash, DeltaReaderNode::Leaf { key, value });
}
Ok(Self { nodes })
}
pub fn get_missing_hashes(&self, root_hash: Hash) -> HashSet<Hash> {
let mut missing_hashes: HashSet<Hash> = HashSet::new();
for node in self.nodes.values() {
let DeltaReaderNode::Internal { left, right } = node else {
continue;
};
for hash in [left, right] {
if !self.nodes.contains_key(hash) {
missing_hashes.insert(*hash);
}
}
}
if !self.nodes.contains_key(&root_hash) {
missing_hashes.insert(root_hash);
}
missing_hashes
}
pub fn collect_from_merkle_blob(
&mut self,
path: &PathBuf,
indexes: &Vec<TreeIndex>,
) -> Result<(), Error> {
let vector = crate::zstd_decode_path(path)?;
for (hash, (_index, node)) in crate::get_internal_terminal(&vector, indexes)? {
self.nodes.insert(hash, node);
}
Ok(())
}
pub fn collect_and_return_from_merkle_blobs(
&mut self,
jobs: &Vec<(Hash, PathBuf)>,
hashes: &HashSet<Hash>,
) -> Result<Vec<(Hash, NodeHashToIndex)>, Error> {
let mut grouped_results = Vec::new();
grouped_results.par_extend(jobs.into_par_iter().map(
|(hash, path)| -> Result<(Hash, (NodeHashToDeltaReaderNode, NodeHashToIndex)), Error> {
Ok((
*hash,
crate::collect_and_return_from_merkle_blob(path, hashes, |key| {
self.nodes.contains_key(key)
})?,
))
},
));
let mut results: Vec<(Hash, NodeHashToIndex)> = Vec::new();
let mut seen_hashes: HashSet<Hash> = HashSet::new();
for result in grouped_results {
let (hash, (nodes, node_hash_to_index)) = result?;
self.nodes.extend(nodes);
let mut filtered = HashMap::new();
for (hash, index) in node_hash_to_index {
if seen_hashes.insert(hash) {
filtered.insert(hash, index);
}
}
results.push((hash, filtered));
}
Ok(results)
}
pub fn collect_from_merkle_blobs(
&mut self,
jobs: &Vec<(PathBuf, Vec<TreeIndex>)>,
) -> Result<(), Error> {
let mut results = Vec::new();
results.par_extend(jobs.into_par_iter().map(
|(path, indexes)| -> Result<HashMap<Hash, (TreeIndex, DeltaReaderNode)>, Error> {
let vector = crate::zstd_decode_path(path)?;
crate::get_internal_terminal(&vector, indexes)
},
));
for result in results {
for (hash, (_index, node)) in result? {
self.nodes.insert(hash, node);
}
}
Ok(())
}
pub fn create_merkle_blob_and_filter_unused_nodes(
&mut self,
root_hash: Hash,
interested_hashes: &HashSet<Hash>,
) -> Result<MerkleBlob, Error> {
let mut all_used_hashes: HashSet<Hash> = HashSet::new();
let merkle_blob = MerkleBlob::build_blob_from_node_list(
&self.nodes,
root_hash,
interested_hashes,
&mut all_used_hashes,
)?;
self.nodes.retain(|k, _v| all_used_hashes.contains(k));
Ok(merkle_blob)
}
}
#[cfg(feature = "py-bindings")]
#[pymethods]
impl DeltaReader {
#[new]
pub fn py_init(internal_nodes: InternalNodesMap, leaf_nodes: LeafNodesMap) -> PyResult<Self> {
Ok(Self::new(internal_nodes, leaf_nodes)?)
}
#[pyo3(name = "get_missing_hashes")]
pub fn py_get_missing_hashes(&self, root_hash: Hash) -> PyResult<HashSet<Hash>> {
Ok(self.get_missing_hashes(root_hash))
}
#[pyo3(name = "add_internal_nodes")]
pub fn py_add_internal_nodes(&mut self, internal_nodes: InternalNodesMap) {
for (hash, (left, right)) in internal_nodes {
self.nodes
.insert(hash, DeltaReaderNode::Internal { left, right });
}
}
#[pyo3(name = "add_leaf_nodes")]
pub fn py_add_leaf_nodes(&mut self, leaf_nodes: LeafNodesMap) {
for (hash, (key, value)) in leaf_nodes {
self.nodes
.insert(hash, DeltaReaderNode::Leaf { key, value });
}
}
#[allow(clippy::needless_pass_by_value)]
#[pyo3(name = "collect_from_merkle_blob")]
pub fn py_collect_from_merkle_blob(
&mut self,
path: PathBuf,
indexes: Vec<TreeIndex>,
) -> PyResult<()> {
self.collect_from_merkle_blob(&path, &indexes)?;
Ok(())
}
#[allow(clippy::needless_pass_by_value)]
#[pyo3(name = "collect_and_return_from_merkle_blobs")]
pub fn py_collect_and_return_from_merkle_blobs(
&mut self,
py: Python<'_>,
jobs: Vec<(Hash, PathBuf)>,
hashes: HashSet<Hash>,
) -> PyResult<Vec<(Hash, NodeHashToIndex)>> {
let mut extracted_jobs: Vec<(Hash, PathBuf)> = Vec::new();
for (hash, path) in jobs {
extracted_jobs.push((hash, path));
}
Ok(py.detach(|| self.collect_and_return_from_merkle_blobs(&extracted_jobs, &hashes))?)
}
#[allow(clippy::needless_pass_by_value)]
#[pyo3(name = "collect_from_merkle_blobs")]
pub fn py_collect_from_merkle_blobs(
&mut self,
py: Python<'_>,
jobs: Vec<(PathBuf, Vec<TreeIndex>)>,
) -> PyResult<()> {
let mut pathed_jobs: Vec<(PathBuf, Vec<TreeIndex>)> = Vec::new();
for (path, indexes) in jobs {
pathed_jobs.push((path, indexes));
}
py.detach(|| self.collect_from_merkle_blobs(&pathed_jobs))?;
Ok(())
}
#[allow(clippy::needless_pass_by_value)]
#[pyo3(name = "create_merkle_blob_and_filter_unused_nodes")]
pub fn py_create_merkle_blob_and_filter_unused_nodes(
&mut self,
root_hash: Hash,
interested_hashes: HashSet<Hash>,
) -> Result<MerkleBlob, Error> {
self.create_merkle_blob_and_filter_unused_nodes(root_hash, &interested_hashes)
}
}
#[cfg(feature = "py-bindings")]
#[pymethods]
impl DeltaFileCache {
#[allow(clippy::needless_pass_by_value)]
#[new]
fn py_new(path: PathBuf) -> PyResult<Self> {
Ok(Self::new(&path)?)
}
#[allow(clippy::needless_pass_by_value)]
#[pyo3(name = "load_previous_hashes")]
pub fn py_load_previous_hashes(&mut self, path: PathBuf) -> PyResult<()> {
Ok(self.load_previous_hashes(&path)?)
}
#[pyo3(name = "get_index")]
pub fn py_get_index(&self, hash: Hash) -> PyResult<TreeIndex> {
Ok(self.get_index(hash)?)
}
#[pyo3(name = "seen_previous_hash")]
pub fn py_seen_previous_hash(&self, hash: Hash) -> bool {
self.seen_previous_hash(hash)
}
#[pyo3(name = "get_raw_node")]
pub fn py_get_raw_node(&mut self, index: TreeIndex) -> PyResult<Node> {
Ok(self.get_raw_node(index)?)
}
#[pyo3(name = "get_hash_at_index")]
pub fn py_get_hash_at_index(&self, index: TreeIndex) -> PyResult<Option<Hash>> {
Ok(self.get_hash_at_index(index)?)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::merkle::test_util::traversal_blob;
use crate::merkle::test_util::{HASH_ONE, HASH_TWO, HASH_ZERO, generate_hash, generate_kvid};
use crate::{InsertLocation, InternalNodesMap, LeafNodesMap};
use expect_test::expect;
use rstest::rstest;
use std::iter::zip;
use tempfile;
fn incomplete_delta_reader() -> DeltaReader {
let mut internal_nodes_map = InternalNodesMap::new();
let mut leaf_nodes_map = LeafNodesMap::new();
internal_nodes_map.insert(HASH_ZERO, (HASH_ONE, HASH_TWO));
leaf_nodes_map.insert(HASH_ONE, (KeyId(0), ValueId(1)));
DeltaReader::new(internal_nodes_map, leaf_nodes_map).unwrap()
}
fn complete_delta_reader() -> DeltaReader {
let mut delta_reader = incomplete_delta_reader();
delta_reader.nodes.insert(
HASH_TWO,
DeltaReaderNode::Leaf {
key: KeyId(2),
value: ValueId(3),
},
);
delta_reader
}
#[test]
fn test_root_hash_missing() {
let internal_nodes_map = InternalNodesMap::new();
let leaf_nodes_map = LeafNodesMap::new();
let delta_reader = DeltaReader::new(internal_nodes_map, leaf_nodes_map).unwrap();
let missing = delta_reader.get_missing_hashes(HASH_ZERO);
let expected = expect![[r"
{
Hash(
0000000000000000000000000000000000000000000000000000000000000000,
),
}
"]];
expected.assert_debug_eq(&missing);
}
#[test]
fn test_delta_reader_get_missing_hashes_one_known_one_unknown() {
let delta_reader = incomplete_delta_reader();
let missing = delta_reader.get_missing_hashes(HASH_ZERO);
let expected = expect![[r"
{
Hash(
0202020202020202020202020202020202020202020202020202020202020202,
),
}
"]];
expected.assert_debug_eq(&missing);
}
#[test]
fn test_delta_reader_collect_from_merkle_blob_completes_incomplete() {
let mut delta_reader = incomplete_delta_reader();
let mut merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
merkle_blob
.insert(KeyId(2), ValueId(3), &HASH_TWO, InsertLocation::AsRoot {})
.unwrap();
let dir_path = tempfile::tempdir().unwrap();
let leaf_blob_path = dir_path.path().join("merkle_blob");
merkle_blob.to_path(&leaf_blob_path).unwrap();
delta_reader
.collect_from_merkle_blob(&leaf_blob_path, &vec![TreeIndex(0)])
.unwrap();
let missing = delta_reader.get_missing_hashes(HASH_ZERO);
#[allow(clippy::needless_raw_string_hashes)]
let expected = expect![[r#"
{}
"#]];
expected.assert_debug_eq(&missing);
}
#[rstest]
fn test_delta_reader_collect_from_merkle_blob_is_complete(traversal_blob: MerkleBlob) {
let mut delta_reader = DeltaReader {
nodes: HashMap::new(),
};
let dir_path = tempfile::tempdir().unwrap();
let blob_path = dir_path.path().join("merkle_blob");
traversal_blob.to_path(&blob_path).unwrap();
delta_reader
.collect_from_merkle_blob(&blob_path, &vec![TreeIndex(0)])
.unwrap();
let missing = delta_reader.get_missing_hashes(
traversal_blob
.get_hash_at_index(TreeIndex(0))
.unwrap()
.expect("Expected root hash"),
);
#[allow(clippy::needless_raw_string_hashes)]
let expected = expect![[r#"
{}
"#]];
expected.assert_debug_eq(&missing);
}
#[test]
fn test_delta_reader_collect_from_merkle_blobs() {
let mut delta_reader = incomplete_delta_reader();
let mut merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
merkle_blob
.insert(KeyId(2), ValueId(3), &HASH_TWO, InsertLocation::AsRoot {})
.unwrap();
let dir_path = tempfile::tempdir().unwrap();
let leaf_blob_path = dir_path.path().join("merkle_blob");
merkle_blob.to_path(&leaf_blob_path).unwrap();
delta_reader
.collect_from_merkle_blobs(&vec![(leaf_blob_path, vec![TreeIndex(0)])])
.unwrap();
let missing = delta_reader.get_missing_hashes(
merkle_blob
.get_hash_at_index(TreeIndex(0))
.unwrap()
.expect("Expected root hash"),
);
#[allow(clippy::needless_raw_string_hashes)]
let expected = expect![[r#"
{}
"#]];
expected.assert_debug_eq(&missing);
}
#[test]
#[should_panic(expected = "integrity check failed while dropping merkle blob: CycleFound")]
fn test_delta_reader_create_merkle_blob_incomplete_fails() {
let mut delta_reader = incomplete_delta_reader();
delta_reader
.create_merkle_blob_and_filter_unused_nodes(HASH_ZERO, &HashSet::new())
.expect_err("incomplete so should fail");
}
#[test]
fn test_delta_reader_create_merkle_blob_works() {
let mut delta_reader = complete_delta_reader();
let interested_hashes: HashSet<Hash> = delta_reader.nodes.keys().copied().collect();
let complete_blob = delta_reader
.create_merkle_blob_and_filter_unused_nodes(HASH_ZERO, &interested_hashes)
.unwrap();
complete_blob.check_integrity().unwrap();
}
#[rstest]
fn test_collect_and_return(traversal_blob: MerkleBlob) {
let dir_path = tempfile::tempdir().unwrap();
let file_path = dir_path.path().join("blob");
traversal_blob.to_path(&file_path).unwrap();
let hashes = traversal_blob
.get_hashes_indexes(false)
.unwrap()
.into_keys()
.collect::<HashSet<Hash>>();
let root_hash = traversal_blob.get_hash(TreeIndex(0)).unwrap();
let mut delta_reader = DeltaReader {
nodes: HashMap::new(),
};
let mut root_hash_to_node_hash_to_index = delta_reader
.collect_and_return_from_merkle_blobs(&vec![(root_hash, file_path)], &hashes)
.unwrap();
let (collected_root_hash, collected_node_hash_to_index) =
root_hash_to_node_hash_to_index.pop().unwrap();
assert_eq!(root_hash_to_node_hash_to_index.len(), 0);
assert_eq!(collected_root_hash, root_hash);
assert_eq!(
collected_node_hash_to_index,
traversal_blob.get_hashes_indexes(false).unwrap()
);
}
#[rstest]
fn test_delta_file_cache() {
let num_inserts = 500;
let mut merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
let mut kv_ids: Vec<(KeyId, ValueId)> = Vec::new();
let mut hashes: Vec<Hash> = Vec::new();
let mut previous_merkle_blob = MerkleBlob::new(Vec::new()).unwrap();
let mut prev_kv_ids: Vec<(KeyId, ValueId)> = Vec::new();
let mut prev_hashes: Vec<Hash> = Vec::new();
for seed in 1..=num_inserts {
let (key, value) = generate_kvid(seed);
kv_ids.push((key, value));
hashes.push(generate_hash(seed));
let (key, value) = generate_kvid(num_inserts + seed);
prev_kv_ids.push((key, value));
prev_hashes.push(generate_hash(num_inserts + seed));
}
merkle_blob
.batch_insert(zip(kv_ids, hashes.clone()).collect())
.unwrap();
merkle_blob.calculate_lazy_hashes().unwrap();
previous_merkle_blob
.batch_insert(zip(prev_kv_ids, prev_hashes.clone()).collect())
.unwrap();
previous_merkle_blob.calculate_lazy_hashes().unwrap();
let dir_path = tempfile::tempdir().unwrap();
let blob_path = dir_path.path().join("merkle_blob");
merkle_blob.to_path(&blob_path).unwrap();
let previous_blob_path = dir_path.path().join("previous_merkle_blob");
previous_merkle_blob.to_path(&previous_blob_path).unwrap();
let mut delta_cache_file = DeltaFileCache::new(&blob_path).unwrap();
for hash in &hashes {
let index = delta_cache_file.get_index(*hash).unwrap();
let received_hash = delta_cache_file.get_hash_at_index(index).unwrap();
assert_eq!(received_hash, Some(*hash));
let node = delta_cache_file.get_raw_node(index).unwrap();
assert_eq!(node.hash(), *hash);
}
delta_cache_file
.load_previous_hashes(&previous_blob_path)
.unwrap();
for hash in &prev_hashes {
let exists = delta_cache_file.seen_previous_hash(*hash);
assert!(exists);
}
}
}