From 41f72ecf2f0547febbfa18de0ed854725e5925e8 Mon Sep 17 00:00:00 2001 From: Qyriad Date: Sat, 22 Mar 2025 20:42:07 +0100 Subject: [PATCH] smt: impl SimpleSmt::get_path() which returns a sparse path --- miden-crypto/src/merkle/smt/simple/mod.rs | 37 ++++++++++++++++++- miden-crypto/src/merkle/smt/simple/tests.rs | 10 ++++- miden-crypto/src/merkle/sparse_path.rs | 41 ++++++++++++++++++++- 3 files changed, 85 insertions(+), 3 deletions(-) diff --git a/miden-crypto/src/merkle/smt/simple/mod.rs b/miden-crypto/src/merkle/smt/simple/mod.rs index 6773cd1..737b335 100644 --- a/miden-crypto/src/merkle/smt/simple/mod.rs +++ b/miden-crypto/src/merkle/smt/simple/mod.rs @@ -1,10 +1,11 @@ -use alloc::collections::BTreeSet; +use alloc::{collections::BTreeSet, vec::Vec}; use super::{ super::ValuePath, EMPTY_WORD, EmptySubtreeRoots, InnerNode, InnerNodeInfo, InnerNodes, LeafIndex, MerkleError, MerklePath, MutationSet, NodeIndex, RpoDigest, SMT_MAX_DEPTH, SMT_MIN_DEPTH, SparseMerkleTree, Word, }; +use crate::merkle::{SparseMerklePath, sparse_path::SparseValuePath}; #[cfg(test)] mod tests; @@ -173,6 +174,40 @@ impl SimpleSmt { >::open(self, key) } + /// Returns a path (but not an opening) to the leaf associated with `key`. + /// + /// Unlike [`SimpleSmt::open()`], this returns a [SparseValuePath] which has a more efficient + /// memory representation optimized for paths containing empty nodes. See [SparseMerklePath] + /// for more information. + pub fn get_path(&self, key: &LeafIndex) -> SparseValuePath { + let value: RpoDigest = self.get_value(key).into(); + + // This is a partial re-implementation of `SparseMerklePath::from_sized_iter()`, which + // constructs in place instead of cloning and immediately dropping the entire vec returned + // by `self.open()`. + + let mut nodes: Vec = Default::default(); + let mut index = NodeIndex::from(*key); + let mut empty_nodes: u64 = 0; + + for _ in 0..DEPTH { + let is_right = index.is_value_odd(); + index.move_up(); + + match self.inner_nodes.get(&index) { + Some(InnerNode { left, right }) => { + let value = if is_right { left } else { right }; + nodes.push(*value); + }, + None => empty_nodes |= u64::checked_shl(1, index.depth().into()).unwrap(), + } + } + + let path = SparseMerklePath::from_raw_parts(empty_nodes, nodes); + + SparseValuePath { value, path } + } + /// Returns a boolean value indicating whether the SMT is empty. pub fn is_empty(&self) -> bool { debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT); diff --git a/miden-crypto/src/merkle/smt/simple/tests.rs b/miden-crypto/src/merkle/smt/simple/tests.rs index 919fed2..a24ba5b 100644 --- a/miden-crypto/src/merkle/smt/simple/tests.rs +++ b/miden-crypto/src/merkle/smt/simple/tests.rs @@ -1,4 +1,5 @@ use alloc::vec::Vec; +use core::iter; use assert_matches::assert_matches; @@ -237,9 +238,16 @@ fn small_tree_opening_is_consistent() { ]; for (key, path) in cases { - let opening = tree.open(&LeafIndex::<3>::new(key).unwrap()); + let index = LeafIndex::<3>::new(key).unwrap(); + let opening = tree.open(&index); assert_eq!(path, *opening.path); + + // Also check that the sparse versions work the same way. + let sparse_path = tree.get_path(&index); + for (path_node, sparse_path_node) in iter::zip(path, sparse_path.path) { + assert_eq!(path_node, sparse_path_node); + } } } diff --git a/miden-crypto/src/merkle/sparse_path.rs b/miden-crypto/src/merkle/sparse_path.rs index eb9fd31..dbce1d4 100644 --- a/miden-crypto/src/merkle/sparse_path.rs +++ b/miden-crypto/src/merkle/sparse_path.rs @@ -3,7 +3,7 @@ use core::iter; use winter_utils::{Deserializable, DeserializationError, Serializable}; -use super::{EmptySubtreeRoots, MerkleError, MerklePath, RpoDigest, SMT_MAX_DEPTH}; +use super::{EmptySubtreeRoots, MerkleError, MerklePath, RpoDigest, SMT_MAX_DEPTH, Word}; /// A different representation of [`MerklePath`] designed for memory efficiency for Merkle paths /// with empty nodes. @@ -87,6 +87,18 @@ impl SparseMerklePath { MerklePath::from(nodes) } + /// Creates a [SparseMerklePath] directly from a bitmask representing empty nodes, and a vec + /// containing all non-empty nodes. + pub fn from_raw_parts(empty_nodes: u64, nodes: Vec) -> Self { + Self { empty_nodes, nodes } + } + + /// Decomposes a [SparseMerklePath] into its raw components: `(empty_nodes, nodes)`. + pub fn into_raw_parts(self) -> (u64, Vec) { + let SparseMerklePath { empty_nodes, nodes } = self; + (empty_nodes, nodes) + } + /// Returns the total depth of this path, i.e., the number of nodes this path represents. pub fn depth(&self) -> u8 { (self.nodes.len() + self.empty_nodes.count_ones() as usize) as u8 @@ -254,6 +266,33 @@ impl DoubleEndedIterator for SparseMerkleIter { } } +// SPARSE MERKLE PATH CONTAINERS +// ================================================================================================ +/// A container for a [crate::Word] value and its [SparseMerklePath] opening. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SparseValuePath { + /// The node value opening for `path`. + pub value: RpoDigest, + /// The path from `value` to `root` (exclusive), using an efficient memory representation for + /// empty nodes. + pub path: SparseMerklePath, +} + +impl SparseValuePath { + /// Convenience function to construct a [SparseValuePath]. + /// + /// `value` is the value `path` leads to, in the tree. + pub fn new(value: RpoDigest, path: SparseMerklePath) -> Self { + Self { value, path } + } +} + +impl From<(SparseMerklePath, Word)> for SparseValuePath { + fn from((path, value): (SparseMerklePath, Word)) -> Self { + SparseValuePath::new(value.into(), path) + } +} + // SERIALIZATION // ================================================================================================