jmt: allow incremental tree overwrites for a given version (#110)

* jmt: prototype for `append_value_sets`

* jmt(tree_cache): construct node key internally

* jmt(tree_cache): handle no-op writes

* jmt: fix needless borrow

* jmt(tree_cache): conservative guards for version mismatches

* jmt: simplify append api

* jmt(tree_cache): fix compile warning

Cargo.toml

@@ -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"

src/lib.rs

@@ -255,23 +255,23 @@ impl KeyHash {
 impl core::fmt::Debug for KeyHash {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_tuple("KeyHash")
-            .field(&hex::encode(&self.0))
+            .field(&hex::encode(self.0))
             .finish()
     }
 }
 
 impl core::fmt::Debug for ValueHash {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_tuple("ValueHash")
-            .field(&hex::encode(&self.0))
+            .field(&hex::encode(self.0))
             .finish()
     }
 }
 
 impl core::fmt::Debug for RootHash {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_tuple("RootHash")
-            .field(&hex::encode(&self.0))
+            .field(&hex::encode(self.0))
             .finish()
     }
 }

src/tree.rs

@@ -445,6 +445,39 @@ where
         Ok(tree_cache.into())
     }
 
+    #[cfg(feature = "migration")]
+    /// Append value sets to the latest version of the tree, without incrementing its version.
+    pub fn append_value_set(
+        &self,
+        value_set: impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>,
+        latest_version: Version,
+    ) -> Result<(RootHash, TreeUpdateBatch)> {
+        let mut tree_cache = TreeCache::new_overwrite(self.reader, latest_version)?;
+        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(latest_version, key, value);
+            self.put(key, value_hash, latest_version, &mut tree_cache, false)
+                .with_context(|| {
+                    format!(
+                        "failed to {} key {} for version {}, key = {:?}",
+                        action, i, latest_version, key
+                    )
+                })?;
+        }
+
+        // Freezes the current cache to make all contents in the current cache immutable.
+        tree_cache.freeze::<H>()?;
+        let (root_hash_vec, tree_batch) = tree_cache.into();
+        if root_hash_vec.len() != 1 {
+            bail!(
+                "appending a value set failed, we expected a single root hash, but got {}",
+                root_hash_vec.len()
+            );
+        }
+        Ok((root_hash_vec[0], tree_batch))
+    }
+
     /// 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.

src/tree_cache.rs

@@ -188,6 +188,38 @@ where
         })
     }
 
+    #[cfg(feature = "migration")]
+    /// Instantiate a [`TreeCache`] over the a [`TreeReader`] that is defined
+    /// against a root node key at version `current_version`.
+    ///
+    /// # Usage
+    /// This method is used to perform incremental addition to a tree without
+    /// increasing the tree version's number.
+    pub fn new_overwrite(reader: &'a R, current_version: Version) -> Result<Self> {
+        let node_cache = HashMap::new();
+        let Some((node_key, _)) = reader.get_rightmost_leaf()? else {
+            bail!("creating an overwrite cache for an empty tree is not supported")
+        };
+
+        anyhow::ensure!(
+            node_key.version() == current_version,
+            "the supplied version is not the latest version of the tree"
+        );
+
+        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,
+            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) {