feat: add simple sparse merkle tree

Cargo.toml

@@ -14,6 +14,10 @@ edition = "2021"
 name = "hash"
 harness = false
 
+[[bench]]
+name = "smt"
+harness = false
+
 [features]
 default = ["blake3/default", "std", "winter_crypto/default", "winter_math/default", "winter_utils/default"]
 std = ["blake3/std", "winter_crypto/std", "winter_math/std", "winter_utils/std"]
@@ -25,6 +29,6 @@ winter_math = { version = "0.4.1", package = "winter-math", default-features = f
 winter_utils = { version = "0.4.1", package = "winter-utils", default-features = false }
 
 [dev-dependencies]
-criterion = "0.4"
+criterion = { version = "0.4", features = ["html_reports"] }
 proptest = "1.0.0"
 rand_utils = { version = "0.4", package = "winter-rand-utils" }

benches/smt.rs

@@ -0,0 +1,84 @@
+use core::mem::swap;
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use miden_crypto::{merkle::SimpleSmt, Felt, Word};
+use rand_utils::prng_array;
+
+fn smt_rpo(c: &mut Criterion) {
+    // setup trees
+
+    let mut seed = [0u8; 32];
+    let mut trees = vec![];
+
+    for depth in 14..=20 {
+        let leaves = ((1 << depth) - 1) as u64;
+        for count in [1, leaves / 2, leaves] {
+            let entries: Vec<_> = (0..count)
+                .map(|i| {
+                    let word = generate_word(&mut seed);
+                    (i, word)
+                })
+                .collect();
+            let tree = SimpleSmt::new(entries, depth).unwrap();
+            trees.push(tree);
+        }
+    }
+
+    let leaf = generate_word(&mut seed);
+
+    // benchmarks
+
+    let mut insert = c.benchmark_group(format!("smt update_leaf"));
+
+    for tree in trees.iter_mut() {
+        let depth = tree.depth();
+        let count = tree.leaves_count() as u64;
+        let key = count >> 2;
+        insert.bench_with_input(
+            format!("simple smt(depth:{depth},count:{count})"),
+            &(key, leaf),
+            |b, (key, leaf)| {
+                b.iter(|| {
+                    tree.update_leaf(black_box(*key), black_box(*leaf)).unwrap();
+                });
+            },
+        );
+    }
+
+    insert.finish();
+
+    let mut path = c.benchmark_group(format!("smt get_leaf_path"));
+
+    for tree in trees.iter_mut() {
+        let depth = tree.depth();
+        let count = tree.leaves_count() as u64;
+        let key = count >> 2;
+        path.bench_with_input(
+            format!("simple smt(depth:{depth},count:{count})"),
+            &key,
+            |b, key| {
+                b.iter(|| {
+                    tree.get_leaf_path(black_box(*key)).unwrap();
+                });
+            },
+        );
+    }
+
+    path.finish();
+}
+
+criterion_group!(smt_group, smt_rpo);
+criterion_main!(smt_group);
+
+// HELPER FUNCTIONS
+// --------------------------------------------------------------------------------------------
+
+fn generate_word(seed: &mut [u8; 32]) -> Word {
+    swap(seed, &mut prng_array(*seed));
+    let nums: [u64; 4] = prng_array(*seed);
+    [
+        Felt::new(nums[0]),
+        Felt::new(nums[1]),
+        Felt::new(nums[2]),
+        Felt::new(nums[3]),
+    ]
+}

src/lib.rs

@@ -23,11 +23,14 @@ pub mod utils {
 // ================================================================================================
 
 /// A group of four field elements in the Miden base field.
-pub type Word = [Felt; 4];
+pub type Word = [Felt; WORD_SIZE];
 
 // CONSTANTS
 // ================================================================================================
 
+/// Number of field elements in a word.
+pub const WORD_SIZE: usize = 4;
+
 /// Field element representing ZERO in the Miden base filed.
 pub const ZERO: Felt = Felt::ZERO;
 

src/merkle/merkle_tree.rs

@@ -1,4 +1,4 @@
-use super::{Digest, Felt, MerkleError, Rpo256, Vec, Word};
+use super::{Felt, MerkleError, Rpo256, RpoDigest, Vec, Word};
 use crate::{utils::uninit_vector, FieldElement};
 use core::slice;
 use winter_math::log2;
@@ -22,7 +22,7 @@ impl MerkleTree {
     pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> {
         let n = leaves.len();
         if n <= 1 {
-            return Err(MerkleError::DepthTooSmall);
+            return Err(MerkleError::DepthTooSmall(n as u32));
         } else if !n.is_power_of_two() {
             return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
         }
@@ -35,7 +35,8 @@ impl MerkleTree {
         nodes[n..].copy_from_slice(&leaves);
 
         // re-interpret nodes as an array of two nodes fused together
-        let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [Digest; 2], n) };
+        let two_nodes =
+            unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [RpoDigest; 2], n) };
 
         // calculate all internal tree nodes
         for i in (1..n).rev() {
@@ -68,7 +69,7 @@ impl MerkleTree {
     /// * The specified index not valid for the specified depth.
     pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> {
         if depth == 0 {
-            return Err(MerkleError::DepthTooSmall);
+            return Err(MerkleError::DepthTooSmall(depth));
         } else if depth > self.depth() {
             return Err(MerkleError::DepthTooBig(depth));
         }
@@ -89,7 +90,7 @@ impl MerkleTree {
     /// * The specified index not valid for the specified depth.
     pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
         if depth == 0 {
-            return Err(MerkleError::DepthTooSmall);
+            return Err(MerkleError::DepthTooSmall(depth));
         } else if depth > self.depth() {
             return Err(MerkleError::DepthTooBig(depth));
         }
@@ -123,7 +124,7 @@ impl MerkleTree {
 
         let n = self.nodes.len() / 2;
         let two_nodes =
-            unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [Digest; 2], n) };
+            unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [RpoDigest; 2], n) };
 
         for _ in 0..depth {
             index /= 2;

src/merkle/mod.rs

@@ -1,29 +1,61 @@
 use super::{
-    hash::rpo::{Rpo256, RpoDigest as Digest},
+    hash::rpo::{Rpo256, RpoDigest},
     utils::collections::{BTreeMap, Vec},
     Felt, Word, ZERO,
 };
+use core::fmt;
 
 mod merkle_tree;
 pub use merkle_tree::MerkleTree;
 
 mod merkle_path_set;
 pub use merkle_path_set::MerklePathSet;
 
+mod simple_smt;
+pub use simple_smt::SimpleSmt;
+
 // ERRORS
 // ================================================================================================
 
 #[derive(Clone, Debug)]
 pub enum MerkleError {
-    DepthTooSmall,
+    DepthTooSmall(u32),
     DepthTooBig(u32),
     NumLeavesNotPowerOfTwo(usize),
     InvalidIndex(u32, u64),
     InvalidDepth(u32, u32),
     InvalidPath(Vec<Word>),
+    InvalidEntriesCount(usize, usize),
     NodeNotInSet(u64),
 }
 
+impl fmt::Display for MerkleError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use MerkleError::*;
+        match self {
+            DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"),
+            DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"),
+            NumLeavesNotPowerOfTwo(leaves) => {
+                write!(f, "the leaves count {leaves} is not a power of 2")
+            }
+            InvalidIndex(depth, index) => write!(
+                f,
+                "the leaf index {index} is not valid for the depth {depth}"
+            ),
+            InvalidDepth(expected, provided) => write!(
+                f,
+                "the provided depth {provided} is not valid for {expected}"
+            ),
+            InvalidPath(_path) => write!(f, "the provided path is not valid"),
+            InvalidEntriesCount(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
+            NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for MerkleError {}
+
 // HELPER FUNCTIONS
 // ================================================================================================