From b585f9cac602c2bc68026955b21290eb373aa3ce Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 13 Nov 2024 15:17:40 -0700 Subject: [PATCH 1/3] merkle: add parent() helper function on NodeIndex --- src/merkle/index.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/merkle/index.rs b/src/merkle/index.rs index 104ceb4..24b3808 100644 --- a/src/merkle/index.rs +++ b/src/merkle/index.rs @@ -97,6 +97,14 @@ impl NodeIndex { self } + /// Returns the parent of the current node. This is the same as [`Self::move_up()`], but returns + /// a new value instead of mutating `self`. + pub const fn parent(mut self) -> Self { + self.depth = self.depth.saturating_sub(1); + self.value >>= 1; + self + } + // PROVIDERS // -------------------------------------------------------------------------------------------- From f5d3dea50d0694ff532ee76efb4aa118d27c401c Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 13 Nov 2024 15:19:05 -0700 Subject: [PATCH 2/3] smt: add pairs_to_leaf() to trait --- src/merkle/smt/full/mod.rs | 24 ++++++++++++++++++++++++ src/merkle/smt/mod.rs | 4 ++++ src/merkle/smt/simple/mod.rs | 12 +++++++++++- 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 226a8b1..1e2c574 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -344,6 +344,30 @@ impl SparseMerkleTree for Smt { fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof { SmtProof::new_unchecked(path, leaf) } + + fn pairs_to_leaf(mut pairs: Vec<(RpoDigest, Word)>) -> SmtLeaf { + assert!(!pairs.is_empty()); + + // FIXME + if cfg!(debug_assertions) { + let mut control = pairs.clone(); + control.sort_by_key(|(key, _)| Self::key_to_leaf_index(key).index.value()); + assert_eq!(control, pairs); + } + + if pairs.len() > 1 { + SmtLeaf::new_multiple(pairs).unwrap() + } else { + let (key, value) = pairs.pop().unwrap(); + // FIXME: should we ever be constructing empty leaves from pairs? + if value == Self::EMPTY_VALUE { + let index = Self::key_to_leaf_index(&key); + SmtLeaf::new_empty(index) + } else { + SmtLeaf::new_single(key, value) + } + } + } } impl Default for Smt { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 056c221..03d9d45 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -338,6 +338,10 @@ pub(crate) trait SparseMerkleTree { /// Maps a key to a leaf index fn key_to_leaf_index(key: &Self::Key) -> LeafIndex; + /// Constructs a single leaf from an arbitrary amount of key-value pairs. + /// Those pairs must all have the same leaf index. + fn pairs_to_leaf(pairs: Vec<(Self::Key, Self::Value)>) -> Self::Leaf; + /// Maps a (MerklePath, Self::Leaf) to an opening. /// /// The length `path` is guaranteed to be equal to `DEPTH` diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 6229ac2..04476a0 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, @@ -370,4 +373,11 @@ impl SparseMerkleTree for SimpleSmt { fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath { (path, leaf).into() } + + fn pairs_to_leaf(mut pairs: Vec<(LeafIndex, Word)>) -> Word { + // SimpleSmt can't have more than one value per key. + assert_eq!(pairs.len(), 1); + let (_key, value) = pairs.pop().unwrap(); + value + } } From 46f951b0f1033e3ec273e783253f6e03c4781617 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Wed, 13 Nov 2024 15:32:48 -0700 Subject: [PATCH 3/3] smt: add sorted_pairs_to_leaves() and test for it --- src/merkle/smt/mod.rs | 227 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 03d9d45..6a86c28 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -1,4 +1,7 @@ use alloc::{collections::BTreeMap, vec::Vec}; +use core::mem; + +use num::Integer; use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex}; use crate::{ @@ -346,6 +349,49 @@ pub(crate) trait SparseMerkleTree { /// /// The length `path` is guaranteed to be equal to `DEPTH` fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening; + + fn sorted_pairs_to_leaves( + pairs: Vec<(Self::Key, Self::Value)>, + ) -> PairComputations { + let mut accumulator: PairComputations = Default::default(); + let mut accumulated_leaves: Vec = Default::default(); + + // The kv-pairs we've seen so far that correspond to a single leaf. + let mut current_leaf_buffer: Vec<(Self::Key, Self::Value)> = Default::default(); + + let mut iter = pairs.into_iter().peekable(); + while let Some((key, value)) = iter.next() { + let col = Self::key_to_leaf_index(&key).index.value(); + let peeked_col = iter.peek().map(|(key, _v)| { + let index = Self::key_to_leaf_index(key); + let next_col = index.index.value(); + // We panic if `pairs` is not sorted by column. + debug_assert!(next_col >= col); + next_col + }); + current_leaf_buffer.push((key, value)); + + // If the next pair is the same column as this one, then we're done after adding this + // pair to the buffer. + if peeked_col == Some(col) { + continue; + } + + // Otherwise, the next pair is a different column, or there is no next pair. Either way + // it's time to swap out our buffer. + let leaf_pairs = mem::take(&mut current_leaf_buffer); + let leaf = Self::pairs_to_leaf(leaf_pairs); + let hash = Self::hash_leaf(&leaf); + + accumulator.nodes.insert(col, leaf); + accumulated_leaves.push(SubtreeLeaf { col, hash }); + + debug_assert!(current_leaf_buffer.is_empty()); + } + + accumulator.leaves = SubtreeLeavesIter::from_leaves(&mut accumulated_leaves).collect(); + accumulator + } } // INNER NODE @@ -463,3 +509,184 @@ impl MutationSet { self.new_root } } + +// SUBTREES +// ================================================================================================ +/// 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, +} + +/// Helper struct to organize the return value of [`SparseMerkleTree::sorted_pairs_to_leaves()`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PairComputations { + /// Literal leaves to be added to the sparse Merkle tree's internal mapping. + pub nodes: BTreeMap, + /// "Conceptual" leaves that will be used for computations. + pub leaves: Vec>, +} + +// Derive requires `L` to impl Default, even though we don't actually need that. +impl Default for PairComputations { + fn default() -> Self { + Self { + nodes: Default::default(), + leaves: Default::default(), + } + } +} + +#[derive(Debug)] +struct SubtreeLeavesIter<'s> { + leaves: core::iter::Peekable>, +} +impl<'s> SubtreeLeavesIter<'s> { + fn from_leaves(leaves: &'s mut Vec) -> Self { + Self { leaves: leaves.drain(..).peekable() } + } +} +impl<'s> core::iter::Iterator for SubtreeLeavesIter<'s> { + type Item = Vec; + + /// Each `next()` collects an entire subtree. + fn next(&mut self) -> Option> { + let mut subtree: Vec = Default::default(); + + let mut last_subtree_col = 0; + + while let Some(leaf) = self.leaves.peek() { + last_subtree_col = u64::max(1, last_subtree_col); + let is_exact_multiple = Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE); + let next_subtree_col = if is_exact_multiple { + u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE) + } else { + last_subtree_col.next_multiple_of(COLS_PER_SUBTREE) + }; + + last_subtree_col = leaf.col; + if leaf.col < next_subtree_col { + subtree.push(self.leaves.next().unwrap()); + } else if subtree.is_empty() { + continue; + } else { + break; + } + } + + if subtree.is_empty() { + debug_assert!(self.leaves.peek().is_none()); + return None; + } + + Some(subtree) + } +} + +// TESTS +// ================================================================================================ +#[cfg(test)] +mod test { + use alloc::{collections::BTreeMap, vec::Vec}; + + use super::{PairComputations, SmtLeaf, SparseMerkleTree, SubtreeLeaf, SubtreeLeavesIter}; + + use crate::{hash::rpo::RpoDigest, merkle::Smt, Felt, Word, ONE}; + + fn smtleaf_to_subtree_leaf(leaf: &SmtLeaf) -> SubtreeLeaf { + SubtreeLeaf { + col: leaf.index().index.value(), + hash: leaf.hash(), + } + } + + #[test] + fn test_sorted_pairs_to_leaves() { + let entries: Vec<(RpoDigest, Word)> = vec![ + // Subtree 0. + (RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]), + (RpoDigest::new([ONE, ONE, ONE, Felt::new(17)]), [ONE; 4]), + // Leaf index collision. + (RpoDigest::new([ONE, ONE, Felt::new(10), Felt::new(20)]), [ONE; 4]), + (RpoDigest::new([ONE, ONE, Felt::new(20), Felt::new(20)]), [ONE; 4]), + // Subtree 1. Normal single leaf again. + (RpoDigest::new([ONE, ONE, ONE, Felt::new(400)]), [ONE; 4]), // Subtree boundary. + (RpoDigest::new([ONE, ONE, ONE, Felt::new(401)]), [ONE; 4]), + // Subtree 2. Another normal leaf. + (RpoDigest::new([ONE, ONE, ONE, Felt::new(1024)]), [ONE; 4]), + ]; + + let control = Smt::with_entries(entries.clone()).unwrap(); + + let control_leaves: Vec = { + let mut entries_iter = entries.iter().cloned(); + let mut next_entry = || entries_iter.next().unwrap(); + let control_leaves = vec![ + // Subtree 0. + SmtLeaf::Single(next_entry()), + SmtLeaf::Single(next_entry()), + SmtLeaf::new_multiple(vec![next_entry(), next_entry()]).unwrap(), + // Subtree 1. + SmtLeaf::Single(next_entry()), + SmtLeaf::Single(next_entry()), + // Subtree 2. + SmtLeaf::Single(next_entry()), + ]; + assert_eq!(entries_iter.next(), None); + control_leaves + }; + + let control_subtree_leaves: Vec> = { + let mut control_leaves_iter = control_leaves.iter(); + let mut next_leaf = || control_leaves_iter.next().unwrap(); + + let control_subtree_leaves: Vec> = [ + // Subtree 0. + vec![next_leaf(), next_leaf(), next_leaf()], + // Subtree 1. + vec![next_leaf(), next_leaf()], + // Subtree 2. + vec![next_leaf()], + ] + .map(|subtree| subtree.into_iter().map(smtleaf_to_subtree_leaf).collect()) + .to_vec(); + assert_eq!(control_leaves_iter.next(), None); + control_subtree_leaves + }; + + let subtrees: PairComputations = Smt::sorted_pairs_to_leaves(entries); + // This will check that the hashes, columns, and subtree assignments all match. + assert_eq!(subtrees.leaves, control_subtree_leaves); + + // Flattening and re-separating out the leaves into subtrees should have the same result. + let mut all_leaves: Vec = + subtrees.leaves.clone().into_iter().flatten().collect(); + let re_grouped: Vec> = SubtreeLeavesIter::from_leaves(&mut all_leaves).collect(); + assert_eq!(subtrees.leaves, re_grouped); + + // Then finally we might as well check the computed leaf nodes too. + let control_leaves: BTreeMap = control + .leaves() + .map(|(index, value)| (index.index.value(), value.clone())) + .collect(); + + for (column, test_leaf) in subtrees.nodes { + if test_leaf.is_empty() { + continue; + } + let control_leaf = control_leaves + .get(&column) + .unwrap_or_else(|| panic!("no leaf node found for column {column}")); + assert_eq!(control_leaf, &test_leaf); + } + } +}