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

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<const DEPTH: u8> SimpleSmt<DEPTH> {
         <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.
     pub fn is_empty(&self) -> bool {
         debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);

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);
+        }
     }
 }
 

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<RpoDigest>) -> Self {
+        Self { empty_nodes, nodes }
+    }
+
+    /// Decomposes a [SparseMerklePath] into its raw components: `(empty_nodes, nodes)`.
+    pub fn into_raw_parts(self) -> (u64, Vec<RpoDigest>) {
+        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
 // ================================================================================================