From b534ab507a4a9adf5a14afa7c4a90839ba91d722 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Mon, 28 Oct 2024 15:38:42 -0600 Subject: [PATCH] add sorted_pairs_to_leaves() and test for it --- src/merkle/smt/full/mod.rs | 24 +++++++++ src/merkle/smt/mod.rs | 100 ++++++++++++++++++++++++++++++++++- src/merkle/smt/simple/mod.rs | 7 +++ 3 files changed, 130 insertions(+), 1 deletion(-) diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 1be483a..89d0702 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -351,6 +351,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 063003b..9a82708 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -341,11 +341,54 @@ 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` fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening; + fn sorted_pairs_to_leaves( + pairs: Vec<(Self::Key, Self::Value)>, + ) -> PrecomputedLeaves { + let mut all_leaves = PrecomputedLeaves::default(); + + let mut 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 next_col = iter.peek().map(|(key, _)| { + let index = Self::key_to_leaf_index(key); + index.index.value() + }); + + buffer.push((key, value)); + + if let Some(next_col) = next_col { + assert!(next_col >= col); + } + + if next_col == Some(col) { + // Keep going in our buffer. + continue; + } + + // Whether the next pair is a different column, or non-existent, we break off. + let leaf_pairs = mem::take(&mut buffer); + let leaf = Self::pairs_to_leaf(leaf_pairs); + let hash = Self::hash_leaf(&leaf); + + all_leaves.nodes.insert(col, leaf); + all_leaves.subtrees.push(SubtreeLeaf { col, hash }); + } + assert_eq!(buffer.len(), 0); + + all_leaves + } + /// Builds Merkle nodes from a bottom layer of tuples of horizontal indices and their hashes, /// sorted by their position. /// @@ -572,6 +615,24 @@ pub struct SubtreeLeaf { pub hash: RpoDigest, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PrecomputedLeaves { + /// 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 subtrees: Vec, +} + +// Derive requires `L` to impl Default, even though we don't actually need that. +impl Default for PrecomputedLeaves { + fn default() -> Self { + Self { + nodes: Default::default(), + subtrees: Default::default(), + } + } +} + // TESTS // ================================================================================================ #[cfg(test)] @@ -582,9 +643,46 @@ mod test { use crate::{ hash::rpo::RpoDigest, merkle::{smt::SubtreeLeaf, Smt, SmtLeaf, SMT_DEPTH}, - Felt, Word, ONE, + Felt, Word, EMPTY_WORD, ONE, }; + #[test] + fn test_sorted_pairs_to_leaves() { + let entries: Vec<(RpoDigest, Word)> = vec![ + (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]), + // Normal single leaf again. + (RpoDigest::new([ONE, ONE, ONE, Felt::new(400)]), [ONE; 4]), + // Empty leaf. + (RpoDigest::new([ONE, ONE, ONE, Felt::new(500)]), EMPTY_WORD), + ]; + let mut entries_iter = entries.iter().cloned(); + let mut next_entry = || entries_iter.next().unwrap(); + + let control_leaves: Vec = vec![ + SmtLeaf::Single(next_entry()), + SmtLeaf::Single(next_entry()), + SmtLeaf::new_multiple(vec![next_entry(), next_entry()]).unwrap(), + SmtLeaf::Single(next_entry()), + SmtLeaf::new_empty(Smt::key_to_leaf_index(&next_entry().0)), + ]; + + let control_subtree_leaves: Vec = control_leaves + .iter() + .map(|leaf| { + let col = leaf.index().index.value(); + let hash = leaf.hash(); + SubtreeLeaf { col, hash } + }) + .collect(); + + let test_subtree_leaves = Smt::sorted_pairs_to_leaves(entries).subtrees; + assert_eq!(control_subtree_leaves, test_subtree_leaves); + } + #[test] fn test_build_subtree_from_leaves() { const PAIR_COUNT: u64 = u64::pow(2, 8); diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index edd9849..04476a0 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -373,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 + } }