smt: impl SimpleSmt::get_path() which returns a sparse path

This commit is contained in:
Qyriad 2025-03-22 20:42:07 +01:00
parent 22b33ed8b6
commit 81ce40ce49
2 changed files with 45 additions and 2 deletions

View file

@ -1,10 +1,11 @@
use alloc::collections::BTreeSet; use alloc::{collections::BTreeSet, vec::Vec};
use super::{ use super::{
super::ValuePath, EMPTY_WORD, EmptySubtreeRoots, InnerNode, InnerNodeInfo, InnerNodes, super::ValuePath, EMPTY_WORD, EmptySubtreeRoots, InnerNode, InnerNodeInfo, InnerNodes,
LeafIndex, MerkleError, MerklePath, MutationSet, NodeIndex, RpoDigest, SMT_MAX_DEPTH, LeafIndex, MerkleError, MerklePath, MutationSet, NodeIndex, RpoDigest, SMT_MAX_DEPTH,
SMT_MIN_DEPTH, SparseMerkleTree, Word, SMT_MIN_DEPTH, SparseMerkleTree, Word,
}; };
use crate::merkle::{SparseMerklePath, sparse_path::SparseValuePath};
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
@ -173,6 +174,40 @@ impl<const DEPTH: u8> SimpleSmt<DEPTH> {
<Self as SparseMerkleTree<DEPTH>>::open(self, key) <Self as SparseMerkleTree<DEPTH>>::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<DEPTH>) -> 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<RpoDigest> = 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. /// Returns a boolean value indicating whether the SMT is empty.
pub fn is_empty(&self) -> bool { pub fn is_empty(&self) -> bool {
debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT); debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);

View file

@ -1,4 +1,5 @@
use alloc::vec::Vec; use alloc::vec::Vec;
use core::iter;
use assert_matches::assert_matches; use assert_matches::assert_matches;
@ -237,9 +238,16 @@ fn small_tree_opening_is_consistent() {
]; ];
for (key, path) in cases { 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); 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);
}
} }
} }