From 5d551cf0984a3692a228ce6e3cdebed0834673cd Mon Sep 17 00:00:00 2001
From: Erwan Or <erwanor@penumbralabs.xyz>
Date: Sun, 17 Mar 2024 16:29:15 -0400
Subject: [PATCH] jmt: simplify append api

---
 src/tree.rs | 47 +++++++++++++++++++++++++----------------------
 1 file changed, 25 insertions(+), 22 deletions(-)

diff --git a/src/tree.rs b/src/tree.rs
index 5d3eede..5662c3a 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -446,33 +446,36 @@ where
     }
 
     #[cfg(feature = "migration")]
-    /// Append value sets to the latest version of the tree.
-    pub fn append_value_sets(
+    /// Append value sets to the latest version of the tree, without incrementing its version.
+    pub fn append_value_set(
         &self,
-        value_sets: impl IntoIterator<Item = impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>>,
+        value_set: impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>,
         latest_version: Version,
-    ) -> Result<(Vec<RootHash>, TreeUpdateBatch)> {
+    ) -> Result<(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>()?;
+        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
+                    )
+                })?;
         }
 
-        Ok(tree_cache.into())
+        // 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.