From 1c5fc8a8300e8ae347da54210fa75bf56f3632f8 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Wed, 13 Nov 2024 12:32:21 -0700
Subject: [PATCH] add SubtreeLeavesIter

---
 src/merkle/smt/mod.rs | 48 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index b8abdba..e97e6ee 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -710,6 +710,54 @@ fn add_subtree_leaf(subtrees: &mut Vec<Vec<SubtreeLeaf>>, leaf: SubtreeLeaf) {
     }
 }
 
+#[derive(Debug)]
+struct SubtreeLeavesIter<'s> {
+    leaves: core::iter::Peekable<alloc::vec::Drain<'s, SubtreeLeaf>>,
+}
+
+impl<'s> SubtreeLeavesIter<'s> {
+    fn from_leaves(leaves: &'s mut Vec<SubtreeLeaf>) -> Self {
+        Self { leaves: leaves.drain(..).peekable() }
+    }
+}
+
+impl<'s> core::iter::Iterator for SubtreeLeavesIter<'s> {
+    type Item = Vec<SubtreeLeaf>;
+
+    /// Each `next()` collects an entire subtree.
+    fn next(&mut self) -> Option<Vec<SubtreeLeaf>> {
+        let mut subtree: Vec<SubtreeLeaf> = Default::default();
+
+        let mut last_subtree_col = 0;
+
+        while let Some(leaf) = self.leaves.peek() {
+            last_subtree_col = u64::max(1, last_subtree_col);
+            let is_exact_multiple = Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE);
+            let next_subtree_col = if is_exact_multiple {
+                u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE)
+            } else {
+                last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
+            };
+
+            last_subtree_col = leaf.col;
+            if leaf.col < next_subtree_col {
+                subtree.push(self.leaves.next().unwrap());
+            } else if subtree.is_empty() {
+                continue;
+            } else {
+                break;
+            }
+        }
+
+        if subtree.is_empty() {
+            debug_assert!(self.leaves.peek().is_none());
+            return None;
+        }
+
+        Some(subtree)
+    }
+}
+
 // TESTS
 // ================================================================================================
 #[cfg(test)]