From 4939eb6720039eaf7c1c12678d5473b9cb931437 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Fri, 21 Mar 2025 12:18:08 +0100
Subject: [PATCH] SparseMerklePath: implement random access

---
 src/merkle/sparse_path.rs | 56 ++++++++++++++++++++++++++++++---------
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/src/merkle/sparse_path.rs b/src/merkle/sparse_path.rs
index f958c69..a7dece8 100644
--- a/src/merkle/sparse_path.rs
+++ b/src/merkle/sparse_path.rs
@@ -78,20 +78,28 @@ impl SparseMerklePath {
         (self.nodes.len() + self.empty_nodes.count_ones() as usize) as u8
     }
 
+    /// Get a specific node in this path at a given depth.
     pub fn get(&self, tree_depth: u8, node_depth: u8) -> RpoDigest {
-        let empty_bit = u64::checked_shl(1, node_depth as u32).unwrap();
+        let empty_bit = u64::checked_shl(1, node_depth.into()).unwrap();
         let is_empty = (self.empty_nodes & empty_bit) != 0;
+
         if is_empty {
-            *EmptySubtreeRoots::entry(tree_depth, node_depth)
-        } else {
-            // Get rid of 1s that are more significant than the one that indicates our empty bit.
-            let mask = u64::unbounded_shl(u64::MAX, node_depth as u32);
-            let empty_before = u64::count_ones(self.empty_nodes & mask);
-            std::dbg!(node_depth, empty_before);
-            //let index = node_depth - empty_before as u8;
-            let index = empty_before as u8 - node_depth;
-            self.nodes[index as usize]
+            return *EmptySubtreeRoots::entry(tree_depth, node_depth);
         }
+
+        // Our index needs to account for all the empty nodes that aren't in `self.nodes`.
+        let nonempty_index: usize = {
+            // TODO: this could also be u64::unbounded_shl(1, node_depth + 1).wrapping_sub(1).
+            // We should check if that has any performance benefits over using 128-bit integers.
+            let mask: u64 = ((1u128 << (node_depth + 1)) - 1u128).try_into().unwrap();
+
+            let empty_before = u64::count_ones(self.empty_nodes & mask);
+            u64::checked_sub(node_depth as u64, empty_before as u64)
+                .unwrap()
+                .try_into()
+                .unwrap()
+        };
+        self.nodes[nonempty_index]
     }
 }
 
@@ -156,6 +164,19 @@ mod tests {
         Felt, Word, ONE,
     };
 
+    fn make_smt(pair_count: u64) -> Smt {
+        let entries: Vec<(RpoDigest, Word)> = (0..pair_count)
+            .map(|n| {
+                let leaf_index = ((n as f64 / pair_count as f64) * 255.0) as u64;
+                let key = RpoDigest::new([ONE, ONE, Felt::new(n), Felt::new(leaf_index)]);
+                let value = [ONE, ONE, ONE, ONE];
+                (key, value)
+            })
+            .collect();
+
+        Smt::with_entries(entries).unwrap()
+    }
+
     #[test]
     fn roundtrip() {
         let pair_count: u64 = 8192;
@@ -179,7 +200,18 @@ mod tests {
     }
 
     #[test]
-    fn get() {
-        todo!();
+    fn random_access() {
+        let tree = make_smt(8192);
+
+        for (i, (key, _value)) in tree.entries().enumerate() {
+            let control_path = tree.path(key);
+            let sparse_path = SparseMerklePath::from_path(SMT_DEPTH, control_path.clone());
+            assert_eq!(control_path.depth(), sparse_path.depth());
+
+            for (depth, control_node) in control_path.iter().enumerate() {
+                let test_node = sparse_path.get(SMT_DEPTH, depth as u8);
+                assert_eq!(*control_node, test_node, "at depth {depth} for entry {i}");
+            }
+        }
     }
 }