From 5de20ade48ba03b7ab6b0d00126b05f4da0ea7e2 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Tue, 5 Nov 2024 13:04:24 -0700
Subject: [PATCH] convert test_singlethreaded_subtree to use SubtreeLeavesIter

---
 src/merkle/smt/mod.rs | 75 ++++++++++++++++++++++++-------------------
 1 file changed, 42 insertions(+), 33 deletions(-)

diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index 7f9df46..dab9501 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -765,14 +765,15 @@ impl<'s> core::iter::Iterator for SubtreeLeavesIter<'s> {
 // ================================================================================================
 #[cfg(test)]
 mod test {
-    use core::mem;
-
     use alloc::{collections::BTreeMap, vec::Vec};
 
-    use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf};
+    use super::{
+        InnerNode, LeafIndex, PairComputations, SmtLeaf, SparseMerkleTree, SubtreeLeaf,
+        SubtreeLeavesIter,
+    };
     use crate::{
         hash::rpo::RpoDigest,
-        merkle::{LeafIndex, NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
+        merkle::{NodeIndex, Smt, SMT_DEPTH},
         Felt, Word, ONE,
     };
 
@@ -964,38 +965,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}");
         }