From 5990e5fb1be18b8a6e2e8dc4419dcf1c2402ae3f Mon Sep 17 00:00:00 2001 From: Qyriad Date: Thu, 27 Feb 2025 17:26:28 +0100 Subject: [PATCH] smt: add SparseMerklePath --- CHANGELOG.md | 1 + src/merkle/mod.rs | 3 + src/merkle/sparse_path.rs | 115 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+) create mode 100644 src/merkle/sparse_path.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 14fd3a3..24b3a2c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ - [BREAKING] Updated Winterfell dependency to v0.12 (#374). - Added debug-only duplicate column check in `build_subtree` (#378). - Filter out empty values in concurrent version of `Smt::with_entries` to fix a panic (#383). +- Added `SparseMerklePath`, a compact representation of `MerklePath` which compacts empty nodes into a bitmask (#389). ## 0.13.3 (2025-02-18) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 39c58c5..ac305c3 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -20,6 +20,9 @@ pub use merkle_tree::{path_to_text, tree_to_text, MerkleTree}; mod path; pub use path::{MerklePath, RootPath, ValuePath}; +mod sparse_path; +pub use sparse_path::SparseMerklePath; + mod smt; #[cfg(feature = "internal")] pub use smt::{build_subtree_for_bench, SubtreeLeaf}; diff --git a/src/merkle/sparse_path.rs b/src/merkle/sparse_path.rs new file mode 100644 index 0000000..765dc05 --- /dev/null +++ b/src/merkle/sparse_path.rs @@ -0,0 +1,115 @@ +use alloc::vec::Vec; +use core::iter; + +use super::{EmptySubtreeRoots, MerklePath, RpoDigest, SMT_MAX_DEPTH}; + +/// A different representation of [`MerklePath`] designed for memory efficiency for Merkle paths +/// with empty nodes. +/// +/// Empty nodes in the path are stored only as their position, represented with a bitmask. A +/// maximum of 64 nodes in the path can be empty. The number of empty nodes has no effect on memory +/// usage by this struct, but will incur overhead during iteration or conversion to a +/// [`MerklePath`], for each empty node. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] +pub struct SparseMerklePath { + /// A bitmask representing empty nodes. The set bit corresponds to the depth of an empty node. + empty_nodes: u64, + /// The non-empty nodes, stored in depth-order, but not contiguous across depth. + nodes: Vec, +} + +impl SparseMerklePath { + /// Converts a Merkle path to a sparse representation. + /// + /// Returns `None` if `path` contains more elements than we can represent ([`SMT_MAX_DEPTH`]). + pub fn from_path(tree_depth: u8, path: MerklePath) -> Option { + // Note that the path does not include the node itself that it is a path to. + // That is to say, the path is not inclusive of its ending. + + if path.len() > SMT_MAX_DEPTH.into() { + return None; + } + let path_depth: u8 = path.len().try_into().unwrap(); + + let mut nodes: Vec = Default::default(); + let mut empty_nodes: u64 = 0; + + for (depth, node) in iter::zip(0..path_depth, path) { + let &equivalent_empty_node = EmptySubtreeRoots::entry(tree_depth, depth); + if node == equivalent_empty_node { + // FIXME: should we just fallback to the Vec if we're out of bits? + assert!(depth < 64, "SparseMerklePath may have at most 64 empty nodes"); + empty_nodes |= u64::checked_shl(1, depth.into()).unwrap(); + } else { + nodes.push(node); + } + } + + Some(Self { empty_nodes, nodes }) + } + + /// Converts this sparse representation back to a normal [`MerklePath`]. + pub fn into_path(mut self, tree_depth: u8) -> MerklePath { + let path_depth = self.depth(); + let mut nodes: Vec = Default::default(); + let mut sparse_nodes = self.nodes.drain(..); + + for depth in 0..path_depth { + let empty_bit = u64::checked_shl(1, depth.into()).unwrap(); + let is_empty = (self.empty_nodes & empty_bit) != 0; + if is_empty { + let &equivalent_empty_node = EmptySubtreeRoots::entry(tree_depth, depth); + nodes.push(equivalent_empty_node); + } else { + nodes.push(sparse_nodes.next().unwrap()); + } + } + + debug_assert_eq!(sparse_nodes.next(), None); + drop(sparse_nodes); + + debug_assert!(self.nodes.is_empty()); + + MerklePath::from(nodes) + } + + pub fn depth(&self) -> u8 { + (self.nodes.len() + self.empty_nodes.count_ones() as usize) as u8 + } +} + +#[cfg(test)] +mod tests { + use alloc::vec::Vec; + + use super::SparseMerklePath; + use crate::{ + hash::rpo::RpoDigest, + merkle::{smt::SparseMerkleTree, Smt, SMT_DEPTH}, + Felt, Word, ONE, + }; + + #[test] + fn roundtrip() { + let pair_count: u64 = 8192; + 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(); + let tree = Smt::with_entries(entries).unwrap(); + + for (key, _value) in tree.entries() { + 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()); + let test_path = sparse_path.into_path(SMT_DEPTH); + + assert_eq!(control_path, test_path); + } + } +}