working test_singlethreaded_subtrees()

src/merkle/smt/mod.rs

@@ -611,6 +611,10 @@ impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
 
 // HELPERS
 // ================================================================================================
+
+/// A depth-8 subtree contains 256 "columns" that can possibly be occupied.
+const COLS_PER_SUBTREE: u64 = u64::pow(2, 8);
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
 pub struct SubtreeLeaf {
     pub col: u64,
@@ -637,9 +641,6 @@ pub struct PairComputations<L> {
 
 impl<L> PairComputations<L> {
     pub fn add_leaf(&mut self, leaf: SubtreeLeaf) {
-        // A depth-8 subtree contains 256 "columns" that can possibly be occupied.
-        const COLS_PER_SUBTREE: u64 = u64::pow(2, 8);
-
         let last_subtree = match self.leaves.last_mut() {
             // Base case.
             None => {
@@ -654,7 +655,7 @@ impl<L> PairComputations<L> {
 
         // The multiple of 256 after 0 is 1, but 0 and 1 do not belong to different subtrees.
         let last_subtree_col = u64::max(1, last_subtree.last().unwrap().col);
-        let next_subtree_col = if last_subtree_col.is_multiple_of(&COLS_PER_SUBTREE) {
+        let next_subtree_col = if Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE) {
             u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE)
         } else {
             last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
@@ -662,7 +663,6 @@ impl<L> PairComputations<L> {
         if leaf.col < next_subtree_col {
             last_subtree.push(leaf);
         } else {
-            //std::eprintln!("\tcreating new subtree for column {}", leaf.col);
             let next_subtree = vec![leaf];
             self.leaves.push(next_subtree);
         }
@@ -729,15 +729,16 @@ impl<'s> core::iter::Iterator for SubtreeLeavesIter<'s> {
 // ================================================================================================
 #[cfg(test)]
 mod test {
+    use core::mem;
+
     use alloc::{collections::BTreeMap, vec::Vec};
 
-    use super::SparseMerkleTree;
+    use num::Integer;
+
+    use super::{InnerNode, PairComputations, SparseMerkleTree, SubtreeLeaf, COLS_PER_SUBTREE};
     use crate::{
         hash::rpo::RpoDigest,
-        merkle::{
-            smt::{PairComputations, SubtreeLeaf},
-            Smt, SmtLeaf, SMT_DEPTH,
-        },
+        merkle::{NodeIndex, Smt, SmtLeaf, SMT_DEPTH},
         Felt, Word, ONE,
     };
 
@@ -811,7 +812,7 @@ mod test {
             }
             let control_leaf = control_leaves
                 .get(&column)
-                .expect(&format!("no leaf node found for column {column}"));
+                .unwrap_or_else(|| panic!("no leaf node found for column {column}"));
             assert_eq!(control_leaf, &test_leaf);
         }
     }
@@ -905,4 +906,92 @@ mod test {
             );
         }
     }
+
+    #[test]
+    fn test_singlethreaded_subtrees() {
+        const PAIR_COUNT: u64 = 4096 * 4;
+
+        let entries = generate_entries(PAIR_COUNT);
+
+        let control = Smt::with_entries(entries.clone()).unwrap();
+
+        let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
+
+        let starting_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.
+                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,
+                    );
+                }
+
+                // Update state.
+                accumulated_nodes.extend(nodes);
+
+                for subtree_leaf in next_leaves {
+                    if leaf_subtrees.is_empty() {
+                        leaf_subtrees.push(vec![subtree_leaf]);
+                        continue;
+                    }
+
+                    let buffer_max_col =
+                        u64::max(1, leaf_subtrees.last().unwrap().last().unwrap().col);
+                    let next_subtree_col =
+                        if Integer::is_multiple_of(&buffer_max_col, &COLS_PER_SUBTREE) {
+                            u64::next_multiple_of(buffer_max_col + 1, COLS_PER_SUBTREE)
+                        } else {
+                            buffer_max_col.next_multiple_of(COLS_PER_SUBTREE)
+                        };
+
+                    if subtree_leaf.col < next_subtree_col {
+                        leaf_subtrees.last_mut().unwrap().push(subtree_leaf);
+                    } else {
+                        leaf_subtrees.push(vec![subtree_leaf]);
+                    }
+                }
+            }
+
+            assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
+        }
+
+        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();
+        assert_eq!(control.root(), root_leaf.hash);
+
+        // Do we have a root?
+        assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
+
+        // And does it match?
+        let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
+        assert_eq!(control.root(), test_root.hash());
+        assert_eq!(control.root(), root_leaf.hash);
+    }
 }