From bd0a69bcfa11bb1c03609229b238420828232c89 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Mon, 28 Oct 2024 14:06:49 -0600
Subject: [PATCH] make build_subtree also return the next leaf row

---
 benches/smt-subtree.rs     | 4 ++--
 src/merkle/smt/full/mod.rs | 2 +-
 src/merkle/smt/mod.rs      | 6 +++---
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/benches/smt-subtree.rs b/benches/smt-subtree.rs
index df97f79..ca5556b 100644
--- a/benches/smt-subtree.rs
+++ b/benches/smt-subtree.rs
@@ -52,7 +52,7 @@ fn smt_subtree_even(c: &mut Criterion) {
                 },
                 |leaves| {
                     // Benchmarked function.
-                    let subtree =
+                    let (subtree, _) =
                         Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH));
                     assert!(!subtree.is_empty());
                 },
@@ -100,7 +100,7 @@ fn smt_subtree_random(c: &mut Criterion) {
                     leaves
                 },
                 |leaves| {
-                    let subtree =
+                    let (subtree, _) =
                         Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH));
                     assert!(!subtree.is_empty());
                 },
diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs
index 90cc026..5e12014 100644
--- a/src/merkle/smt/full/mod.rs
+++ b/src/merkle/smt/full/mod.rs
@@ -253,7 +253,7 @@ impl Smt {
     pub fn build_subtree(
         leaves: Vec<(u64, RpoDigest)>,
         bottom_depth: u8,
-    ) -> BTreeMap<NodeIndex, InnerNode> {
+    ) -> (BTreeMap<NodeIndex, InnerNode>, Vec<(u64, RpoDigest)>) {
         <Self as SparseMerkleTree<SMT_DEPTH>>::build_subtree(leaves, bottom_depth)
     }
 }
diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index ca3c492..950463d 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -361,7 +361,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     fn build_subtree(
         mut leaves: Vec<(u64, RpoDigest)>,
         bottom_depth: u8,
-    ) -> BTreeMap<NodeIndex, InnerNode> {
+    ) -> (BTreeMap<NodeIndex, InnerNode>, Vec<(u64, RpoDigest)>) {
         debug_assert!(bottom_depth <= DEPTH);
         debug_assert!(bottom_depth.is_multiple_of(&8));
         debug_assert!(leaves.len() <= usize::pow(2, 8));
@@ -445,7 +445,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
             mem::swap(&mut leaves, &mut next_leaves);
         }
 
-        inner_nodes
+        (inner_nodes, leaves)
     }
 }
 
@@ -618,7 +618,7 @@ mod test {
         leaves.sort();
         leaves.dedup_by_key(|leaf| leaf.0);
 
-        let first_subtree = Smt::build_subtree(leaves, SMT_DEPTH);
+        let (first_subtree, _) = Smt::build_subtree(leaves, SMT_DEPTH);
         assert!(!first_subtree.is_empty());
 
         for (index, node) in first_subtree.into_iter() {