diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 8bbc38ac..edd75d40 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -24,14 +24,8 @@ pub use merkle_tree::{path_to_text, tree_to_text, MerkleTree}; mod path; pub use path::{MerklePath, RootPath, ValuePath}; -mod simple_smt; -pub use simple_smt::SimpleSmt; - -pub(crate) mod sparse_merkle_tree; -pub use sparse_merkle_tree::{LeafIndex, SMT_MAX_DEPTH, SMT_MIN_DEPTH}; - mod smt; -pub use smt::{Smt, SmtLeaf, SMT_DEPTH}; +pub use smt::{LeafIndex, SimpleSmt, Smt, SmtLeaf, SMT_DEPTH, SMT_MAX_DEPTH, SMT_MIN_DEPTH}; mod tiered_smt; pub use tiered_smt::{TieredSmt, TieredSmtProof, TieredSmtProofError}; diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs new file mode 100644 index 00000000..f8c17aca --- /dev/null +++ b/src/merkle/smt/full/mod.rs @@ -0,0 +1,367 @@ +use core::cmp::Ordering; + +use winter_math::StarkField; + +use crate::hash::rpo::Rpo256; +use crate::merkle::EmptySubtreeRoots; +use crate::utils::{ + collections::{BTreeMap, BTreeSet, Vec}, + vec, +}; +use crate::{Felt, EMPTY_WORD}; + +use super::{ + InnerNode, LeafIndex, MerkleError, MerklePath, NodeIndex, RpoDigest, SparseMerkleTree, Word, +}; + +#[cfg(test)] +mod tests; + +// CONSTANTS +// ================================================================================================ + +pub const SMT_DEPTH: u8 = 64; + +// SMT +// ================================================================================================ + +#[derive(Debug, Clone, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] +pub struct Smt { + root: RpoDigest, + leaves: BTreeMap, + inner_nodes: BTreeMap, +} + +impl Smt { + // CONSTRUCTORS + // -------------------------------------------------------------------------------------------- + + /// Returns a new [NewSmt]. + /// + /// All leaves in the returned tree are set to [ZERO; 4]. + pub fn new() -> Self { + let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); + + Self { + root, + leaves: BTreeMap::new(), + inner_nodes: BTreeMap::new(), + } + } + + /// Returns a new [SimpleSmt] instantiated with leaves set as specified by the provided entries. + /// + /// All leaves omitted from the entries list are set to [ZERO; 4]. + /// + /// # Errors + /// Returns an error if: + /// - The number of entries exceeds 2^63 entries. + /// - The provided entries contain multiple values for the same key. + pub fn with_entries( + entries: impl IntoIterator, + ) -> Result { + // create an empty tree + let mut tree = Self::new(); + + // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so + // entries with the empty value need additional tracking. + let mut key_set_to_zero = BTreeSet::new(); + + for (key, value) in entries { + let old_value = tree.insert(key, value); + + if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) { + return Err(MerkleError::DuplicateValuesForIndex( + LeafIndex::::from(key).value(), + )); + } + + if value == EMPTY_WORD { + key_set_to_zero.insert(key); + }; + } + Ok(tree) + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the root of the tree + pub fn root(&self) -> RpoDigest { + >::root(self) + } + + /// Returns the leaf at the specified index. + pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf { + >::get_leaf(self, key) + } + + /// Returns the depth of the tree + pub const fn depth(&self) -> u8 { + SMT_DEPTH + } + + /// Returns a Merkle path from the leaf node specified by the key to the root. + /// + /// The node itself is not included in the path. + pub fn open(&self, key: &RpoDigest) -> (MerklePath, SmtLeaf) { + >::open(self, key) + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Inserts a value at the specified key, returning the previous value associated with that key. + /// Recall that by definition, any key that hasn't been updated is associated with + /// [`EMPTY_WORD`]. + /// + /// This also recomputes all hashes between the leaf (associated with the key) and the root, + /// updating the root itself. + pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word { + >::insert(self, key, value) + } + + // HELPERS + // -------------------------------------------------------------------------------------------- + + /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty + /// value, such that this is indeed an insertion. + fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option { + debug_assert_ne!(value, Self::EMPTY_VALUE); + + let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); + + match self.leaves.get_mut(&leaf_index.value()) { + Some(leaf) => leaf.insert(key, value), + None => { + self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value))); + + None + } + } + } + + /// Removes key-value pair at leaf index pointed to by `key` if it exists. + fn perform_remove(&mut self, key: RpoDigest) -> Option { + let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); + + if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) { + let (old_value, is_empty) = leaf.remove(key); + if is_empty { + self.leaves.remove(&leaf_index.value()); + } + old_value + } else { + // there's nothing stored at the leaf; nothing to update + None + } + } +} + +impl SparseMerkleTree for Smt { + type Key = RpoDigest; + type Value = Word; + type Leaf = SmtLeaf; + type Opening = (MerklePath, SmtLeaf); + + const EMPTY_VALUE: Self::Value = EMPTY_WORD; + + fn root(&self) -> RpoDigest { + self.root + } + + fn set_root(&mut self, root: RpoDigest) { + self.root = root; + } + + fn get_inner_node(&self, index: NodeIndex) -> InnerNode { + self.inner_nodes.get(&index).cloned().unwrap_or_else(|| { + let node = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1); + + InnerNode { left: *node, right: *node } + }) + } + + fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) { + self.inner_nodes.insert(index, inner_node); + } + + fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option { + // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key` + if value != Self::EMPTY_VALUE { + self.perform_insert(key, value) + } else { + self.perform_remove(key) + } + } + + fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf { + let leaf_pos = LeafIndex::::from(*key).value(); + + match self.leaves.get(&leaf_pos) { + Some(leaf) => leaf.clone(), + None => SmtLeaf::Single((*key, Self::EMPTY_VALUE)), + } + } + + fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest { + leaf.hash() + } + + fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex { + let most_significant_felt = key[3]; + LeafIndex::new_max_depth(most_significant_felt.as_int()) + } +} + +impl Default for Smt { + fn default() -> Self { + Self::new() + } +} + +// LEAF +// ================================================================================================ + +#[derive(Clone, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] +pub enum SmtLeaf { + Single((RpoDigest, Word)), + Multiple(Vec<(RpoDigest, Word)>), +} + +impl SmtLeaf { + /// Converts a leaf to a list of field elements + pub fn to_elements(&self) -> Vec { + self.clone().into_elements() + } + + /// Converts a leaf to a list of field elements + pub fn into_elements(self) -> Vec { + match self { + SmtLeaf::Single(kv_pair) => kv_to_elements(kv_pair).collect(), + SmtLeaf::Multiple(kv_pairs) => kv_pairs.into_iter().flat_map(kv_to_elements).collect(), + } + } + + /// Compute the hash of the leaf + pub fn hash(&self) -> RpoDigest { + match self { + SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]), + SmtLeaf::Multiple(kvs) => { + let elements: Vec = kvs.iter().copied().flat_map(kv_to_elements).collect(); + Rpo256::hash_elements(&elements) + } + } + } + + // HELPERS + // --------------------------------------------------------------------------------------------- + + /// Insert key-value pair into the leaf; return the previous value associated with `key`, if + /// any. + fn insert(&mut self, key: RpoDigest, value: Word) -> Option { + match self { + SmtLeaf::Single(kv_pair) => { + if kv_pair.0 == key { + // the key is already in this leaf. Update the value and return the previous + // value + let old_value = kv_pair.1; + kv_pair.1 = value; + Some(old_value) + } else { + // Another entry is present in this leaf. Transform the entry into a list + // entry, and make sure the key-value pairs are sorted by key + let mut pairs = vec![*kv_pair, (key, value)]; + pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2)); + + *self = SmtLeaf::Multiple(pairs); + + None + } + } + SmtLeaf::Multiple(kv_pairs) => { + match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { + Ok(pos) => { + let old_value = kv_pairs[pos].1; + kv_pairs[pos].1 = value; + + Some(old_value) + } + Err(pos) => { + kv_pairs.insert(pos, (key, value)); + + None + } + } + } + } + } + + /// Remove key-value pair into the leaf stored at key; return the previous value associated with + /// `key`, if any. We also return an `is_empty` flag which indicates whether the leaf became + /// empty, and must be removed from the data structure it is contained in. + fn remove(&mut self, key: RpoDigest) -> (Option, bool) { + match self { + SmtLeaf::Single((key_at_leaf, value_at_leaf)) => { + if *key_at_leaf == key { + // our key was indeed stored in the leaf, so we return the value that was stored + // in it, and indicate that the leaf should be removed + let old_value = *value_at_leaf; + + (Some(old_value), true) + } else { + // another key is stored at leaf; nothing to update + (None, false) + } + } + SmtLeaf::Multiple(kv_pairs) => { + match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { + Ok(pos) => { + let old_value = kv_pairs[pos].1; + + kv_pairs.remove(pos); + debug_assert!(!kv_pairs.is_empty()); + + if kv_pairs.len() == 1 { + // convert the leaf into `Single` + *self = SmtLeaf::Single(kv_pairs[0]); + } + + (Some(old_value), false) + } + Err(_) => { + // other keys are stored at leaf; nothing to update + (None, false) + } + } + } + } + } +} + +// HELPER FUNCTIONS +// ================================================================================================ + +/// Converts a key-value tuple to an iterator of `Felt`s +fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { + let key_elements = key.into_iter(); + let value_elements = value.into_iter(); + + key_elements.chain(value_elements) +} + +/// Compares two keys, compared element-by-element using their integer representations starting with +/// the most significant element. +fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering { + for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() { + let v1 = v1.as_int(); + let v2 = v2.as_int(); + if v1 != v2 { + return v1.cmp(&v2); + } + } + + Ordering::Equal +} diff --git a/src/merkle/smt/tests.rs b/src/merkle/smt/full/tests.rs similarity index 99% rename from src/merkle/smt/tests.rs rename to src/merkle/smt/full/tests.rs index 66b40ec5..f068e22f 100644 --- a/src/merkle/smt/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -1,6 +1,8 @@ -use crate::{merkle::MerkleStore, ONE, WORD_SIZE}; - use super::*; +use crate::{ + merkle::{EmptySubtreeRoots, MerkleStore}, + ONE, WORD_SIZE, +}; /// This test checks that inserting twice at the same key functions as expected. The test covers /// only the case where the key is alone in its leaf diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index b5f0624b..4e0e153c 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -1,365 +1,241 @@ -use core::cmp::Ordering; - use winter_math::StarkField; -use crate::hash::rpo::Rpo256; -use crate::utils::{collections::Vec, vec}; -use crate::{Felt, EMPTY_WORD}; - -use super::sparse_merkle_tree::{InnerNode, SparseMerkleTree}; -use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, LeafIndex, MerkleError, MerklePath, NodeIndex, - RpoDigest, Word, +use crate::{ + hash::rpo::{Rpo256, RpoDigest}, + Word, }; -#[cfg(test)] -mod tests; +use super::{MerkleError, MerklePath, NodeIndex, Vec}; -// CONSTANTS -// ================================================================================================ +mod full; +pub use full::{Smt, SmtLeaf, SMT_DEPTH}; -pub const SMT_DEPTH: u8 = 64; +mod simple; +pub use simple::SimpleSmt; -// SMT +// CONSTANTS // ================================================================================================ -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub struct Smt { - root: RpoDigest, - leaves: BTreeMap, - inner_nodes: BTreeMap, -} - -impl Smt { - // CONSTRUCTORS - // -------------------------------------------------------------------------------------------- - - /// Returns a new [NewSmt]. - /// - /// All leaves in the returned tree are set to [ZERO; 4]. - pub fn new() -> Self { - let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0); - - Self { - root, - leaves: BTreeMap::new(), - inner_nodes: BTreeMap::new(), - } - } - - /// Returns a new [SimpleSmt] instantiated with leaves set as specified by the provided entries. - /// - /// All leaves omitted from the entries list are set to [ZERO; 4]. - /// - /// # Errors - /// Returns an error if: - /// - The number of entries exceeds 2^63 entries. - /// - The provided entries contain multiple values for the same key. - pub fn with_entries( - entries: impl IntoIterator, - ) -> Result { - // create an empty tree - let mut tree = Self::new(); - - // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so - // entries with the empty value need additional tracking. - let mut key_set_to_zero = BTreeSet::new(); - - for (key, value) in entries { - let old_value = tree.insert(key, value); - - if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) { - return Err(MerkleError::DuplicateValuesForIndex( - LeafIndex::::from(key).value(), - )); - } - - if value == EMPTY_WORD { - key_set_to_zero.insert(key); - }; - } - Ok(tree) - } +/// Minimum supported depth. +pub const SMT_MIN_DEPTH: u8 = 1; - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- +/// Maximum supported depth. +pub const SMT_MAX_DEPTH: u8 = 64; - /// Returns the root of the tree - pub fn root(&self) -> RpoDigest { - >::root(self) - } - - /// Returns the leaf at the specified index. - pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf { - >::get_leaf(self, key) - } +// SPARSE MERKLE TREE +// ================================================================================================ - /// Returns the depth of the tree - pub const fn depth(&self) -> u8 { - SMT_DEPTH - } +/// 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 +/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's +/// value was not explicitly set, then its value is the default value. Typically, the vast majority +/// of leaves will store the default value (hence it is "sparse"), and therefore the internal +/// representation of the tree will only keep track of the leaves that have a different value from +/// the default. +/// +/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the +/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a +/// tree is just a single value, and is probably a programming mistake. +/// +/// Every key maps to one leaf. If there are as many keys as there are leaves, then +/// [Self::Leaf] should be the same type as [Self::Value], as is the case with +/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`] +/// must accomodate all keys that map to the same leaf. +/// +/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. +pub(crate) trait SparseMerkleTree { + /// The type for a key + type Key: Clone; + /// The type for a value + type Value: Clone + PartialEq; + /// The type for a leaf + type Leaf; + /// The type for an opening (i.e. a "proof") of a leaf + type Opening: From<(MerklePath, Self::Leaf)>; + + /// The default value used to compute the hash of empty leaves + const EMPTY_VALUE: Self::Value; + + // PROVIDED METHODS + // --------------------------------------------------------------------------------------------- /// Returns a Merkle path from the leaf node specified by the key to the root. /// /// The node itself is not included in the path. - pub fn open(&self, key: &RpoDigest) -> (MerklePath, SmtLeaf) { - >::open(self, key) - } + fn open(&self, key: &Self::Key) -> Self::Opening { + let leaf = self.get_leaf(key); + + let mut index: NodeIndex = { + let leaf_index: LeafIndex = Self::key_to_leaf_index(key); + leaf_index.into() + }; + + let merkle_path = { + let mut path = Vec::with_capacity(index.depth() as usize); + for _ in 0..index.depth() { + let is_right = index.is_value_odd(); + index.move_up(); + let InnerNode { left, right } = self.get_inner_node(index); + let value = if is_right { left } else { right }; + path.push(value); + } - // STATE MUTATORS - // -------------------------------------------------------------------------------------------- + MerklePath::new(path) + }; + + (merkle_path, leaf).into() + } /// Inserts a value at the specified key, returning the previous value associated with that key. /// Recall that by definition, any key that hasn't been updated is associated with - /// [`EMPTY_WORD`]. + /// [`Self::EMPTY_VALUE`]. /// /// This also recomputes all hashes between the leaf (associated with the key) and the root, /// updating the root itself. - pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word { - >::insert(self, key, value) - } - - // HELPERS - // -------------------------------------------------------------------------------------------- - - /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty - /// value, such that this is indeed an insertion. - fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option { - debug_assert_ne!(value, Self::EMPTY_VALUE); - - let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); - - match self.leaves.get_mut(&leaf_index.value()) { - Some(leaf) => leaf.insert(key, value), - None => { - self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value))); + fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value { + let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE); - None - } + // if the old value and new value are the same, there is nothing to update + if value == old_value { + return value; } - } - /// Removes key-value pair at leaf index pointed to by `key` if it exists. - fn perform_remove(&mut self, key: RpoDigest) -> Option { - let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); - - if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) { - let (old_value, is_empty) = leaf.remove(key); - if is_empty { - self.leaves.remove(&leaf_index.value()); - } - old_value - } else { - // there's nothing stored at the leaf; nothing to update - None + let leaf = self.get_leaf(&key); + let node_index = { + let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); + leaf_index.into() + }; + + self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf)); + + old_value + } + + /// Recomputes the branch nodes (including the root) from `index` all the way to the root. + /// `node_hash_at_index` is the hash of the node stored at index. + fn recompute_nodes_from_index_to_root( + &mut self, + mut index: NodeIndex, + node_hash_at_index: RpoDigest, + ) { + let mut value = node_hash_at_index; + for _ in 0..index.depth() { + let is_right = index.is_value_odd(); + index.move_up(); + let InnerNode { left, right } = self.get_inner_node(index); + let (left, right) = if is_right { (left, value) } else { (value, right) }; + self.insert_inner_node(index, InnerNode { left, right }); + value = Rpo256::merge(&[left, right]); } + self.set_root(value); } -} -impl SparseMerkleTree for Smt { - type Key = RpoDigest; - type Value = Word; - type Leaf = SmtLeaf; - type Opening = (MerklePath, SmtLeaf); + // REQUIRED METHODS + // --------------------------------------------------------------------------------------------- - const EMPTY_VALUE: Self::Value = EMPTY_WORD; + /// The root of the tree + fn root(&self) -> RpoDigest; - fn root(&self) -> RpoDigest { - self.root - } - - fn set_root(&mut self, root: RpoDigest) { - self.root = root; - } + /// Sets the root of the tree + fn set_root(&mut self, root: RpoDigest); - fn get_inner_node(&self, index: NodeIndex) -> InnerNode { - self.inner_nodes.get(&index).cloned().unwrap_or_else(|| { - let node = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1); + /// Retrieves an inner node at the given index + fn get_inner_node(&self, index: NodeIndex) -> InnerNode; - InnerNode { left: *node, right: *node } - }) - } + /// Inserts an inner node at the given index + fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode); - fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) { - self.inner_nodes.insert(index, inner_node); - } + /// Inserts a leaf node, and returns the value at the key if already exists + fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option; - fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option { - // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key` - if value != Self::EMPTY_VALUE { - self.perform_insert(key, value) - } else { - self.perform_remove(key) - } - } + /// Returns the leaf at the specified index. + fn get_leaf(&self, key: &Self::Key) -> Self::Leaf; - fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf { - let leaf_pos = LeafIndex::::from(*key).value(); + /// Returns the hash of a leaf + fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; - match self.leaves.get(&leaf_pos) { - Some(leaf) => leaf.clone(), - None => SmtLeaf::Single((*key, Self::EMPTY_VALUE)), - } - } + fn key_to_leaf_index(key: &Self::Key) -> LeafIndex; +} - fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest { - leaf.hash() - } +// INNER NODE +// ================================================================================================ - fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex { - let most_significant_felt = key[3]; - LeafIndex::new_max_depth(most_significant_felt.as_int()) - } +#[derive(Debug, Default, Clone, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] +pub(crate) struct InnerNode { + pub left: RpoDigest, + pub right: RpoDigest, } -impl Default for Smt { - fn default() -> Self { - Self::new() +impl InnerNode { + pub fn hash(&self) -> RpoDigest { + Rpo256::merge(&[self.left, self.right]) } } -// LEAF +// LEAF INDEX // ================================================================================================ -#[derive(Clone, Debug, PartialEq, Eq)] +/// The index of a leaf, at a depth known at compile-time. +#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub enum SmtLeaf { - Single((RpoDigest, Word)), - Multiple(Vec<(RpoDigest, Word)>), +pub struct LeafIndex { + index: NodeIndex, } -impl SmtLeaf { - /// Converts a leaf to a list of field elements - pub fn to_elements(&self) -> Vec { - self.clone().into_elements() - } - - /// Converts a leaf to a list of field elements - pub fn into_elements(self) -> Vec { - match self { - SmtLeaf::Single(kv_pair) => kv_to_elements(kv_pair).collect(), - SmtLeaf::Multiple(kv_pairs) => kv_pairs.into_iter().flat_map(kv_to_elements).collect(), +impl LeafIndex { + pub fn new(value: u64) -> Result { + if DEPTH < SMT_MIN_DEPTH { + return Err(MerkleError::DepthTooSmall(DEPTH)); } - } - /// Compute the hash of the leaf - pub fn hash(&self) -> RpoDigest { - match self { - SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]), - SmtLeaf::Multiple(kvs) => { - let elements: Vec = kvs.iter().copied().flat_map(kv_to_elements).collect(); - Rpo256::hash_elements(&elements) - } - } + Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? }) } - // HELPERS - // --------------------------------------------------------------------------------------------- + pub fn value(&self) -> u64 { + self.index.value() + } +} - /// Insert key-value pair into the leaf; return the previous value associated with `key`, if - /// any. - fn insert(&mut self, key: RpoDigest, value: Word) -> Option { - match self { - SmtLeaf::Single(kv_pair) => { - if kv_pair.0 == key { - // the key is already in this leaf. Update the value and return the previous - // value - let old_value = kv_pair.1; - kv_pair.1 = value; - Some(old_value) - } else { - // Another entry is present in this leaf. Transform the entry into a list - // entry, and make sure the key-value pairs are sorted by key - let mut pairs = vec![*kv_pair, (key, value)]; - pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2)); - - *self = SmtLeaf::Multiple(pairs); - - None - } - } - SmtLeaf::Multiple(kv_pairs) => { - match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { - Ok(pos) => { - let old_value = kv_pairs[pos].1; - kv_pairs[pos].1 = value; - - Some(old_value) - } - Err(pos) => { - kv_pairs.insert(pos, (key, value)); - - None - } - } - } +impl LeafIndex { + pub fn new_max_depth(value: u64) -> Self { + LeafIndex { + index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value), } } +} - /// Remove key-value pair into the leaf stored at key; return the previous value associated with - /// `key`, if any. We also return an `is_empty` flag which indicates whether the leaf became - /// empty, and must be removed from the data structure it is contained in. - fn remove(&mut self, key: RpoDigest) -> (Option, bool) { - match self { - SmtLeaf::Single((key_at_leaf, value_at_leaf)) => { - if *key_at_leaf == key { - // our key was indeed stored in the leaf, so we return the value that was stored - // in it, and indicate that the leaf should be removed - let old_value = *value_at_leaf; - - (Some(old_value), true) - } else { - // another key is stored at leaf; nothing to update - (None, false) - } - } - SmtLeaf::Multiple(kv_pairs) => { - match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { - Ok(pos) => { - let old_value = kv_pairs[pos].1; - - kv_pairs.remove(pos); - debug_assert!(!kv_pairs.is_empty()); - - if kv_pairs.len() == 1 { - // convert the leaf into `Single` - *self = SmtLeaf::Single(kv_pairs[0]); - } - - (Some(old_value), false) - } - Err(_) => { - // other keys are stored at leaf; nothing to update - (None, false) - } - } - } - } +impl From> for NodeIndex { + fn from(value: LeafIndex) -> Self { + value.index } } -// HELPER FUNCTIONS -// ================================================================================================ +impl TryFrom for LeafIndex { + type Error = MerkleError; -/// Converts a key-value tuple to an iterator of `Felt`s -fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { - let key_elements = key.into_iter(); - let value_elements = value.into_iter(); + fn try_from(node_index: NodeIndex) -> Result { + if node_index.depth() != DEPTH { + return Err(MerkleError::InvalidDepth { + expected: DEPTH, + provided: node_index.depth(), + }); + } - key_elements.chain(value_elements) + Self::new(node_index.value()) + } } -/// Compares two keys, compared element-by-element using their integer representations starting with -/// the most significant element. -fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering { - for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() { - let v1 = v1.as_int(); - let v2 = v2.as_int(); - if v1 != v2 { - return v1.cmp(&v2); - } +impl From for LeafIndex { + fn from(value: Word) -> Self { + // We use the most significant `Felt` of a `Word` as the leaf index. + Self::new_max_depth(value[3].as_int()) } +} - Ordering::Equal +impl From for LeafIndex { + fn from(value: RpoDigest) -> Self { + Word::from(value).into() + } } diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/smt/simple/mod.rs similarity index 97% rename from src/merkle/simple_smt/mod.rs rename to src/merkle/smt/simple/mod.rs index 8a167755..c097fa95 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -1,10 +1,14 @@ -use crate::EMPTY_WORD; +use crate::{ + merkle::{EmptySubtreeRoots, InnerNodeInfo, MerkleTreeDelta, StoreNode}, + utils::collections::TryApplyDiff, + EMPTY_WORD, +}; use super::{ - sparse_merkle_tree::{InnerNode, SparseMerkleTree, SMT_MAX_DEPTH}, - BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, LeafIndex, MerkleError, MerklePath, - MerkleTreeDelta, NodeIndex, RpoDigest, StoreNode, TryApplyDiff, Word, SMT_MIN_DEPTH, + InnerNode, LeafIndex, MerkleError, MerklePath, NodeIndex, RpoDigest, SparseMerkleTree, Word, + SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; +use crate::utils::collections::{BTreeMap, BTreeSet}; #[cfg(test)] mod tests; diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/smt/simple/tests.rs similarity index 98% rename from src/merkle/simple_smt/tests.rs rename to src/merkle/smt/simple/tests.rs index 3c74cef6..4ef501c7 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/smt/simple/tests.rs @@ -1,15 +1,15 @@ use super::{ - super::{InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt, EMPTY_WORD}, + super::{MerkleError, RpoDigest, SimpleSmt}, NodeIndex, }; use crate::{ hash::rpo::Rpo256, merkle::{ - digests_to_words, int_to_leaf, int_to_node, sparse_merkle_tree::SparseMerkleTree, - EmptySubtreeRoots, LeafIndex, + digests_to_words, int_to_leaf, int_to_node, smt::SparseMerkleTree, EmptySubtreeRoots, + InnerNodeInfo, LeafIndex, MerkleTree, }, utils::collections::Vec, - Word, + Word, EMPTY_WORD, }; // TEST DATA diff --git a/src/merkle/sparse_merkle_tree.rs b/src/merkle/sparse_merkle_tree.rs deleted file mode 100644 index 5cc7a0ff..00000000 --- a/src/merkle/sparse_merkle_tree.rs +++ /dev/null @@ -1,235 +0,0 @@ -use winter_math::StarkField; - -use crate::{ - hash::rpo::{Rpo256, RpoDigest}, - Word, -}; - -use super::{MerkleError, MerklePath, NodeIndex, Vec}; - -// CONSTANTS -// ================================================================================================ - -/// Minimum supported depth. -pub const SMT_MIN_DEPTH: u8 = 1; - -/// Maximum supported depth. -pub const SMT_MAX_DEPTH: u8 = 64; - -// SPARSE MERKLE TREE -// ================================================================================================ - -/// 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 -/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's -/// value was not explicitly set, then its value is the default value. Typically, the vast majority -/// of leaves will store the default value (hence it is "sparse"), and therefore the internal -/// representation of the tree will only keep track of the leaves that have a different value from -/// the default. -/// -/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the -/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a -/// tree is just a single value, and is probably a programming mistake. -/// -/// Every key maps to one leaf. If there are as many keys as there are leaves, then -/// [Self::Leaf] should be the same type as [Self::Value], as is the case with -/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`] -/// must accomodate all keys that map to the same leaf. -/// -/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs. -pub(crate) trait SparseMerkleTree { - /// The type for a key - type Key: Clone; - /// The type for a value - type Value: Clone + PartialEq; - /// The type for a leaf - type Leaf; - /// The type for an opening (i.e. a "proof") of a leaf - type Opening: From<(MerklePath, Self::Leaf)>; - - /// The default value used to compute the hash of empty leaves - const EMPTY_VALUE: Self::Value; - - // PROVIDED METHODS - // --------------------------------------------------------------------------------------------- - - /// Returns a Merkle path from the leaf node specified by the key to the root. - /// - /// The node itself is not included in the path. - fn open(&self, key: &Self::Key) -> Self::Opening { - let leaf = self.get_leaf(key); - - let mut index: NodeIndex = { - let leaf_index: LeafIndex = Self::key_to_leaf_index(key); - leaf_index.into() - }; - - let merkle_path = { - let mut path = Vec::with_capacity(index.depth() as usize); - for _ in 0..index.depth() { - let is_right = index.is_value_odd(); - index.move_up(); - let InnerNode { left, right } = self.get_inner_node(index); - let value = if is_right { left } else { right }; - path.push(value); - } - - MerklePath::new(path) - }; - - (merkle_path, leaf).into() - } - - /// Inserts a value at the specified key, returning the previous value associated with that key. - /// Recall that by definition, any key that hasn't been updated is associated with - /// [`Self::EMPTY_VALUE`]. - /// - /// This also recomputes all hashes between the leaf (associated with the key) and the root, - /// updating the root itself. - fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value { - let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE); - - // if the old value and new value are the same, there is nothing to update - if value == old_value { - return value; - } - - let leaf = self.get_leaf(&key); - let node_index = { - let leaf_index: LeafIndex = Self::key_to_leaf_index(&key); - leaf_index.into() - }; - - self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf)); - - old_value - } - - /// Recomputes the branch nodes (including the root) from `index` all the way to the root. - /// `node_hash_at_index` is the hash of the node stored at index. - fn recompute_nodes_from_index_to_root( - &mut self, - mut index: NodeIndex, - node_hash_at_index: RpoDigest, - ) { - let mut value = node_hash_at_index; - for _ in 0..index.depth() { - let is_right = index.is_value_odd(); - index.move_up(); - let InnerNode { left, right } = self.get_inner_node(index); - let (left, right) = if is_right { (left, value) } else { (value, right) }; - self.insert_inner_node(index, InnerNode { left, right }); - value = Rpo256::merge(&[left, right]); - } - self.set_root(value); - } - - // REQUIRED METHODS - // --------------------------------------------------------------------------------------------- - - /// The root of the tree - fn root(&self) -> RpoDigest; - - /// Sets the root of the tree - fn set_root(&mut self, root: RpoDigest); - - /// Retrieves an inner node at the given index - fn get_inner_node(&self, index: NodeIndex) -> InnerNode; - - /// Inserts an inner node at the given index - fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode); - - /// Inserts a leaf node, and returns the value at the key if already exists - fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option; - - /// Returns the leaf at the specified index. - fn get_leaf(&self, key: &Self::Key) -> Self::Leaf; - - /// Returns the hash of a leaf - fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest; - - fn key_to_leaf_index(key: &Self::Key) -> LeafIndex; -} - -// INNER NODE -// ================================================================================================ - -#[derive(Debug, Default, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub(crate) struct InnerNode { - pub left: RpoDigest, - pub right: RpoDigest, -} - -impl InnerNode { - pub fn hash(&self) -> RpoDigest { - Rpo256::merge(&[self.left, self.right]) - } -} - -// LEAF INDEX -// ================================================================================================ - -/// The index of a leaf, at a depth known at compile-time. -#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub struct LeafIndex { - index: NodeIndex, -} - -impl LeafIndex { - pub fn new(value: u64) -> Result { - if DEPTH < SMT_MIN_DEPTH { - return Err(MerkleError::DepthTooSmall(DEPTH)); - } - - Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? }) - } - - pub fn value(&self) -> u64 { - self.index.value() - } -} - -impl LeafIndex { - pub fn new_max_depth(value: u64) -> Self { - LeafIndex { - index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value), - } - } -} - -impl From> for NodeIndex { - fn from(value: LeafIndex) -> Self { - value.index - } -} - -impl TryFrom for LeafIndex { - type Error = MerkleError; - - fn try_from(node_index: NodeIndex) -> Result { - if node_index.depth() != DEPTH { - return Err(MerkleError::InvalidDepth { - expected: DEPTH, - provided: node_index.depth(), - }); - } - - Self::new(node_index.value()) - } -} - -impl From for LeafIndex { - fn from(value: Word) -> Self { - // We use the most significant `Felt` of a `Word` as the leaf index. - Self::new_max_depth(value[3].as_int()) - } -} - -impl From for LeafIndex { - fn from(value: RpoDigest) -> Self { - Word::from(value).into() - } -}