diff --git a/src/merkle/sparse_path.rs b/src/merkle/sparse_path.rs index 765dc05..cb97903 100644 --- a/src/merkle/sparse_path.rs +++ b/src/merkle/sparse_path.rs @@ -77,6 +77,44 @@ impl SparseMerklePath { pub fn depth(&self) -> u8 { (self.nodes.len() + self.empty_nodes.count_ones() as usize) as u8 } + + /// Get a specific node in this path at a given depth. + /// + /// # Panics + /// With debug assertions enabled, this function panics if `node_depth` is greater than + /// `tree_depth` (as it is impossible to have a node of greater depth than the tree it is + /// contained in). + pub fn get(&self, tree_depth: u8, node_depth: u8) -> Option { + if node_depth == tree_depth || node_depth > self.depth() { + return None; + } + + debug_assert!( + tree_depth >= node_depth, + "tree depth {tree_depth} must be greater than node depth {node_depth}", + ); + + let empty_bit = u64::checked_shl(1, node_depth.into()).unwrap(); + let is_empty = (self.empty_nodes & empty_bit) != 0; + + if is_empty { + return Some(*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() + }; + Some(self.nodes[nonempty_index]) + } } #[cfg(test)] @@ -90,6 +128,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; @@ -112,4 +163,20 @@ mod tests { assert_eq!(control_path, test_path); } } + + #[test] + 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()).unwrap(); + 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).unwrap(); + assert_eq!(*control_node, test_node, "at depth {depth} for entry {i}"); + } + } + } }