convert test_singlethreaded_subtree to use an iterator adapter instead of state mutation

src/merkle/smt/mod.rs

@@ -684,19 +684,65 @@ 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)]
 mod test {
-    use core::mem;
-
     use alloc::{collections::BTreeMap, vec::Vec};
 
     use super::{SparseMerkleTree, SubtreeLeaf};
     use crate::{
         hash::rpo::RpoDigest,
         merkle::{
-            smt::{InnerNode, PairComputations},
+            smt::{InnerNode, PairComputations, SubtreeLeavesIter},
             LeafIndex, NodeIndex, Smt, SmtLeaf, SMT_DEPTH,
         },
         Felt, Word, ONE,
@@ -890,38 +936,46 @@ mod test {
         } = Smt::sorted_pairs_to_leaves(entries);
 
         for current_depth in (8..=SMT_DEPTH).step_by(8).rev() {
-            for (i, subtree) in mem::take(&mut leaf_subtrees).into_iter().enumerate() {
-                // Pre-assertions.
-                assert!(
-                    subtree.is_sorted(),
-                    "subtree {i} at bottom-depth {current_depth} is not sorted",
-                );
-                assert!(
-                    !subtree.is_empty(),
-                    "subtree {i} at bottom-depth {current_depth} is empty!",
-                );
-
-                // Do actual things.
-                let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth);
-
-                // Post-assertions.
-                assert!(next_leaves.is_sorted());
-                for (&index, test_node) in nodes.iter() {
-                    let control_node = control.get_inner_node(index);
-                    assert_eq!(
-                        test_node, &control_node,
-                        "depth {} subtree {}: test node does not match control at index {:?}",
-                        current_depth, i, index,
+            // There's no flat_map_unzip(), so this is the best we can do.
+            let (nodes, subtrees): (Vec<BTreeMap<_, _>>, Vec<Vec<SubtreeLeaf>>) = leaf_subtrees
+                .into_iter()
+                .enumerate()
+                .map(|(i, subtree)| {
+                    // Pre-assertions.
+                    assert!(
+                        subtree.is_sorted(),
+                        "subtree {i} at bottom-depth {current_depth} is not sorted",
+                    );
+                    assert!(
+                        !subtree.is_empty(),
+                        "subtree {i} at bottom-depth {current_depth} is empty!",
                     );
-                }
 
-                // Update state.
-                accumulated_nodes.extend(nodes);
+                    // Do actual things.
+                    let (nodes, next_leaves) = Smt::build_subtree(subtree, current_depth);
+                    // Post-assertions.
+                    assert!(next_leaves.is_sorted());
 
-                for subtree_leaf in next_leaves {
-                    super::add_subtree_leaf(&mut leaf_subtrees, subtree_leaf);
-                }
-            }
+                    for (&index, test_node) in nodes.iter() {
+                        let control_node = control.get_inner_node(index);
+                        assert_eq!(
+                            test_node, &control_node,
+                            "depth {} subtree {}: test node does not match control at index {:?}",
+                            current_depth, i, index,
+                        );
+                    }
+
+                    (nodes, next_leaves)
+                })
+                .unzip();
+
+            // Update state between each depth iteration.
+
+            // FIXME: is this flatten or Box<dyn Iterator> better?
+            let mut all_leaves: Vec<SubtreeLeaf> = subtrees.into_iter().flatten().collect();
+            leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut all_leaves).collect();
+
+            accumulated_nodes.extend(nodes.into_iter().flatten());
 
             assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
         }