From aa3197fcc1184310d5b631f8f5ba98351646168d Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Mon, 4 Nov 2024 13:38:55 -0700
Subject: [PATCH] cleanup test_singlethreaded_subtrees a bit

---
 src/merkle/smt/mod.rs | 52 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 41 insertions(+), 11 deletions(-)

diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs
index e97e6ee..d3db29f 100644
--- a/src/merkle/smt/mod.rs
+++ b/src/merkle/smt/mod.rs
@@ -615,9 +615,15 @@ impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
 /// A depth-8 subtree contains 256 "columns" that can possibly be occupied.
 const COLS_PER_SUBTREE: u64 = u64::pow(2, 8);
 
+/// Helper struct for organizing the data we care about when computing Merkle subtrees.
+///
+/// Note that these represet "conceptual" leaves of some subtree, not necessarily
+/// [`SparseMerkleTree::Leaf`].
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
 pub struct SubtreeLeaf {
+    /// The 'value' field of [`NodeIndex`]. When computing a subtree, the depth is already known.
     pub col: u64,
+    /// The hash of the node this `SubtreeLeaf` represents.
     pub hash: RpoDigest,
 }
 
@@ -769,7 +775,7 @@ mod test {
     use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf};
     use crate::{
         hash::rpo::RpoDigest,
-        merkle::{NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
+        merkle::{LeafIndex, NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
         Felt, Word, ONE,
     };
 
@@ -948,9 +954,11 @@ mod test {
 
         let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
 
-        let starting_leaves = Smt::sorted_pairs_to_leaves(entries);
+        let PairComputations {
+            leaves: mut leaf_subtrees,
+            nodes: test_leaves,
+        } = Smt::sorted_pairs_to_leaves(entries);
 
-        let mut leaf_subtrees = starting_leaves.leaves;
         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.
@@ -988,23 +996,45 @@ mod test {
             assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
         }
 
+        // Make sure the true leaves match, checking length first and then each individual leaf.
+        let control_leaves: BTreeMap<_, _> = control.leaves().collect();
+        let control_leaves_len = control_leaves.len();
+        let test_leaves_len = test_leaves.len();
+        assert_eq!(test_leaves_len, control_leaves_len);
+        for (col, ref test_leaf) in test_leaves {
+            let index = LeafIndex::new_max_depth(col);
+            let &control_leaf = control_leaves.get(&index).unwrap();
+            assert_eq!(test_leaf, control_leaf);
+        }
+
+        // Make sure the inner nodes match, checking length first and then each individual leaf.
+        let control_nodes_len = control.inner_nodes().count();
+        let test_nodes_len = accumulated_nodes.len();
+        assert_eq!(test_nodes_len, control_nodes_len);
         for (index, test_node) in accumulated_nodes.clone() {
             let control_node = control.get_inner_node(index);
             assert_eq!(test_node, control_node, "test node does not match control at {index:?}");
         }
 
-        assert_eq!(leaf_subtrees.len(), 1);
-        let mut leaf_subtree = leaf_subtrees.pop().unwrap();
-        assert_eq!(leaf_subtree.len(), 1);
-        let root_leaf = leaf_subtree.pop().unwrap();
+        // After the last iteration of the above for loop, we should have the new root node actually
+        // in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from
+        // `build_subtree()`. So let's check both!
+
+        let control_root = control.get_inner_node(NodeIndex::root());
+
+        // That for loop should have left us with only one leaf subtree...
+        let [leaf_subtree]: [_; 1] = leaf_subtrees.try_into().unwrap();
+        // which itself contains only one 'leaf'...
+        let [root_leaf]: [_; 1] = leaf_subtree.try_into().unwrap();
+        // which matches the expected root.
         assert_eq!(control.root(), root_leaf.hash);
 
-        // Do we have a root?
+        // Likewise `accumulated_nodes` should contain a node at the root index...
         assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
-
-        // And does it match?
+        // and it should match our actual root.
         let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
-        assert_eq!(control.root(), test_root.hash());
+        assert_eq!(control_root, *test_root);
+        // And of course the root we got from each place should match.
         assert_eq!(control.root(), root_leaf.hash);
     }
 }