From 501bb1477bcf69970d3ee0879ca1635c725a6cc3 Mon Sep 17 00:00:00 2001
From: Erwan Or <erwan.ounn.84@gmail.com>
Date: Thu, 7 Mar 2024 15:45:59 -0400
Subject: [PATCH] jmt: prototype for `append_value_sets`

---
 Cargo.toml        |  1 +
 src/tree.rs       | 30 ++++++++++++++++++++++++++++++
 src/tree_cache.rs | 21 +++++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index 61d13c7..906937d 100644
--- a/Cargo.toml
+++ b/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"
diff --git a/src/tree.rs b/src/tree.rs
index 9d9b339..5d3eede 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -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.
diff --git a/src/tree_cache.rs b/src/tree_cache.rs
index 64246be..0f24821 100644
--- a/src/tree_cache.rs
+++ b/src/tree_cache.rs
@@ -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) {