Skip to content

Commit

Permalink
jmt: prototype for append_value_sets
Browse files Browse the repository at this point in the history
  • Loading branch information
erwanor committed Mar 7, 2024
1 parent 041ad5c commit 501bb14
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 0 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ default = ["ics23", "std", "sha2"]
mocks = ["dep:parking_lot"]
blake3_tests = ["dep:blake3"]
std = ["dep:thiserror"]
migration = []

[dependencies]
anyhow = "1.0.38"
Expand Down
30 changes: 30 additions & 0 deletions src/tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -445,6 +445,36 @@ where
Ok(tree_cache.into())
}

#[cfg(feature = "migration")]
/// Append value sets to the latest version of the tree.
pub fn append_value_sets(
&self,
value_sets: impl IntoIterator<Item = impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>>,
latest_version: Version,
) -> Result<(Vec<RootHash>, TreeUpdateBatch)> {
let mut tree_cache = TreeCache::new_overwrite(self.reader, latest_version)?;
for (idx, value_set) in value_sets.into_iter().enumerate() {
let version = latest_version + idx as u64;
for (i, (key, value)) in value_set.into_iter().enumerate() {
let action = if value.is_some() { "insert" } else { "delete" };
let value_hash = value.as_ref().map(|v| ValueHash::with::<H>(v));
tree_cache.put_value(version, key, value);
self.put(key, value_hash, version, &mut tree_cache, false)
.with_context(|| {
format!(
"failed to {} key {} for version {}, key = {:?}",
action, i, version, key
)
})?;
}

// Freezes the current cache to make all contents in the current cache immutable.
tree_cache.freeze::<H>()?;
}

Ok(tree_cache.into())
}

/// Same as [`put_value_sets`], this method returns a Merkle proof for every update of the Merkle tree.
/// The proofs can be verified using the [`verify_update`] method, which requires the old `root_hash`, the `merkle_proof` and the new `root_hash`
/// The first argument contains all the root hashes that were stored in the tree cache so far. The last one is the new root hash of the tree.
Expand Down
21 changes: 21 additions & 0 deletions src/tree_cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,27 @@ where
})
}

#[cfg(feature = "migration")]
pub fn new_overwrite(
reader: &'a R,
current_version: Version,
root_node_key: NodeKey,
) -> Result<Self> {
let mut node_cache = HashMap::new();
let root_node_key = NodeKey::new_empty_path(current_version);
Ok(Self {
node_cache,
stale_node_index_cache: HashSet::new(),
frozen_cache: FrozenTreeCache::new(),
root_node_key,
next_version: current_version + 1,
reader,
num_stale_leaves: 0,
num_new_leaves: 0,
value_cache: Default::default(),
})
}

/// Gets a node with given node key. If it doesn't exist in node cache, read from `reader`.
pub fn get_node(&self, node_key: &NodeKey) -> Result<Node> {
Ok(if let Some(node) = self.node_cache.get(node_key) {
Expand Down

0 comments on commit 501bb14

Please sign in to comment.