From 81ce40ce497f1ef24b042051ef2a4e5f499c2497 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
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 +++++-
 2 files changed, 45 insertions(+), 2 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<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);
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);
+        }
     }
 }