Added support for hashmaps in Smt and SimpleSmt

.github/workflows/test.yml

@@ -17,7 +17,7 @@ jobs:
       matrix:
         toolchain: [stable, nightly]
         os: [ubuntu]
-        args: [default, no-std]
+        args: [default, smt-hashmaps, no-std]
     timeout-minutes: 30
     steps:
       - uses: actions/checkout@main

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Fixed a bug in the implementation of `draw_integers` for `RpoRandomCoin` (#343).
 - [BREAKING] Refactor error messages and use `thiserror` to derive errors (#344).
 - [BREAKING] Updated Winterfell dependency to v0.11 (#346).
+- Added support for hashmaps in `Smt` and `SimpleSmt` which gives up to 10x boost in some operations (#363).
 
 
 ## 0.12.0 (2024-10-30)

Cargo.toml

@@ -48,6 +48,7 @@ harness = false
 concurrent = ["dep:rayon"]
 default = ["std", "concurrent"]
 executable = ["dep:clap", "dep:rand-utils", "std"]
+smt_hashmaps = ["dep:hashbrown"]
 internal = []
 serde = ["dep:serde", "serde?/alloc", "winter-math/serde"]
 std = [
@@ -63,6 +64,7 @@ std = [
 [dependencies]
 blake3 = { version = "1.5", default-features = false }
 clap = { version = "4.5", optional = true, features = ["derive"] }
+hashbrown = { version = "0.15", optional = true, features = ["serde"] }
 num = { version = "0.4", default-features = false, features = ["alloc", "libm"] }
 num-complex = { version = "0.4", default-features = false }
 rand = { version = "0.8", default-features = false }

Makefile

@@ -46,14 +46,17 @@ doc: ## Generate and check documentation
 test-default: ## Run tests with default features
 	$(DEBUG_OVERFLOW_INFO) cargo nextest run --profile default --release --all-features
 
+.PHONY: test-smt-hashmaps
+test-smt-hashmaps: ## Run tests with `smt_hashmaps` feature enabled
+	$(DEBUG_OVERFLOW_INFO) cargo nextest run --profile default --release --features smt_hashmaps
 
 .PHONY: test-no-std
 test-no-std: ## Run tests with `no-default-features` (std)
 	$(DEBUG_OVERFLOW_INFO) cargo nextest run --profile default --release --no-default-features
 
 
 .PHONY: test
-test: test-default test-no-std ## Run all tests
+test: test-default test-smt-hashmaps test-no-std ## Run all tests
 
 # --- checking ------------------------------------------------------------------------------------
 

README.md

@@ -63,6 +63,7 @@ This crate can be compiled with the following features:
 - `concurrent`- enabled by default; enables multi-threaded implementation of `Smt::with_entries()` which significantly improves performance on multi-core CPUs.
 - `std` - enabled by default and relies on the Rust standard library.
 - `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly.
+- `smt_hashmaps` - uses hashbrown hashmaps in SMT implementation which significantly improves performance of SMT updating. Keys ordering in SMT iterators is not guarantied when this feature is enabled.
 
 All of these features imply the use of [alloc](https://doc.rust-lang.org/alloc/) to support heap-allocated collections.
 

src/hash/rescue/rpo/digest.rs

@@ -1,5 +1,11 @@
 use alloc::string::String;
-use core::{cmp::Ordering, fmt::Display, ops::Deref, slice};
+use core::{
+    cmp::Ordering,
+    fmt::Display,
+    hash::{Hash, Hasher},
+    ops::Deref,
+    slice,
+};
 
 use thiserror::Error;
 
@@ -55,6 +61,12 @@ impl RpoDigest {
     }
 }
 
+impl Hash for RpoDigest {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        state.write(&self.as_bytes());
+    }
+}
+
 impl Digest for RpoDigest {
     fn as_bytes(&self) -> [u8; DIGEST_BYTES] {
         let mut result = [0; DIGEST_BYTES];

src/merkle/node.rs

@@ -3,6 +3,7 @@ use super::RpoDigest;
 /// Representation of a node with two children used for iterating over containers.
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+#[cfg_attr(test, derive(PartialOrd, Ord))]
 pub struct InnerNodeInfo {
     pub value: RpoDigest,
     pub left: RpoDigest,

src/merkle/smt/full/mod.rs

@@ -1,12 +1,8 @@
-use alloc::{
-    collections::{BTreeMap, BTreeSet},
-    string::ToString,
-    vec::Vec,
-};
+use alloc::{collections::BTreeSet, string::ToString, vec::Vec};
 
 use super::{
-    EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath,
-    MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
+    EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, InnerNodes, LeafIndex, MerkleError,
+    MerklePath, MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
 };
 
 mod error;
@@ -30,6 +26,8 @@ pub const SMT_DEPTH: u8 = 64;
 // SMT
 // ================================================================================================
 
+type Leaves = super::Leaves<SmtLeaf>;
+
 /// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
 /// by 4 field elements.
 ///
@@ -43,8 +41,8 @@ pub const SMT_DEPTH: u8 = 64;
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct Smt {
     root: RpoDigest,
-    leaves: BTreeMap<u64, SmtLeaf>,
-    inner_nodes: BTreeMap<NodeIndex, InnerNode>,
+    inner_nodes: InnerNodes,
+    leaves: Leaves,
 }
 
 impl Smt {
@@ -64,8 +62,8 @@ impl Smt {
 
         Self {
             root,
-            leaves: BTreeMap::new(),
-            inner_nodes: BTreeMap::new(),
+            inner_nodes: Default::default(),
+            leaves: Default::default(),
         }
     }
 
@@ -148,11 +146,7 @@ impl Smt {
     /// # Panics
     /// With debug assertions on, this function panics if `root` does not match the root node in
     /// `inner_nodes`.
-    pub fn from_raw_parts(
-        inner_nodes: BTreeMap<NodeIndex, InnerNode>,
-        leaves: BTreeMap<u64, SmtLeaf>,
-        root: RpoDigest,
-    ) -> Self {
+    pub fn from_raw_parts(inner_nodes: InnerNodes, leaves: Leaves, root: RpoDigest) -> Self {
         // Our particular implementation of `from_raw_parts()` never returns `Err`.
         <Self as SparseMerkleTree<SMT_DEPTH>>::from_raw_parts(inner_nodes, leaves, root).unwrap()
     }
@@ -339,8 +333,8 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
     const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
 
     fn from_raw_parts(
-        inner_nodes: BTreeMap<NodeIndex, InnerNode>,
-        leaves: BTreeMap<u64, SmtLeaf>,
+        inner_nodes: InnerNodes,
+        leaves: Leaves,
         root: RpoDigest,
     ) -> Result<Self, MerkleError> {
         if cfg!(debug_assertions) {

src/merkle/smt/full/tests.rs

@@ -1,9 +1,9 @@
-use alloc::{collections::BTreeMap, vec::Vec};
+use alloc::vec::Vec;
 
 use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH};
 use crate::{
     merkle::{
-        smt::{NodeMutation, SparseMerkleTree},
+        smt::{NodeMutation, SparseMerkleTree, UnorderedMap},
         EmptySubtreeRoots, MerkleStore, MutationSet,
     },
     utils::{Deserializable, Serializable},
@@ -420,7 +420,7 @@ fn test_prospective_insertion() {
     assert_eq!(revert.root(), root_empty, "reverse mutations new root did not match");
     assert_eq!(
         revert.new_pairs,
-        BTreeMap::from_iter([(key_1, EMPTY_WORD)]),
+        UnorderedMap::from_iter([(key_1, EMPTY_WORD)]),
         "reverse mutations pairs did not match"
     );
     assert_eq!(
@@ -440,7 +440,7 @@ fn test_prospective_insertion() {
     assert_eq!(revert.root(), old_root, "reverse mutations new root did not match");
     assert_eq!(
         revert.new_pairs,
-        BTreeMap::from_iter([(key_2, EMPTY_WORD), (key_3, EMPTY_WORD)]),
+        UnorderedMap::from_iter([(key_2, EMPTY_WORD), (key_3, EMPTY_WORD)]),
         "reverse mutations pairs did not match"
     );
 
@@ -454,7 +454,7 @@ fn test_prospective_insertion() {
     assert_eq!(revert.root(), old_root, "reverse mutations new root did not match");
     assert_eq!(
         revert.new_pairs,
-        BTreeMap::from_iter([(key_3, value_3)]),
+        UnorderedMap::from_iter([(key_3, value_3)]),
         "reverse mutations pairs did not match"
     );
 
@@ -474,7 +474,7 @@ fn test_prospective_insertion() {
     assert_eq!(revert.root(), old_root, "reverse mutations new root did not match");
     assert_eq!(
         revert.new_pairs,
-        BTreeMap::from_iter([(key_1, value_1), (key_2, value_2), (key_3, value_3)]),
+        UnorderedMap::from_iter([(key_1, value_1), (key_2, value_2), (key_3, value_3)]),
         "reverse mutations pairs did not match"
     );
 
@@ -603,21 +603,21 @@ fn test_smt_get_value() {
 /// Tests that `entries()` works as expected
 #[test]
 fn test_smt_entries() {
-    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
-    let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
+    let key_1 = RpoDigest::from([ONE, ONE, ONE, ONE]);
+    let key_2 = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
 
     let value_1 = [ONE; WORD_SIZE];
     let value_2 = [2_u32.into(); WORD_SIZE];
+    let entries = [(key_1, value_1), (key_2, value_2)];
 
-    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
+    let smt = Smt::with_entries(entries).unwrap();
 
-    let mut entries = smt.entries();
+    let mut expected = Vec::from_iter(entries);
+    expected.sort_by_key(|(k, _)| *k);
+    let mut actual: Vec<_> = smt.entries().cloned().collect();
+    actual.sort_by_key(|(k, _)| *k);
 
-    // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
-    // switches the order, it is OK to modify the order here as well.
-    assert_eq!(&(key_1, value_1), entries.next().unwrap());
-    assert_eq!(&(key_2, value_2), entries.next().unwrap());
-    assert!(entries.next().is_none());
+    assert_eq!(actual, expected);
 }
 
 /// Tests that `EMPTY_ROOT` constant generated in the `Smt` equals to the root of the empty tree of