From 25b8cb64baf62f399e23b796a827016480deb31d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philippe=20Laferri=C3=A8re?= <plafer@protonmail.com>
Date: Mon, 27 Nov 2023 14:19:01 -0500
Subject: [PATCH] Introduce `SimpleSmt::with_contiguous_leaves()` (#227)

* with_contiguous_leaves

* test
---
 src/merkle/simple_smt/mod.rs   | 16 ++++++++++++++++
 src/merkle/simple_smt/tests.rs | 15 +++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs
index 39f52930d..e6e11d12f 100644
--- a/src/merkle/simple_smt/mod.rs
+++ b/src/merkle/simple_smt/mod.rs
@@ -104,6 +104,22 @@ impl SimpleSmt {
         Ok(tree)
     }
 
+    /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
+    /// starting at index 0.
+    pub fn with_contiguous_leaves<R, I>(depth: u8, entries: R) -> Result<Self, MerkleError>
+    where
+        R: IntoIterator<IntoIter = I>,
+        I: Iterator<Item = Word> + ExactSizeIterator,
+    {
+        Self::with_leaves(
+            depth,
+            entries
+                .into_iter()
+                .enumerate()
+                .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
+        )
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs
index 0bd818b05..5070d91eb 100644
--- a/src/merkle/simple_smt/tests.rs
+++ b/src/merkle/simple_smt/tests.rs
@@ -71,6 +71,21 @@ fn build_sparse_tree() {
     assert_eq!(old_value, EMPTY_WORD);
 }
 
+/// Tests that [`SimpleSmt::with_contiguous_leaves`] works as expected
+#[test]
+fn build_contiguous_tree() {
+    let tree_with_leaves = SimpleSmt::with_leaves(
+        2,
+        [0, 1, 2, 3].into_iter().zip(digests_to_words(&VALUES4).into_iter()),
+    )
+    .unwrap();
+
+    let tree_with_contiguous_leaves =
+        SimpleSmt::with_contiguous_leaves(2, digests_to_words(&VALUES4).into_iter()).unwrap();
+
+    assert_eq!(tree_with_leaves, tree_with_contiguous_leaves);
+}
+
 #[test]
 fn test_depth2_tree() {
     let tree =