From 8c61bb224cd1038358277a494b3407855fceba90 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 28 Oct 2024 14:06:49 -0600 Subject: [PATCH] make build_subtree also return the next leaf row convert (col, hash) tuples to a dedicated struct --- benches/smt-subtree.rs | 12 ++++++------ src/merkle/mod.rs | 2 +- src/merkle/smt/full/mod.rs | 6 +++--- src/merkle/smt/mod.rs | 34 ++++++++++++++-------------------- src/merkle/smt/simple/mod.rs | 5 ++++- 5 files changed, 28 insertions(+), 31 deletions(-) diff --git a/benches/smt-subtree.rs b/benches/smt-subtree.rs index df97f79..cd7454a 100644 --- a/benches/smt-subtree.rs +++ b/benches/smt-subtree.rs @@ -3,7 +3,7 @@ use std::{fmt::Debug, hint, mem, time::Duration}; use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion}; use miden_crypto::{ hash::rpo::RpoDigest, - merkle::{NodeIndex, Smt, SmtLeaf, SMT_DEPTH}, + merkle::{NodeIndex, Smt, SmtLeaf, SubtreeLeaf, SMT_DEPTH}, Felt, Word, ONE, }; use rand_utils::prng_array; @@ -43,16 +43,16 @@ fn smt_subtree_even(c: &mut Criterion) { let leaf = SmtLeaf::new_single(*key, *value); let col = NodeIndex::from(leaf.index()).value(); let hash = leaf.hash(); - (col, hash) + SubtreeLeaf { col, hash } }) .collect(); leaves.sort(); - leaves.dedup_by_key(|leaf| leaf.0); + leaves.dedup_by_key(|leaf| leaf.col); leaves }, |leaves| { // Benchmarked function. - let subtree = + let (subtree, _) = Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH)); assert!(!subtree.is_empty()); }, @@ -93,14 +93,14 @@ fn smt_subtree_random(c: &mut Criterion) { let leaf = SmtLeaf::new_single(*key, *value); let col = NodeIndex::from(leaf.index()).value(); let hash = leaf.hash(); - (col, hash) + SubtreeLeaf { col, hash } }) .collect(); leaves.sort(); leaves }, |leaves| { - let subtree = + let (subtree, _) = Smt::build_subtree(hint::black_box(leaves), hint::black_box(SMT_DEPTH)); assert!(!subtree.is_empty()); }, diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index a562aa5..dc897dc 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -23,7 +23,7 @@ pub use path::{MerklePath, RootPath, ValuePath}; mod smt; pub use smt::{ LeafIndex, MutationSet, SimpleSmt, Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, - SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH, + SubtreeLeaf, SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; mod mmr; diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 90cc026..1be483a 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -6,7 +6,7 @@ use alloc::{ use super::{ EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, - MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, + MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, SubtreeLeaf, Word, EMPTY_WORD, }; mod error; @@ -251,9 +251,9 @@ impl Smt { } pub fn build_subtree( - leaves: Vec<(u64, RpoDigest)>, + leaves: Vec, bottom_depth: u8, - ) -> BTreeMap { + ) -> (BTreeMap, Vec) { >::build_subtree(leaves, bottom_depth) } } diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index ca3c492..063003b 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -357,11 +357,10 @@ pub(crate) trait SparseMerkleTree { /// more entries than can fit in a depth-8 subtree (more than 256), if `bottom_depth` is /// lower in the tree than the specified maximum depth (`DEPTH`), or if `leaves` is not sorted. // FIXME: more complete docstring. - #[cfg_attr(not(test), allow(dead_code))] fn build_subtree( - mut leaves: Vec<(u64, RpoDigest)>, + mut leaves: Vec, bottom_depth: u8, - ) -> BTreeMap { + ) -> (BTreeMap, Vec) { debug_assert!(bottom_depth <= DEPTH); debug_assert!(bottom_depth.is_multiple_of(&8)); debug_assert!(leaves.len() <= usize::pow(2, 8)); @@ -370,7 +369,7 @@ pub(crate) trait SparseMerkleTree { let mut inner_nodes: BTreeMap = Default::default(); - let mut next_leaves: Vec<(u64, RpoDigest)> = Vec::with_capacity(leaves.len() / 2); + let mut next_leaves: Vec = Vec::with_capacity(leaves.len() / 2); for next_depth in (subtree_root..bottom_depth).rev() { debug_assert!(next_depth <= bottom_depth); @@ -379,7 +378,7 @@ pub(crate) trait SparseMerkleTree { // `current_depth` is the stuff we have. let current_depth = next_depth + 1; - let mut iter = leaves.drain(..).map(SubtreeLeaf::from_tuple).peekable(); + let mut iter = leaves.drain(..).peekable(); while let Some(first) = iter.next() { // On non-continuous iterations, including the first iteration, `first_column` may // be a left or right node. On subsequent continuous iterations, we will always call @@ -432,7 +431,7 @@ pub(crate) trait SparseMerkleTree { if hash != equivalent_empty_hash { inner_nodes.insert(index, node); // FIXME: is it possible for this to end up not being sorted? I don't think so. - next_leaves.push((index.value(), hash)); + next_leaves.push(SubtreeLeaf { col: index.value(), hash }); } } @@ -445,7 +444,7 @@ pub(crate) trait SparseMerkleTree { mem::swap(&mut leaves, &mut next_leaves); } - inner_nodes + (inner_nodes, leaves) } } @@ -568,14 +567,9 @@ impl MutationSet { // HELPERS // ================================================================================================ #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default)] -struct SubtreeLeaf { - col: u64, - hash: RpoDigest, -} -impl SubtreeLeaf { - const fn from_tuple((col, hash): (u64, RpoDigest)) -> Self { - Self { col, hash } - } +pub struct SubtreeLeaf { + pub col: u64, + pub hash: RpoDigest, } // TESTS @@ -587,7 +581,7 @@ mod test { use super::SparseMerkleTree; use crate::{ hash::rpo::RpoDigest, - merkle::{Smt, SmtLeaf, SMT_DEPTH}, + merkle::{smt::SubtreeLeaf, Smt, SmtLeaf, SMT_DEPTH}, Felt, Word, ONE, }; @@ -606,19 +600,19 @@ mod test { let control = Smt::with_entries(entries.clone()).unwrap(); - let mut leaves: Vec<(u64, RpoDigest)> = entries + let mut leaves: Vec = entries .iter() .map(|(key, value)| { let leaf = SmtLeaf::new_single(*key, *value); let col = leaf.index().index.value(); let hash = leaf.hash(); - (col, hash) + SubtreeLeaf { col, hash } }) .collect(); leaves.sort(); - leaves.dedup_by_key(|leaf| leaf.0); + leaves.dedup_by_key(|leaf| leaf.col); - let first_subtree = Smt::build_subtree(leaves, SMT_DEPTH); + let (first_subtree, _) = Smt::build_subtree(leaves, SMT_DEPTH); assert!(!first_subtree.is_empty()); for (index, node) in first_subtree.into_iter() { diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 6229ac2..edd9849 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -1,4 +1,7 @@ -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::{ + collections::{BTreeMap, BTreeSet}, + vec::Vec, +}; use super::{ super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,