From a1929c3d2f8b239575decc7d63c4014c12a97a98 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Thu, 31 Oct 2024 18:58:18 -0600 Subject: [PATCH] working test_singlethreaded_subtrees() --- src/merkle/smt/mod.rs | 110 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 101 insertions(+), 9 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 6587d70..becb3f2 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -407,7 +407,7 @@ pub(crate) trait SparseMerkleTree { bottom_depth: u8, ) -> (BTreeMap, Vec) { debug_assert!(bottom_depth <= DEPTH); - debug_assert!(bottom_depth.is_multiple_of(&8)); + debug_assert!(Integer::is_multiple_of(&bottom_depth, &8)); debug_assert!(leaves.len() <= usize::pow(2, 8)); let subtree_root = bottom_depth - 8; @@ -611,6 +611,10 @@ impl MutationSet { // 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 PrecomputedSubtrees { impl PrecomputedSubtrees { 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 PrecomputedSubtrees { // 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 PrecomputedSubtrees { 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); } @@ -683,12 +683,16 @@ impl Default for PrecomputedSubtrees { // ================================================================================================ #[cfg(test)] mod test { + use core::mem; + use alloc::{collections::BTreeMap, vec::Vec}; - use super::SparseMerkleTree; + use num::Integer; + + use super::{SparseMerkleTree, SubtreeLeaf, COLS_PER_SUBTREE}; use crate::{ hash::rpo::RpoDigest, - merkle::{smt::SubtreeLeaf, Smt, SmtLeaf, SMT_DEPTH}, + merkle::{smt::InnerNode, NodeIndex, Smt, SmtLeaf, SMT_DEPTH}, Felt, Word, ONE, }; @@ -762,7 +766,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); } } @@ -856,4 +860,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 = 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); + } }