From 327499095121a66a320e516659ad6a7f9daf2801 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Mon, 4 Nov 2024 17:09:35 -0700
Subject: [PATCH] smt: make with_entries() a trait method

---
 src/merkle/smt/full/mod.rs   | 48 ++++++++++++++++++++----------------
 src/merkle/smt/mod.rs        |  7 ++++++
 src/merkle/smt/simple/mod.rs |  8 ++++++
 3 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs
index 89d0702..0b170d6 100644
--- a/src/merkle/smt/full/mod.rs
+++ b/src/merkle/smt/full/mod.rs
@@ -78,27 +78,7 @@ impl Smt {
     pub fn with_entries(
         entries: impl IntoIterator<Item = (RpoDigest, Word)>,
     ) -> Result<Self, MerkleError> {
-        // create an empty tree
-        let mut tree = Self::new();
-
-        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
-        // entries with the empty value need additional tracking.
-        let mut key_set_to_zero = BTreeSet::new();
-
-        for (key, value) in entries {
-            let old_value = tree.insert(key, value);
-
-            if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
-                return Err(MerkleError::DuplicateValuesForIndex(
-                    LeafIndex::<SMT_DEPTH>::from(key).value(),
-                ));
-            }
-
-            if value == EMPTY_WORD {
-                key_set_to_zero.insert(key);
-            };
-        }
-        Ok(tree)
+        <Self as SparseMerkleTree<SMT_DEPTH>>::with_entries(entries)
     }
 
     // PUBLIC ACCESSORS
@@ -267,6 +247,32 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
     const EMPTY_VALUE: Self::Value = EMPTY_WORD;
     const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
 
+    fn with_entries(
+        entries: impl IntoIterator<Item = (RpoDigest, Word)>,
+    ) -> Result<Self, MerkleError> {
+        // create an empty tree
+        let mut tree = Self::new();
+
+        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
+        // entries with the empty value need additional tracking.
+        let mut key_set_to_zero = BTreeSet::new();
+
+        for (key, value) in entries {
+            let old_value = tree.insert(key, value);
+
+            if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
+                return Err(MerkleError::DuplicateValuesForIndex(
+                    LeafIndex::<SMT_DEPTH>::from(key).value(),
+                ));
+            }
+
+            if value == EMPTY_WORD {
+                key_set_to_zero.insert(key);
+            };
+        }
+        Ok(tree)
+    }
+
     fn root(&self) -> RpoDigest {
         self.root
     }
diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index 79e9faa..7f9df46 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -292,6 +292,13 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     // REQUIRED METHODS
     // ---------------------------------------------------------------------------------------------
 
+    /// Construct a tree from an iterator of its keys and values.
+    fn with_entries(
+        entries: impl IntoIterator<Item = (Self::Key, Self::Value)>,
+    ) -> Result<Self, MerkleError>
+    where
+        Self: Sized;
+
     /// The root of the tree
     fn root(&self) -> RpoDigest;
 
diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs
index 04476a0..0a30b04 100644
--- a/src/merkle/smt/simple/mod.rs
+++ b/src/merkle/smt/simple/mod.rs
@@ -309,6 +309,14 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
     const EMPTY_VALUE: Self::Value = EMPTY_WORD;
     const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0);
 
+    fn with_entries(
+        entries: impl IntoIterator<Item = (LeafIndex<DEPTH>, Word)>,
+    ) -> Result<Self, MerkleError> {
+        <SimpleSmt<DEPTH>>::with_leaves(
+            entries.into_iter().map(|(key, value)| (key.value(), value)),
+        )
+    }
+
     fn root(&self) -> RpoDigest {
         self.root
     }