From bafe5c156d2f51bd71ffe6175f9603bdc818b6eb Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Thu, 19 Dec 2024 19:43:51 +0500 Subject: [PATCH 01/16] feat: rewrite `Smt` and `SimpleSmt` to hashmaps --- Cargo.lock | 31 ++++++++ Cargo.toml | 1 + src/hash/rescue/rpo/digest.rs | 7 ++ src/merkle/node.rs | 2 +- src/merkle/smt/full/mod.rs | 24 +++--- src/merkle/smt/mod.rs | 137 +++++++++++++++++++++++++++++---- src/merkle/smt/simple/mod.rs | 23 +++--- src/merkle/smt/simple/tests.rs | 8 +- 8 files changed, 190 insertions(+), 43 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 00e6328d..e22cc7f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "anes" version = "0.1.6" @@ -349,6 +355,12 @@ version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +[[package]] +name = "equivalent" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" + [[package]] name = "errno" version = "0.3.9" @@ -371,6 +383,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f81ec6369c545a7d40e4589b5597581fa1c441fe1cce96dd1de43159910a36a2" + [[package]] name = "generic-array" version = "0.14.7" @@ -410,6 +428,18 @@ dependencies = [ "crunchy", ] +[[package]] +name = "hashbrown" +version = "0.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", + "serde", +] + [[package]] name = "heck" version = "0.5.0" @@ -534,6 +564,7 @@ dependencies = [ "criterion", "getrandom", "glob", + "hashbrown", "hex", "num", "num-complex", diff --git a/Cargo.toml b/Cargo.toml index 46ac3bdd..f7b8c7f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,6 +63,7 @@ std = [ [dependencies] blake3 = { version = "1.5", default-features = false } clap = { version = "4.5", optional = true, features = ["derive"] } +hashbrown = { version = "0.15", 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 } diff --git a/src/hash/rescue/rpo/digest.rs b/src/hash/rescue/rpo/digest.rs index 44663691..2206aab6 100644 --- a/src/hash/rescue/rpo/digest.rs +++ b/src/hash/rescue/rpo/digest.rs @@ -1,5 +1,6 @@ use alloc::string::String; use core::{cmp::Ordering, fmt::Display, ops::Deref, slice}; +use std::hash::{Hash, Hasher}; use thiserror::Error; @@ -55,6 +56,12 @@ impl RpoDigest { } } +impl Hash for RpoDigest { + fn hash(&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]; diff --git a/src/merkle/node.rs b/src/merkle/node.rs index bf18d386..ed92d7cc 100644 --- a/src/merkle/node.rs +++ b/src/merkle/node.rs @@ -1,7 +1,7 @@ use super::RpoDigest; /// Representation of a node with two children used for iterating over containers. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] pub struct InnerNodeInfo { pub value: RpoDigest, diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 56c08c66..6cc97a6c 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -1,8 +1,6 @@ -use alloc::{ - collections::{BTreeMap, BTreeSet}, - string::ToString, - vec::Vec, -}; +use alloc::{collections::BTreeSet, string::ToString, vec::Vec}; + +use hashbrown::HashMap; use super::{ EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, @@ -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, - inner_nodes: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, } impl Smt { @@ -64,8 +62,8 @@ impl Smt { Self { root, - leaves: BTreeMap::new(), - inner_nodes: BTreeMap::new(), + inner_nodes: Default::default(), + leaves: Default::default(), } } @@ -149,8 +147,8 @@ impl Smt { /// 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, - leaves: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, root: RpoDigest, ) -> Self { // Our particular implementation of `from_raw_parts()` never returns `Err`. @@ -317,8 +315,8 @@ impl SparseMerkleTree for Smt { const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); fn from_raw_parts( - inner_nodes: BTreeMap, - leaves: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, root: RpoDigest, ) -> Result { if cfg!(debug_assertions) { diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 42020fcf..62c061c9 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -1,7 +1,10 @@ use alloc::{collections::BTreeMap, vec::Vec}; use core::mem; +use std::hash::Hash; +use hashbrown::HashMap; use num::Integer; +use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex}; use crate::{ @@ -48,7 +51,7 @@ pub const SMT_MAX_DEPTH: u8 = 64; /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. pub(crate) trait SparseMerkleTree { /// The type for a key - type Key: Clone + Ord; + type Key: Clone + Ord + Eq + Hash; /// The type for a value type Value: Clone + PartialEq; /// The type for a leaf @@ -172,8 +175,8 @@ pub(crate) trait SparseMerkleTree { use NodeMutation::*; let mut new_root = self.root(); - let mut new_pairs: BTreeMap = Default::default(); - let mut node_mutations: BTreeMap = Default::default(); + let mut new_pairs: HashMap = Default::default(); + let mut node_mutations: HashMap = Default::default(); for (key, value) in kv_pairs { // If the old value and the new value are the same, there is nothing to update. @@ -309,8 +312,8 @@ pub(crate) trait SparseMerkleTree { /// Construct this type from already computed leaves and nodes. The caller ensures passed /// arguments are correct and consistent with each other. fn from_raw_parts( - inner_nodes: BTreeMap, - leaves: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, root: RpoDigest, ) -> Result where @@ -441,7 +444,7 @@ pub(crate) trait SparseMerkleTree { #[cfg(feature = "concurrent")] fn build_subtrees( mut entries: Vec<(Self::Key, Self::Value)>, - ) -> (BTreeMap, BTreeMap) { + ) -> (HashMap, HashMap) { entries.sort_by_key(|item| { let index = Self::key_to_leaf_index(&item.0); index.value() @@ -456,10 +459,10 @@ pub(crate) trait SparseMerkleTree { #[cfg(feature = "concurrent")] fn build_subtrees_from_sorted_entries( entries: Vec<(Self::Key, Self::Value)>, - ) -> (BTreeMap, BTreeMap) { + ) -> (HashMap, HashMap) { use rayon::prelude::*; - let mut accumulated_nodes: BTreeMap = Default::default(); + let mut accumulated_nodes: HashMap = Default::default(); let PairComputations { leaves: mut leaf_subtrees, @@ -576,7 +579,7 @@ pub(crate) enum NodeMutation { /// Represents a group of prospective mutations to a `SparseMerkleTree`, created by /// `SparseMerkleTree::compute_mutations()`, and that can be applied with /// `SparseMerkleTree::apply_mutations()`. -#[derive(Debug, Clone, PartialEq, Eq, Default)] +#[derive(Debug, Clone, Default)] pub struct MutationSet { /// The root of the Merkle tree this MutationSet is for, recorded at the time /// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying @@ -587,12 +590,12 @@ pub struct MutationSet { /// index overlayed, if any. Each [`NodeMutation::Addition`] corresponds to a /// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`] /// corresponds to a [`SparseMerkleTree::remove_inner_node()`] call. - node_mutations: BTreeMap, + node_mutations: HashMap, /// The set of top-level key-value pairs we're prospectively adding to the tree, including /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling /// back to the existing value in the Merkle tree. Each entry corresponds to a /// [`SparseMerkleTree::insert_value()`] call. - new_pairs: BTreeMap, + new_pairs: HashMap, /// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with /// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call. new_root: RpoDigest, @@ -606,6 +609,114 @@ impl MutationSet { } } +// SERIALIZATION +// ================================================================================================ + +impl Serializable for InnerNode { + fn write_into(&self, target: &mut W) { + target.write(self.left); + target.write(self.right); + } +} + +impl Deserializable for InnerNode { + fn read_from(source: &mut R) -> Result { + let left = source.read()?; + let right = source.read()?; + + Ok(Self { left, right }) + } +} + +impl Serializable for NodeMutation { + fn write_into(&self, target: &mut W) { + match self { + NodeMutation::Removal => target.write_bool(false), + NodeMutation::Addition(inner_node) => { + target.write_bool(true); + inner_node.write_into(target); + }, + } + } +} + +impl Deserializable for NodeMutation { + fn read_from(source: &mut R) -> Result { + if source.read_bool()? { + let inner_node = source.read()?; + return Ok(NodeMutation::Addition(inner_node)); + } + + Ok(NodeMutation::Removal) + } +} + +impl Serializable for MutationSet { + fn write_into(&self, target: &mut W) { + target.write(self.old_root); + target.write(self.new_root); + + let removals: Vec<_> = self + .node_mutations + .iter() + .filter(|(_, value)| matches!(value, NodeMutation::Removal)) + .map(|(key, _)| key) + .collect(); + let additions: Vec<_> = self + .node_mutations + .iter() + .filter_map(|(key, value)| match value { + NodeMutation::Addition(node) => Some((key, node)), + _ => None, + }) + .collect(); + + target.write_u16( + removals.len().try_into().expect("Number of items to remove must fit in u16"), + ); + target.write_many(removals); + + target + .write_u16(additions.len().try_into().expect("Number of items to add must fit in u16")); + target.write_many(additions); + + target.write_u16(self.new_pairs.len() as u16); + target.write_many(&self.new_pairs); + } +} + +impl Deserializable + for MutationSet +{ + fn read_from(source: &mut R) -> Result { + let old_root = source.read()?; + let new_root = source.read()?; + + let num_removals = source.read_u16()? as usize; + let removals: Vec = source.read_many(num_removals)?; + + let num_additions = source.read_u16()? as usize; + let additions: Vec<(NodeIndex, InnerNode)> = source.read_many(num_additions)?; + + let node_mutations = HashMap::from_iter( + removals.into_iter().map(|index| (index, NodeMutation::Removal)).chain( + additions.into_iter().map(|(index, node)| (index, NodeMutation::Addition(node))), + ), + ); + + let num_new_pairs = source.read_u16()? as usize; + let new_pairs = source.read_many(num_new_pairs)?; + let new_pairs = HashMap::from_iter(new_pairs); + + Ok(Self { + old_root, + node_mutations, + new_pairs, + new_root, + }) + } +} + // SUBTREES // ================================================================================================ /// A subtree is of depth 8. @@ -627,10 +738,10 @@ pub struct SubtreeLeaf { } /// Helper struct to organize the return value of [`SparseMerkleTree::sorted_pairs_to_leaves()`]. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone)] pub(crate) struct PairComputations { /// Literal leaves to be added to the sparse Merkle tree's internal mapping. - pub nodes: BTreeMap, + pub nodes: HashMap, /// "Conceptual" leaves that will be used for computations. pub leaves: Vec>, } diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index c236c9da..ce853686 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -1,7 +1,6 @@ -use alloc::{ - collections::{BTreeMap, BTreeSet}, - vec::Vec, -}; +use alloc::{collections::BTreeSet, vec::Vec}; + +use hashbrown::HashMap; use super::{ super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, @@ -22,8 +21,8 @@ mod tests; #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] pub struct SimpleSmt { root: RpoDigest, - leaves: BTreeMap, - inner_nodes: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, } impl SimpleSmt { @@ -54,8 +53,8 @@ impl SimpleSmt { Ok(Self { root, - leaves: BTreeMap::new(), - inner_nodes: BTreeMap::new(), + inner_nodes: Default::default(), + leaves: Default::default(), }) } @@ -109,8 +108,8 @@ impl SimpleSmt { /// 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, - leaves: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, root: RpoDigest, ) -> Self { // Our particular implementation of `from_raw_parts()` never returns `Err`. @@ -327,8 +326,8 @@ impl SparseMerkleTree for SimpleSmt { const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0); fn from_raw_parts( - inner_nodes: BTreeMap, - leaves: BTreeMap, + inner_nodes: HashMap, + leaves: HashMap, root: RpoDigest, ) -> Result { if cfg!(debug_assertions) { diff --git a/src/merkle/smt/simple/tests.rs b/src/merkle/smt/simple/tests.rs index 84bad47f..e8d5d666 100644 --- a/src/merkle/smt/simple/tests.rs +++ b/src/merkle/smt/simple/tests.rs @@ -1,4 +1,4 @@ -use alloc::vec::Vec; +use alloc::{collections::BTreeSet, vec::Vec}; use assert_matches::assert_matches; @@ -141,12 +141,12 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> { let l2n2 = tree.get_node(NodeIndex::make(2, 2))?; let l2n3 = tree.get_node(NodeIndex::make(2, 3))?; - let nodes: Vec = tree.inner_nodes().collect(); - let expected = vec![ + let nodes: BTreeSet = tree.inner_nodes().collect(); + let expected = BTreeSet::from_iter([ InnerNodeInfo { value: root, left: l1n0, right: l1n1 }, InnerNodeInfo { value: l1n0, left: l2n0, right: l2n1 }, InnerNodeInfo { value: l1n1, left: l2n2, right: l2n3 }, - ]; + ]); assert_eq!(nodes, expected); Ok(()) From afb3e8baa801c7414f0b4f9f8c8a4cbd75fdb920 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Thu, 19 Dec 2024 21:16:54 +0500 Subject: [PATCH 02/16] feat: introduce `hashmaps` feature --- Cargo.toml | 3 ++- src/lib.rs | 8 ++++++ src/merkle/smt/full/mod.rs | 22 +++++++--------- src/merkle/smt/full/tests.rs | 18 ++++++------- src/merkle/smt/mod.rs | 48 +++++++++++++++++++++++----------- src/merkle/smt/simple/mod.rs | 22 +++++++--------- src/merkle/smt/simple/tests.rs | 11 +++++--- 7 files changed, 77 insertions(+), 55 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f7b8c7f9..0eb8331f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,7 @@ harness = false concurrent = ["dep:rayon"] default = ["std", "concurrent"] executable = ["dep:clap", "dep:rand-utils", "std"] +hashmaps = ["dep:hashbrown"] internal = [] serde = ["dep:serde", "serde?/alloc", "winter-math/serde"] std = [ @@ -63,7 +64,7 @@ std = [ [dependencies] blake3 = { version = "1.5", default-features = false } clap = { version = "4.5", optional = true, features = ["derive"] } -hashbrown = { version = "0.15", features = ["serde"] } +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 } diff --git a/src/lib.rs b/src/lib.rs index dc04e728..26f3d9f8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,14 @@ pub use winter_math::{ /// A group of four field elements in the Miden base field. pub type Word = [Felt; WORD_SIZE]; +/// A map whose keys are not guarantied to be ordered. +#[cfg(feature = "hashmaps")] +type UnorderedMap = hashbrown::HashMap; + +/// A map whose keys are not guarantied to be ordered. +#[cfg(not(feature = "hashmaps"))] +type UnorderedMap = alloc::collections::BTreeMap; + // CONSTANTS // ================================================================================================ diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 6cc97a6c..175b61cc 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -1,10 +1,8 @@ use alloc::{collections::BTreeSet, string::ToString, vec::Vec}; -use hashbrown::HashMap; - 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; @@ -28,6 +26,8 @@ pub const SMT_DEPTH: u8 = 64; // SMT // ================================================================================================ +type Leaves = super::Leaves; + /// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented /// by 4 field elements. /// @@ -41,8 +41,8 @@ pub const SMT_DEPTH: u8 = 64; #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] pub struct Smt { root: RpoDigest, - inner_nodes: HashMap, - leaves: HashMap, + inner_nodes: InnerNodes, + leaves: Leaves, } impl Smt { @@ -146,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: HashMap, - leaves: HashMap, - 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`. >::from_raw_parts(inner_nodes, leaves, root).unwrap() } @@ -315,8 +311,8 @@ impl SparseMerkleTree for Smt { const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); fn from_raw_parts( - inner_nodes: HashMap, - leaves: HashMap, + inner_nodes: InnerNodes, + leaves: Leaves, root: RpoDigest, ) -> Result { if cfg!(debug_assertions) { diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 03eb3829..b4a702f8 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -499,21 +499,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 diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 62c061c9..33905ca3 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -2,7 +2,7 @@ use alloc::{collections::BTreeMap, vec::Vec}; use core::mem; use std::hash::Hash; -use hashbrown::HashMap; +use constrains::KeyConstrains; use num::Integer; use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; @@ -30,6 +30,24 @@ pub const SMT_MAX_DEPTH: u8 = 64; // SPARSE MERKLE TREE // ================================================================================================ +type InnerNodes = crate::UnorderedMap; +type Leaves = crate::UnorderedMap; +type NodeMutations = crate::UnorderedMap; + +#[cfg(feature = "hashmaps")] +mod constrains { + use core::hash::Hash; + + pub trait KeyConstrains: Hash + Eq {} + impl KeyConstrains for T {} +} + +#[cfg(not(feature = "hashmaps"))] +mod constrains { + pub trait KeyConstrains: Ord {} + impl KeyConstrains for T {} +} + /// An abstract description of a sparse Merkle tree. /// /// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed @@ -51,7 +69,7 @@ pub const SMT_MAX_DEPTH: u8 = 64; /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. pub(crate) trait SparseMerkleTree { /// The type for a key - type Key: Clone + Ord + Eq + Hash; + type Key: Clone + KeyConstrains; /// The type for a value type Value: Clone + PartialEq; /// The type for a leaf @@ -175,8 +193,8 @@ pub(crate) trait SparseMerkleTree { use NodeMutation::*; let mut new_root = self.root(); - let mut new_pairs: HashMap = Default::default(); - let mut node_mutations: HashMap = Default::default(); + let mut new_pairs: crate::UnorderedMap = Default::default(); + let mut node_mutations: NodeMutations = Default::default(); for (key, value) in kv_pairs { // If the old value and the new value are the same, there is nothing to update. @@ -312,8 +330,8 @@ pub(crate) trait SparseMerkleTree { /// Construct this type from already computed leaves and nodes. The caller ensures passed /// arguments are correct and consistent with each other. fn from_raw_parts( - inner_nodes: HashMap, - leaves: HashMap, + inner_nodes: InnerNodes, + leaves: Leaves, root: RpoDigest, ) -> Result where @@ -444,7 +462,7 @@ pub(crate) trait SparseMerkleTree { #[cfg(feature = "concurrent")] fn build_subtrees( mut entries: Vec<(Self::Key, Self::Value)>, - ) -> (HashMap, HashMap) { + ) -> (InnerNodes, Leaves) { entries.sort_by_key(|item| { let index = Self::key_to_leaf_index(&item.0); index.value() @@ -459,10 +477,10 @@ pub(crate) trait SparseMerkleTree { #[cfg(feature = "concurrent")] fn build_subtrees_from_sorted_entries( entries: Vec<(Self::Key, Self::Value)>, - ) -> (HashMap, HashMap) { + ) -> (InnerNodes, Leaves) { use rayon::prelude::*; - let mut accumulated_nodes: HashMap = Default::default(); + let mut accumulated_nodes: InnerNodes = Default::default(); let PairComputations { leaves: mut leaf_subtrees, @@ -590,12 +608,12 @@ pub struct MutationSet { /// index overlayed, if any. Each [`NodeMutation::Addition`] corresponds to a /// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`] /// corresponds to a [`SparseMerkleTree::remove_inner_node()`] call. - node_mutations: HashMap, + node_mutations: NodeMutations, /// The set of top-level key-value pairs we're prospectively adding to the tree, including /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling /// back to the existing value in the Merkle tree. Each entry corresponds to a /// [`SparseMerkleTree::insert_value()`] call. - new_pairs: HashMap, + new_pairs: crate::UnorderedMap, /// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with /// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call. new_root: RpoDigest, @@ -685,7 +703,7 @@ impl Serializable for Mutatio } } -impl Deserializable +impl Deserializable for MutationSet { fn read_from(source: &mut R) -> Result { @@ -698,7 +716,7 @@ impl Deserial let num_additions = source.read_u16()? as usize; let additions: Vec<(NodeIndex, InnerNode)> = source.read_many(num_additions)?; - let node_mutations = HashMap::from_iter( + let node_mutations = NodeMutations::from_iter( removals.into_iter().map(|index| (index, NodeMutation::Removal)).chain( additions.into_iter().map(|(index, node)| (index, NodeMutation::Addition(node))), ), @@ -706,7 +724,7 @@ impl Deserial let num_new_pairs = source.read_u16()? as usize; let new_pairs = source.read_many(num_new_pairs)?; - let new_pairs = HashMap::from_iter(new_pairs); + let new_pairs = crate::UnorderedMap::from_iter(new_pairs); Ok(Self { old_root, @@ -741,7 +759,7 @@ pub struct SubtreeLeaf { #[derive(Debug, Clone)] pub(crate) struct PairComputations { /// Literal leaves to be added to the sparse Merkle tree's internal mapping. - pub nodes: HashMap, + pub nodes: crate::UnorderedMap, /// "Conceptual" leaves that will be used for computations. pub leaves: Vec>, } diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index ce853686..511deedf 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -1,10 +1,8 @@ use alloc::{collections::BTreeSet, vec::Vec}; -use hashbrown::HashMap; - use super::{ - super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, - MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, + super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, InnerNodes, LeafIndex, + MerkleError, MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD, SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; @@ -14,6 +12,8 @@ mod tests; // SPARSE MERKLE TREE // ================================================================================================ +type Leaves = super::Leaves; + /// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction. /// /// The root of the tree is recomputed on each new leaf update. @@ -21,8 +21,8 @@ mod tests; #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] pub struct SimpleSmt { root: RpoDigest, - inner_nodes: HashMap, - leaves: HashMap, + inner_nodes: InnerNodes, + leaves: Leaves, } impl SimpleSmt { @@ -107,11 +107,7 @@ impl SimpleSmt { /// # 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: HashMap, - leaves: HashMap, - 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`. >::from_raw_parts(inner_nodes, leaves, root).unwrap() } @@ -326,8 +322,8 @@ impl SparseMerkleTree for SimpleSmt { const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0); fn from_raw_parts( - inner_nodes: HashMap, - leaves: HashMap, + inner_nodes: InnerNodes, + leaves: Leaves, root: RpoDigest, ) -> Result { if cfg!(debug_assertions) { diff --git a/src/merkle/smt/simple/tests.rs b/src/merkle/smt/simple/tests.rs index e8d5d666..9078c526 100644 --- a/src/merkle/smt/simple/tests.rs +++ b/src/merkle/smt/simple/tests.rs @@ -1,4 +1,4 @@ -use alloc::{collections::BTreeSet, vec::Vec}; +use alloc::vec::Vec; use assert_matches::assert_matches; @@ -141,12 +141,15 @@ fn test_inner_node_iterator() -> Result<(), MerkleError> { let l2n2 = tree.get_node(NodeIndex::make(2, 2))?; let l2n3 = tree.get_node(NodeIndex::make(2, 3))?; - let nodes: BTreeSet = tree.inner_nodes().collect(); - let expected = BTreeSet::from_iter([ + let mut nodes: Vec = tree.inner_nodes().collect(); + let mut expected = [ InnerNodeInfo { value: root, left: l1n0, right: l1n1 }, InnerNodeInfo { value: l1n0, left: l2n0, right: l2n1 }, InnerNodeInfo { value: l1n1, left: l2n2, right: l2n3 }, - ]); + ]; + nodes.sort(); + expected.sort(); + assert_eq!(nodes, expected); Ok(()) From 9e555e779b3a46972b55f1256857f9e05e3e2594 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Fri, 20 Dec 2024 11:30:37 +0500 Subject: [PATCH 03/16] refactor: rename feature to `smt_hashmaps` and use it only for SMTs --- Cargo.toml | 4 ++-- src/lib.rs | 8 -------- src/merkle/smt/mod.rs | 34 ++++++++++++++++++++-------------- 3 files changed, 22 insertions(+), 24 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0eb8331f..bde034fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,7 +48,7 @@ harness = false concurrent = ["dep:rayon"] default = ["std", "concurrent"] executable = ["dep:clap", "dep:rand-utils", "std"] -hashmaps = ["dep:hashbrown"] +smt_hashmaps = [] internal = [] serde = ["dep:serde", "serde?/alloc", "winter-math/serde"] std = [ @@ -64,7 +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"] } +hashbrown = { version = "0.15", 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 } diff --git a/src/lib.rs b/src/lib.rs index 26f3d9f8..dc04e728 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,14 +26,6 @@ pub use winter_math::{ /// A group of four field elements in the Miden base field. pub type Word = [Felt; WORD_SIZE]; -/// A map whose keys are not guarantied to be ordered. -#[cfg(feature = "hashmaps")] -type UnorderedMap = hashbrown::HashMap; - -/// A map whose keys are not guarantied to be ordered. -#[cfg(not(feature = "hashmaps"))] -type UnorderedMap = alloc::collections::BTreeMap; - // CONSTANTS // ================================================================================================ diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 33905ca3..d2409d32 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -2,8 +2,8 @@ use alloc::{collections::BTreeMap, vec::Vec}; use core::mem; use std::hash::Hash; -use constrains::KeyConstrains; use num::Integer; +use types::{KeyConstrains, UnorderedMap}; use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex}; @@ -30,24 +30,30 @@ pub const SMT_MAX_DEPTH: u8 = 64; // SPARSE MERKLE TREE // ================================================================================================ -type InnerNodes = crate::UnorderedMap; -type Leaves = crate::UnorderedMap; -type NodeMutations = crate::UnorderedMap; - -#[cfg(feature = "hashmaps")] -mod constrains { +#[cfg(feature = "smt_hashmaps")] +mod types { use core::hash::Hash; + /// A map whose keys are not guarantied to be ordered. + pub type UnorderedMap = hashbrown::HashMap; + pub trait KeyConstrains: Hash + Eq {} impl KeyConstrains for T {} } -#[cfg(not(feature = "hashmaps"))] -mod constrains { +#[cfg(not(feature = "smt_hashmaps"))] +mod types { + /// A map whose keys are not guarantied to be ordered. + pub type UnorderedMap = alloc::collections::BTreeMap; + pub trait KeyConstrains: Ord {} impl KeyConstrains for T {} } +type InnerNodes = UnorderedMap; +type Leaves = UnorderedMap; +type NodeMutations = UnorderedMap; + /// An abstract description of a sparse Merkle tree. /// /// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed @@ -193,7 +199,7 @@ pub(crate) trait SparseMerkleTree { use NodeMutation::*; let mut new_root = self.root(); - let mut new_pairs: crate::UnorderedMap = Default::default(); + let mut new_pairs: UnorderedMap = Default::default(); let mut node_mutations: NodeMutations = Default::default(); for (key, value) in kv_pairs { @@ -613,7 +619,7 @@ pub struct MutationSet { /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling /// back to the existing value in the Merkle tree. Each entry corresponds to a /// [`SparseMerkleTree::insert_value()`] call. - new_pairs: crate::UnorderedMap, + new_pairs: UnorderedMap, /// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with /// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call. new_root: RpoDigest, @@ -724,7 +730,7 @@ impl Dese let num_new_pairs = source.read_u16()? as usize; let new_pairs = source.read_many(num_new_pairs)?; - let new_pairs = crate::UnorderedMap::from_iter(new_pairs); + let new_pairs = UnorderedMap::from_iter(new_pairs); Ok(Self { old_root, @@ -759,7 +765,7 @@ pub struct SubtreeLeaf { #[derive(Debug, Clone)] pub(crate) struct PairComputations { /// Literal leaves to be added to the sparse Merkle tree's internal mapping. - pub nodes: crate::UnorderedMap, + pub nodes: UnorderedMap, /// "Conceptual" leaves that will be used for computations. pub leaves: Vec>, } @@ -787,7 +793,7 @@ impl<'s> SubtreeLeavesIter<'s> { Self { leaves: leaves.drain(..).peekable() } } } -impl core::iter::Iterator for SubtreeLeavesIter<'_> { +impl Iterator for SubtreeLeavesIter<'_> { type Item = Vec; /// Each `next()` collects an entire subtree. From a9f7bbe923e91c00336dbb007f6ff5c669d6b2c5 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Fri, 20 Dec 2024 11:33:39 +0500 Subject: [PATCH 04/16] docs: update `CHANGELOG.md` --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2957f2c9..6b502ed3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,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) From dab437fcf598e118c384733a3d409870609e3ecd Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Fri, 20 Dec 2024 11:35:24 +0500 Subject: [PATCH 05/16] fix: no-std build --- src/hash/rescue/rpo/digest.rs | 9 +++++++-- src/merkle/smt/mod.rs | 3 +-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/hash/rescue/rpo/digest.rs b/src/hash/rescue/rpo/digest.rs index 2206aab6..a525892b 100644 --- a/src/hash/rescue/rpo/digest.rs +++ b/src/hash/rescue/rpo/digest.rs @@ -1,6 +1,11 @@ use alloc::string::String; -use core::{cmp::Ordering, fmt::Display, ops::Deref, slice}; -use std::hash::{Hash, Hasher}; +use core::{ + cmp::Ordering, + fmt::Display, + hash::{Hash, Hasher}, + ops::Deref, + slice, +}; use thiserror::Error; diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index d2409d32..9e807f00 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -1,6 +1,5 @@ use alloc::{collections::BTreeMap, vec::Vec}; -use core::mem; -use std::hash::Hash; +use core::{hash::Hash, mem}; use num::Integer; use types::{KeyConstrains, UnorderedMap}; From f9b79b1414856e0a7f4b0784924bb727c19a4c38 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Fri, 27 Dec 2024 17:57:48 +0500 Subject: [PATCH 06/16] fix: compilation errors --- src/merkle/smt/full/tests.rs | 12 +++--- src/merkle/smt/mod.rs | 81 ++++++------------------------------ 2 files changed, 18 insertions(+), 75 deletions(-) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 9f20edb6..d52367db 100644 --- a/src/merkle/smt/full/tests.rs +++ b/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::{types::UnorderedMap, NodeMutation, SparseMerkleTree}, 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" ); diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index f8682fd6..78c33637 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -366,7 +366,7 @@ pub(crate) trait SparseMerkleTree { }); } - let mut reverse_mutations = BTreeMap::new(); + let mut reverse_mutations = NodeMutations::new(); for (index, mutation) in node_mutations { match mutation { Removal => { @@ -384,7 +384,7 @@ pub(crate) trait SparseMerkleTree { } } - let mut reverse_pairs = BTreeMap::new(); + let mut reverse_pairs = UnorderedMap::new(); for (key, value) in new_pairs { if let Some(old_value) = self.insert_value(key.clone(), value) { reverse_pairs.insert(key, old_value); @@ -711,85 +711,27 @@ impl MutationSet { } /// Returns the set of inner nodes that need to be removed or added. - pub fn node_mutations(&self) -> &BTreeMap { + pub fn node_mutations(&self) -> &NodeMutations { &self.node_mutations } /// Returns the set of top-level key-value pairs that need to be added, updated or deleted /// (i.e. set to `EMPTY_WORD`). - pub fn new_pairs(&self) -> &BTreeMap { + pub fn new_pairs(&self) -> &UnorderedMap { &self.new_pairs } } -// SERIALIZATION -// ================================================================================================ - -impl Serializable for InnerNode { - fn write_into(&self, target: &mut W) { - self.left.write_into(target); - self.right.write_into(target); +impl PartialEq for MutationSet { + fn eq(&self, other: &Self) -> bool { + self.old_root == other.old_root + && self.node_mutations == other.node_mutations + && self.new_pairs == other.new_pairs + && self.new_root == other.new_root } } -impl Deserializable for InnerNode { - fn read_from(source: &mut R) -> Result { - let left = source.read()?; - let right = source.read()?; - - Ok(Self { left, right }) - } -} - -impl Serializable for NodeMutation { - fn write_into(&self, target: &mut W) { - match self { - NodeMutation::Removal => target.write_bool(false), - NodeMutation::Addition(inner_node) => { - target.write_bool(true); - inner_node.write_into(target); - }, - } - } -} - -impl Deserializable for NodeMutation { - fn read_from(source: &mut R) -> Result { - if source.read_bool()? { - let inner_node = source.read()?; - return Ok(NodeMutation::Addition(inner_node)); - } - - Ok(NodeMutation::Removal) - } -} - -impl Serializable for MutationSet { - fn write_into(&self, target: &mut W) { - target.write(self.old_root); - target.write(self.new_root); - self.node_mutations.write_into(target); - self.new_pairs.write_into(target); - } -} - -impl Deserializable - for MutationSet -{ - fn read_from(source: &mut R) -> Result { - let old_root = source.read()?; - let new_root = source.read()?; - let node_mutations = source.read()?; - let new_pairs = source.read()?; - - Ok(Self { - old_root, - node_mutations, - new_pairs, - new_root, - }) - } -} +impl Eq for MutationSet {} // SERIALIZATION // ================================================================================================ @@ -901,6 +843,7 @@ impl Dese // SUBTREES // ================================================================================================ + /// A subtree is of depth 8. const SUBTREE_DEPTH: u8 = 8; From c9e198a4f17016cc8ae0110fb0be4cac80fc6d4c Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Sun, 29 Dec 2024 11:11:24 +0500 Subject: [PATCH 07/16] fix: typo --- src/merkle/smt/mod.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 78c33637..7dec0378 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -2,7 +2,7 @@ use alloc::{collections::BTreeMap, vec::Vec}; use core::{hash::Hash, mem}; use num::Integer; -use types::{KeyConstrains, UnorderedMap}; +use types::{KeyConstraints, UnorderedMap}; use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex}; @@ -36,8 +36,8 @@ mod types { /// A map whose keys are not guarantied to be ordered. pub type UnorderedMap = hashbrown::HashMap; - pub trait KeyConstrains: Hash + Eq {} - impl KeyConstrains for T {} + pub trait KeyConstraints: Hash + Eq {} + impl KeyConstraints for T {} } #[cfg(not(feature = "smt_hashmaps"))] @@ -45,8 +45,8 @@ mod types { /// A map whose keys are not guarantied to be ordered. pub type UnorderedMap = alloc::collections::BTreeMap; - pub trait KeyConstrains: Ord {} - impl KeyConstrains for T {} + pub trait KeyConstraints: Ord {} + impl KeyConstraints for T {} } type InnerNodes = UnorderedMap; @@ -74,7 +74,7 @@ type NodeMutations = UnorderedMap; /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. pub(crate) trait SparseMerkleTree { /// The type for a key - type Key: Clone + KeyConstrains; + type Key: Clone + KeyConstraints; /// The type for a value type Value: Clone + PartialEq; /// The type for a leaf @@ -722,7 +722,7 @@ impl MutationSet { } } -impl PartialEq for MutationSet { +impl PartialEq for MutationSet { fn eq(&self, other: &Self) -> bool { self.old_root == other.old_root && self.node_mutations == other.node_mutations @@ -731,7 +731,7 @@ impl PartialEq for MutationSet< } } -impl Eq for MutationSet {} +impl Eq for MutationSet {} // SERIALIZATION // ================================================================================================ @@ -809,7 +809,7 @@ impl Serializable for Mutatio } } -impl Deserializable +impl Deserializable for MutationSet { fn read_from(source: &mut R) -> Result { From da9bc922f7e91c18fd5c08d4cd68d0c3951b4822 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Sun, 29 Dec 2024 11:14:31 +0500 Subject: [PATCH 08/16] refactor: make `InnerNodeInfo` derive `PartialOrd, Ord` only for tests --- src/merkle/node.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/merkle/node.rs b/src/merkle/node.rs index ed92d7cc..b8219033 100644 --- a/src/merkle/node.rs +++ b/src/merkle/node.rs @@ -1,8 +1,9 @@ use super::RpoDigest; /// Representation of a node with two children used for iterating over containers. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +#[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, From a017c970b1d00c0c7ffa44b670ab03352ae3c69a Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Sun, 29 Dec 2024 11:22:57 +0500 Subject: [PATCH 09/16] fix: use `u32` for serialization of `MutationSet` --- src/merkle/smt/mod.rs | 43 +++++++++++++++++++++++++++---------------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 7dec0378..1b4636d5 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -780,13 +780,13 @@ impl Serializable for Mutatio target.write(self.old_root); target.write(self.new_root); - let removals: Vec<_> = self + let inner_removals: Vec<_> = self .node_mutations .iter() .filter(|(_, value)| matches!(value, NodeMutation::Removal)) .map(|(key, _)| key) .collect(); - let additions: Vec<_> = self + let inner_additions: Vec<_> = self .node_mutations .iter() .filter_map(|(key, value)| match value { @@ -795,16 +795,25 @@ impl Serializable for Mutatio }) .collect(); - target.write_u16( - removals.len().try_into().expect("Number of items to remove must fit in u16"), + target.write_u32( + inner_removals + .len() + .try_into() + .expect("Number of items to remove must fit in `u32`"), ); - target.write_many(removals); + target.write_many(inner_removals); - target - .write_u16(additions.len().try_into().expect("Number of items to add must fit in u16")); - target.write_many(additions); + target.write_u32( + inner_additions + .len() + .try_into() + .expect("Number of items to add must fit in `u32`"), + ); + target.write_many(inner_additions); - target.write_u16(self.new_pairs.len() as u16); + target.write_u32( + self.new_pairs.len().try_into().expect("Number of new pairs must fit in `u32`"), + ); target.write_many(&self.new_pairs); } } @@ -816,19 +825,21 @@ impl Des let old_root = source.read()?; let new_root = source.read()?; - let num_removals = source.read_u16()? as usize; - let removals: Vec = source.read_many(num_removals)?; + let num_removals = source.read_u32()? as usize; + let inner_removals: Vec = source.read_many(num_removals)?; - let num_additions = source.read_u16()? as usize; - let additions: Vec<(NodeIndex, InnerNode)> = source.read_many(num_additions)?; + let num_additions = source.read_u32()? as usize; + let inner_additions: Vec<(NodeIndex, InnerNode)> = source.read_many(num_additions)?; let node_mutations = NodeMutations::from_iter( - removals.into_iter().map(|index| (index, NodeMutation::Removal)).chain( - additions.into_iter().map(|(index, node)| (index, NodeMutation::Addition(node))), + inner_removals.into_iter().map(|index| (index, NodeMutation::Removal)).chain( + inner_additions + .into_iter() + .map(|(index, node)| (index, NodeMutation::Addition(node))), ), ); - let num_new_pairs = source.read_u16()? as usize; + let num_new_pairs = source.read_u32()? as usize; let new_pairs = source.read_many(num_new_pairs)?; let new_pairs = UnorderedMap::from_iter(new_pairs); From 2add8dd36471ac1ce6591b9e47d359cb492f3f9f Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Sun, 29 Dec 2024 12:29:15 +0500 Subject: [PATCH 10/16] fix: use `usize` for serialization of `MutationSet` --- src/merkle/smt/mod.rs | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 1b4636d5..3d5a0b7e 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -795,25 +795,13 @@ impl Serializable for Mutatio }) .collect(); - target.write_u32( - inner_removals - .len() - .try_into() - .expect("Number of items to remove must fit in `u32`"), - ); + target.write_usize(inner_removals.len()); target.write_many(inner_removals); - target.write_u32( - inner_additions - .len() - .try_into() - .expect("Number of items to add must fit in `u32`"), - ); + target.write_usize(inner_additions.len()); target.write_many(inner_additions); - target.write_u32( - self.new_pairs.len().try_into().expect("Number of new pairs must fit in `u32`"), - ); + target.write_usize(self.new_pairs.len()); target.write_many(&self.new_pairs); } } @@ -825,10 +813,10 @@ impl Des let old_root = source.read()?; let new_root = source.read()?; - let num_removals = source.read_u32()? as usize; + let num_removals = source.read_usize()?; let inner_removals: Vec = source.read_many(num_removals)?; - let num_additions = source.read_u32()? as usize; + let num_additions = source.read_usize()?; let inner_additions: Vec<(NodeIndex, InnerNode)> = source.read_many(num_additions)?; let node_mutations = NodeMutations::from_iter( @@ -839,7 +827,7 @@ impl Des ), ); - let num_new_pairs = source.read_u32()? as usize; + let num_new_pairs = source.read_usize()?; let new_pairs = source.read_many(num_new_pairs)?; let new_pairs = UnorderedMap::from_iter(new_pairs); From 720c02e9be1cb781142422954a74ae29bcee7326 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Tue, 31 Dec 2024 00:07:09 +0500 Subject: [PATCH 11/16] refactor: get rid of `KeyConstraints` --- src/merkle/smt/full/tests.rs | 2 +- src/merkle/smt/mod.rs | 47 +++++++++--------------------------- 2 files changed, 12 insertions(+), 37 deletions(-) diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index d52367db..787c01a8 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -3,7 +3,7 @@ use alloc::vec::Vec; use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH}; use crate::{ merkle::{ - smt::{types::UnorderedMap, NodeMutation, SparseMerkleTree}, + smt::{NodeMutation, SparseMerkleTree, UnorderedMap}, EmptySubtreeRoots, MerkleStore, MutationSet, }, utils::{Deserializable, Serializable}, diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index 3d5a0b7e..b96c1bd7 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -2,7 +2,6 @@ use alloc::{collections::BTreeMap, vec::Vec}; use core::{hash::Hash, mem}; use num::Integer; -use types::{KeyConstraints, UnorderedMap}; use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex}; @@ -29,26 +28,11 @@ pub const SMT_MAX_DEPTH: u8 = 64; // SPARSE MERKLE TREE // ================================================================================================ +/// A map whose keys are not guarantied to be ordered. #[cfg(feature = "smt_hashmaps")] -mod types { - use core::hash::Hash; - - /// A map whose keys are not guarantied to be ordered. - pub type UnorderedMap = hashbrown::HashMap; - - pub trait KeyConstraints: Hash + Eq {} - impl KeyConstraints for T {} -} - +type UnorderedMap = hashbrown::HashMap; #[cfg(not(feature = "smt_hashmaps"))] -mod types { - /// A map whose keys are not guarantied to be ordered. - pub type UnorderedMap = alloc::collections::BTreeMap; - - pub trait KeyConstraints: Ord {} - impl KeyConstraints for T {} -} - +type UnorderedMap = alloc::collections::BTreeMap; type InnerNodes = UnorderedMap; type Leaves = UnorderedMap; type NodeMutations = UnorderedMap; @@ -74,7 +58,7 @@ type NodeMutations = UnorderedMap; /// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. pub(crate) trait SparseMerkleTree { /// The type for a key - type Key: Clone + KeyConstraints; + type Key: Clone + Ord + Eq + Hash; /// The type for a value type Value: Clone + PartialEq; /// The type for a leaf @@ -676,8 +660,8 @@ pub enum NodeMutation { /// Represents a group of prospective mutations to a `SparseMerkleTree`, created by /// `SparseMerkleTree::compute_mutations()`, and that can be applied with /// `SparseMerkleTree::apply_mutations()`. -#[derive(Debug, Clone, Default)] -pub struct MutationSet { +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct MutationSet { /// The root of the Merkle tree this MutationSet is for, recorded at the time /// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying /// mutations to the wrong tree or applying stale mutations to a tree that has since changed. @@ -698,7 +682,7 @@ pub struct MutationSet { new_root: RpoDigest, } -impl MutationSet { +impl MutationSet { /// Returns the SMT root that was calculated during `SparseMerkleTree::compute_mutations()`. See /// that method for more information. pub fn root(&self) -> RpoDigest { @@ -722,17 +706,6 @@ impl MutationSet { } } -impl PartialEq for MutationSet { - fn eq(&self, other: &Self) -> bool { - self.old_root == other.old_root - && self.node_mutations == other.node_mutations - && self.new_pairs == other.new_pairs - && self.new_root == other.new_root - } -} - -impl Eq for MutationSet {} - // SERIALIZATION // ================================================================================================ @@ -775,7 +748,9 @@ impl Deserializable for NodeMutation { } } -impl Serializable for MutationSet { +impl Serializable + for MutationSet +{ fn write_into(&self, target: &mut W) { target.write(self.old_root); target.write(self.new_root); @@ -806,7 +781,7 @@ impl Serializable for Mutatio } } -impl Deserializable +impl Deserializable for MutationSet { fn read_from(source: &mut R) -> Result { From a6324965da40c1fe8402595dd6cb3150bb31273d Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Tue, 31 Dec 2024 00:09:43 +0500 Subject: [PATCH 12/16] refactor: make `hashbrown` dependency optional --- Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 08b47fea..3469c089 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,7 +48,7 @@ harness = false concurrent = ["dep:rayon"] default = ["std", "concurrent"] executable = ["dep:clap", "dep:rand-utils", "std"] -smt_hashmaps = [] +smt_hashmaps = ["dep:hashbrown"] internal = [] serde = ["dep:serde", "serde?/alloc", "winter-math/serde"] std = [ @@ -64,7 +64,7 @@ std = [ [dependencies] blake3 = { version = "1.5", default-features = false } clap = { version = "4.5", optional = true, features = ["derive"] } -hashbrown = { version = "0.15", features = ["serde"] } +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 } From a797a9e6e5225af89b650ec66b997584d5d749c7 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Tue, 31 Dec 2024 00:17:53 +0500 Subject: [PATCH 13/16] refactor: read/write vectors directly --- src/merkle/smt/mod.rs | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index b96c1bd7..ec439571 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -770,11 +770,8 @@ impl Serializable }) .collect(); - target.write_usize(inner_removals.len()); - target.write_many(inner_removals); - - target.write_usize(inner_additions.len()); - target.write_many(inner_additions); + target.write(inner_removals); + target.write(inner_additions); target.write_usize(self.new_pairs.len()); target.write_many(&self.new_pairs); @@ -788,11 +785,8 @@ impl De let old_root = source.read()?; let new_root = source.read()?; - let num_removals = source.read_usize()?; - let inner_removals: Vec = source.read_many(num_removals)?; - - let num_additions = source.read_usize()?; - let inner_additions: Vec<(NodeIndex, InnerNode)> = source.read_many(num_additions)?; + let inner_removals: Vec = source.read()?; + let inner_additions: Vec<(NodeIndex, InnerNode)> = source.read()?; let node_mutations = NodeMutations::from_iter( inner_removals.into_iter().map(|index| (index, NodeMutation::Removal)).chain( From 393483b792eed8596f095e2419b6fb68001a5bb1 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Thu, 2 Jan 2025 14:26:30 +0500 Subject: [PATCH 14/16] docs: add `smt_hashmaps` feature description --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index c9a758ad..7933e0b0 100644 --- a/README.md +++ b/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. From c43773967c68047ebd398c4e2f720e140e85b5c8 Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Thu, 2 Jan 2025 14:38:47 +0500 Subject: [PATCH 15/16] chore: add `smt_hashmaps` feature testing --- Makefile | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 39274435..a855cc04 100644 --- a/Makefile +++ b/Makefile @@ -46,6 +46,9 @@ 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) @@ -53,7 +56,7 @@ test-no-std: ## Run tests with `no-default-features` (std) .PHONY: test -test: test-default test-no-std ## Run all tests +test: test-default test-smt-hashmaps test-no-std ## Run all tests # --- checking ------------------------------------------------------------------------------------ From 1e88c0870702a1d00c3e2fd60c9a5d027efe03fe Mon Sep 17 00:00:00 2001 From: polydez <155382956+polydez@users.noreply.github.com> Date: Thu, 2 Jan 2025 14:53:57 +0500 Subject: [PATCH 16/16] chore: add `smt-hashmaps` testing to github actions --- .github/workflows/test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index bb9395b1..2194ef66 100644 --- a/.github/workflows/test.yml +++ b/.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