From 552d90429b30328f7cc399ff8a83d3385f1b1439 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philippe=20Laferri=C3=A8re?= Date: Thu, 8 Feb 2024 16:31:27 -0500 Subject: [PATCH] Remove `TieredSmt` (#277) --- README.md | 2 +- src/main.rs | 52 +- src/merkle/mod.rs | 3 - src/merkle/store/mod.rs | 12 +- src/merkle/tiered_smt/error.rs | 48 -- src/merkle/tiered_smt/mod.rs | 509 ----------------- src/merkle/tiered_smt/nodes.rs | 419 -------------- src/merkle/tiered_smt/proof.rs | 170 ------ src/merkle/tiered_smt/tests.rs | 968 -------------------------------- src/merkle/tiered_smt/values.rs | 584 ------------------- 10 files changed, 19 insertions(+), 2748 deletions(-) delete mode 100644 src/merkle/tiered_smt/error.rs delete mode 100644 src/merkle/tiered_smt/mod.rs delete mode 100644 src/merkle/tiered_smt/nodes.rs delete mode 100644 src/merkle/tiered_smt/proof.rs delete mode 100644 src/merkle/tiered_smt/tests.rs delete mode 100644 src/merkle/tiered_smt/values.rs diff --git a/README.md b/README.md index 5fcb8d3..90a7353 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ For performance benchmarks of these hash functions and their comparison to other * `PartialMerkleTree`: a partial view of a Merkle tree where some sub-trees may not be known. This is similar to a collection of Merkle paths all resolving to the same root. The length of the paths can be at most 64. * `PartialMmr`: a partial view of a Merkle mountain range structure. * `SimpleSmt`: a Sparse Merkle Tree (with no compaction), mapping 64-bit keys to 4-element values. -* `TieredSmt`: a Sparse Merkle tree (with compaction), mapping 4-element keys to 4-element values. +* `Smt`: a Sparse Merkle tree (with compaction at depth 64), mapping 4-element keys to 4-element values. The module also contains additional supporting components such as `NodeIndex`, `MerklePath`, and `MerkleError` to assist with tree indexation, opening proofs, and reporting inconsistent arguments/state. diff --git a/src/main.rs b/src/main.rs index 31264d8..7d28493 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,19 +1,14 @@ use clap::Parser; use miden_crypto::{ hash::rpo::{Rpo256, RpoDigest}, - merkle::{MerkleError, TieredSmt}, + merkle::{MerkleError, Smt}, Felt, Word, ONE, }; use rand_utils::rand_value; use std::time::Instant; #[derive(Parser, Debug)] -#[clap( - name = "Benchmark", - about = "Tiered SMT benchmark", - version, - rename_all = "kebab-case" -)] +#[clap(name = "Benchmark", about = "SMT benchmark", version, rename_all = "kebab-case")] pub struct BenchmarkCmd { /// Size of the tree #[clap(short = 's', long = "size")] @@ -21,11 +16,11 @@ pub struct BenchmarkCmd { } fn main() { - benchmark_tsmt(); + benchmark_smt(); } -/// Run a benchmark for the Tiered SMT. -pub fn benchmark_tsmt() { +/// Run a benchmark for [`Smt`]. +pub fn benchmark_smt() { let args = BenchmarkCmd::parse(); let tree_size = args.size; @@ -42,38 +37,25 @@ pub fn benchmark_tsmt() { proof_generation(&mut tree, tree_size).unwrap(); } -/// Runs the construction benchmark for the Tiered SMT, returning the constructed tree. -pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result { +/// Runs the construction benchmark for [`Smt`], returning the constructed tree. +pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result { println!("Running a construction benchmark:"); let now = Instant::now(); - let tree = TieredSmt::with_entries(entries)?; + let tree = Smt::with_entries(entries)?; let elapsed = now.elapsed(); println!( - "Constructed a TSMT with {} key-value pairs in {:.3} seconds", + "Constructed a SMT with {} key-value pairs in {:.3} seconds", size, elapsed.as_secs_f32(), ); - // Count how many nodes end up at each tier - let mut nodes_num_16_32_48 = (0, 0, 0); - - tree.upper_leaf_nodes().for_each(|(index, _)| match index.depth() { - 16 => nodes_num_16_32_48.0 += 1, - 32 => nodes_num_16_32_48.1 += 1, - 48 => nodes_num_16_32_48.2 += 1, - _ => unreachable!(), - }); - - println!("Number of nodes on depth 16: {}", nodes_num_16_32_48.0); - println!("Number of nodes on depth 32: {}", nodes_num_16_32_48.1); - println!("Number of nodes on depth 48: {}", nodes_num_16_32_48.2); - println!("Number of nodes on depth 64: {}\n", tree.bottom_leaves().count()); + println!("Number of leaf nodes: {}\n", tree.leaves().count()); Ok(tree) } -/// Runs the insertion benchmark for the Tiered SMT. -pub fn insertion(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleError> { +/// Runs the insertion benchmark for the [`Smt`]. +pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> { println!("Running an insertion benchmark:"); let mut insertion_times = Vec::new(); @@ -89,7 +71,7 @@ pub fn insertion(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleError> { } println!( - "An average insertion time measured by 20 inserts into a TSMT with {} key-value pairs is {:.3} milliseconds\n", + "An average insertion time measured by 20 inserts into a SMT with {} key-value pairs is {:.3} milliseconds\n", size, // calculate the average by dividing by 20 and convert to milliseconds by multiplying by // 1000. As a result, we can only multiply by 50 @@ -99,8 +81,8 @@ pub fn insertion(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleError> { Ok(()) } -/// Runs the proof generation benchmark for the Tiered SMT. -pub fn proof_generation(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleError> { +/// Runs the proof generation benchmark for the [`Smt`]. +pub fn proof_generation(tree: &mut Smt, size: u64) -> Result<(), MerkleError> { println!("Running a proof generation benchmark:"); let mut insertion_times = Vec::new(); @@ -111,13 +93,13 @@ pub fn proof_generation(tree: &mut TieredSmt, size: u64) -> Result<(), MerkleErr tree.insert(test_key, test_value); let now = Instant::now(); - let _proof = tree.prove(test_key); + let _proof = tree.open(&test_key); let elapsed = now.elapsed(); insertion_times.push(elapsed.as_secs_f32()); } println!( - "An average proving time measured by 20 value proofs in a TSMT with {} key-value pairs in {:.3} microseconds", + "An average proving time measured by 20 value proofs in a SMT with {} key-value pairs in {:.3} microseconds", size, // calculate the average by dividing by 20 and convert to microseconds by multiplying by // 1000000. As a result, we can only multiply by 50000 diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 1570997..9b8c7c5 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -27,9 +27,6 @@ pub use smt::{ SMT_MAX_DEPTH, SMT_MIN_DEPTH, }; -mod tiered_smt; -pub use tiered_smt::{TieredSmt, TieredSmtProof, TieredSmtProofError}; - mod mmr; pub use mmr::{InOrderIndex, Mmr, MmrDelta, MmrError, MmrPeaks, MmrProof, PartialMmr}; diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index a06e292..b060088 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,7 +1,7 @@ use super::{ mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, KvMap, MerkleError, MerklePath, MerkleTree, NodeIndex, PartialMerkleTree, RecordingMap, RootPath, Rpo256, RpoDigest, SimpleSmt, - Smt, TieredSmt, ValuePath, Vec, + Smt, ValuePath, Vec, }; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use core::borrow::Borrow; @@ -361,9 +361,6 @@ impl> MerkleStore { // if the node is not in the store assume it is a leaf } else { - // assert that if we have a leaf that is not at the max depth then it must be - // at the depth of one of the tiers of an TSMT. - debug_assert!(TieredSmt::TIER_DEPTHS[..3].contains(&index.depth())); return Some((index, node_hash)); } } @@ -511,13 +508,6 @@ impl> From<&Mmr> for MerkleStore { } } -impl> From<&TieredSmt> for MerkleStore { - fn from(value: &TieredSmt) -> Self { - let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); - Self { nodes } - } -} - impl> From<&PartialMerkleTree> for MerkleStore { fn from(value: &PartialMerkleTree) -> Self { let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect(); diff --git a/src/merkle/tiered_smt/error.rs b/src/merkle/tiered_smt/error.rs deleted file mode 100644 index 92f7dea..0000000 --- a/src/merkle/tiered_smt/error.rs +++ /dev/null @@ -1,48 +0,0 @@ -use core::fmt::Display; - -#[derive(Debug, PartialEq, Eq)] -pub enum TieredSmtProofError { - EntriesEmpty, - EmptyValueNotAllowed, - MismatchedPrefixes(u64, u64), - MultipleEntriesOutsideLastTier, - NotATierPath(u8), - PathTooLong, -} - -impl Display for TieredSmtProofError { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - match self { - TieredSmtProofError::EntriesEmpty => { - write!(f, "Missing entries for tiered sparse merkle tree proof") - } - TieredSmtProofError::EmptyValueNotAllowed => { - write!( - f, - "The empty value [0, 0, 0, 0] is not allowed inside a tiered sparse merkle tree" - ) - } - TieredSmtProofError::MismatchedPrefixes(first, second) => { - write!(f, "Not all leaves have the same prefix. First {first} second {second}") - } - TieredSmtProofError::MultipleEntriesOutsideLastTier => { - write!(f, "Multiple entries are only allowed for the last tier (depth 64)") - } - TieredSmtProofError::NotATierPath(got) => { - write!( - f, - "Path length does not correspond to a tier. Got {got} Expected one of 16, 32, 48, 64" - ) - } - TieredSmtProofError::PathTooLong => { - write!( - f, - "Path longer than maximum depth of 64 for tiered sparse merkle tree proof" - ) - } - } - } -} - -#[cfg(feature = "std")] -impl std::error::Error for TieredSmtProofError {} diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs deleted file mode 100644 index 0dff818..0000000 --- a/src/merkle/tiered_smt/mod.rs +++ /dev/null @@ -1,509 +0,0 @@ -use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, - Rpo256, RpoDigest, StarkField, Vec, Word, -}; -use crate::utils::vec; -use core::{cmp, ops::Deref}; - -mod nodes; -use nodes::NodeStore; - -mod values; -use values::ValueStore; - -mod proof; -pub use proof::TieredSmtProof; - -mod error; -pub use error::TieredSmtProofError; - -#[cfg(test)] -mod tests; - -// TIERED SPARSE MERKLE TREE -// ================================================================================================ - -/// Tiered (compacted) Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and -/// values are represented by 4 field elements. -/// -/// Leaves in the tree can exist only on specific depths called "tiers". These depths are: 16, 32, -/// 48, and 64. Initially, when a tree is empty, it is equivalent to an empty Sparse Merkle tree -/// of depth 64 (i.e., leaves at depth 64 are set to [ZERO; 4]). As non-empty values are inserted -/// into the tree they are added to the first available tier. -/// -/// For example, when the first key-value pair is inserted, it will be stored in a node at depth -/// 16 such that the 16 most significant bits of the key determine the position of the node at -/// depth 16. If another value with a key sharing the same 16-bit prefix is inserted, both values -/// move into the next tier (depth 32). This process is repeated until values end up at the bottom -/// tier (depth 64). If multiple values have keys with a common 64-bit prefix, such key-value pairs -/// are stored in a sorted list at the bottom tier. -/// -/// To differentiate between internal and leaf nodes, node values are computed as follows: -/// - Internal nodes: hash(left_child, right_child). -/// - Leaf node at depths 16, 32, or 64: hash(key, value, domain=depth). -/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n], domain=64). -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub struct TieredSmt { - root: RpoDigest, - nodes: NodeStore, - values: ValueStore, -} - -impl TieredSmt { - // CONSTANTS - // -------------------------------------------------------------------------------------------- - - /// The number of levels between tiers. - pub const TIER_SIZE: u8 = 16; - - /// Depths at which leaves can exist in a tiered SMT. - pub const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64]; - - /// Maximum node depth. This is also the bottom tier of the tree. - pub const MAX_DEPTH: u8 = 64; - - /// Value of an empty leaf. - pub const EMPTY_VALUE: Word = super::EMPTY_WORD; - - // CONSTRUCTORS - // -------------------------------------------------------------------------------------------- - - /// Returns a new [TieredSmt] instantiated with the specified key-value pairs. - /// - /// # Errors - /// Returns an error if the provided entries contain multiple values for the same key. - pub fn with_entries(entries: R) -> Result - where - R: IntoIterator, - I: Iterator + ExactSizeIterator, - { - // create an empty tree - let mut tree = Self::default(); - - // append leaves to the tree returning an error if a duplicate entry for the same key - // is found - let mut empty_entries = BTreeSet::new(); - for (key, value) in entries { - let old_value = tree.insert(key, value); - if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) { - return Err(MerkleError::DuplicateValuesForKey(key)); - } - // if we've processed an empty entry, add the key to the set of empty entry keys, and - // if this key was already in the set, return an error - if value == Self::EMPTY_VALUE && !empty_entries.insert(key) { - return Err(MerkleError::DuplicateValuesForKey(key)); - } - } - Ok(tree) - } - - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- - - /// Returns the root of this Merkle tree. - pub const fn root(&self) -> RpoDigest { - self.root - } - - /// Returns a node at the specified index. - /// - /// # Errors - /// Returns an error if: - /// - The specified index depth is 0 or greater than 64. - /// - The node with the specified index does not exists in the Merkle tree. This is possible - /// when a leaf node with the same index prefix exists at a tier higher than the requested - /// node. - pub fn get_node(&self, index: NodeIndex) -> Result { - self.nodes.get_node(index) - } - - /// Returns a Merkle path from the node at the specified index to the root. - /// - /// The node itself is not included in the path. - /// - /// # Errors - /// Returns an error if: - /// - The specified index depth is 0 or greater than 64. - /// - The node with the specified index does not exists in the Merkle tree. This is possible - /// when a leaf node with the same index prefix exists at a tier higher than the node to - /// which the path is requested. - pub fn get_path(&self, index: NodeIndex) -> Result { - self.nodes.get_path(index) - } - - /// Returns the value associated with the specified key. - /// - /// If nothing was inserted into this tree for the specified key, [ZERO; 4] is returned. - pub fn get_value(&self, key: RpoDigest) -> Word { - match self.values.get(&key) { - Some(value) => *value, - None => Self::EMPTY_VALUE, - } - } - - /// Returns a proof for a key-value pair defined by the specified key. - /// - /// The proof can be used to attest membership of this key-value pair in a Tiered Sparse Merkle - /// Tree defined by the same root as this tree. - pub fn prove(&self, key: RpoDigest) -> TieredSmtProof { - let (path, index, leaf_exists) = self.nodes.get_proof(&key); - - let entries = if index.depth() == Self::MAX_DEPTH { - match self.values.get_all(index.value()) { - Some(entries) => entries, - None => vec![(key, Self::EMPTY_VALUE)], - } - } else if leaf_exists { - let entry = - self.values.get_first(index_to_prefix(&index)).expect("leaf entry not found"); - debug_assert_eq!(entry.0, key); - vec![*entry] - } else { - vec![(key, Self::EMPTY_VALUE)] - }; - - TieredSmtProof::new(path, entries).expect("Bug detected, TSMT produced invalid proof") - } - - // STATE MUTATORS - // -------------------------------------------------------------------------------------------- - - /// Inserts the provided value into the tree under the specified key and returns the value - /// previously stored under this key. - /// - /// If the value for the specified key was not previously set, [ZERO; 4] is returned. - pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word { - // if an empty value is being inserted, remove the leaf node to make it look as if the - // value was never inserted - if value == Self::EMPTY_VALUE { - return self.remove_leaf_node(key); - } - - // insert the value into the value store, and if the key was already in the store, update - // it with the new value - if let Some(old_value) = self.values.insert(key, value) { - if old_value != value { - // if the new value is different from the old value, determine the location of - // the leaf node for this key, build the node, and update the root - let (index, leaf_exists) = self.nodes.get_leaf_index(&key); - debug_assert!(leaf_exists); - let node = self.build_leaf_node(index, key, value); - self.root = self.nodes.update_leaf_node(index, node); - } - return old_value; - }; - - // determine the location for the leaf node; this index could have 3 different meanings: - // - it points to a root of an empty subtree or an empty node at depth 64; in this case, - // we can replace the node with the value node immediately. - // - it points to an existing leaf at the bottom tier (i.e., depth = 64); in this case, - // we need to process update the bottom leaf. - // - it points to an existing leaf node for a different key with the same prefix (same - // key case was handled above); in this case, we need to move the leaf to a lower tier - let (index, leaf_exists) = self.nodes.get_leaf_index(&key); - - self.root = if leaf_exists && index.depth() == Self::MAX_DEPTH { - // returned index points to a leaf at the bottom tier - let node = self.build_leaf_node(index, key, value); - self.nodes.update_leaf_node(index, node) - } else if leaf_exists { - // returned index points to a leaf for a different key with the same prefix - - // get the key-value pair for the key with the same prefix; since the key-value - // pair has already been inserted into the value store, we need to filter it out - // when looking for the other key-value pair - let (other_key, other_value) = self - .values - .get_first_filtered(index_to_prefix(&index), &key) - .expect("other key-value pair not found"); - - // determine how far down the tree should we move the leaves - let common_prefix_len = get_common_prefix_tier_depth(&key, other_key); - let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH); - - // compute node locations for new and existing key-value paris - let new_index = LeafNodeIndex::from_key(&key, depth); - let other_index = LeafNodeIndex::from_key(other_key, depth); - - // compute node values for the new and existing key-value pairs - let new_node = self.build_leaf_node(new_index, key, value); - let other_node = self.build_leaf_node(other_index, *other_key, *other_value); - - // replace the leaf located at index with a subtree containing nodes for new and - // existing key-value paris - self.nodes.replace_leaf_with_subtree( - index, - [(new_index, new_node), (other_index, other_node)], - ) - } else { - // returned index points to an empty subtree or an empty leaf at the bottom tier - let node = self.build_leaf_node(index, key, value); - self.nodes.insert_leaf_node(index, node) - }; - - Self::EMPTY_VALUE - } - - // ITERATORS - // -------------------------------------------------------------------------------------------- - - /// Returns an iterator over all key-value pairs in this [TieredSmt]. - pub fn iter(&self) -> impl Iterator { - self.values.iter() - } - - /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16 - /// 32, 48, or 64). - /// - /// The iterator order is unspecified. - pub fn inner_nodes(&self) -> impl Iterator + '_ { - self.nodes.inner_nodes() - } - - /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt] - /// where each yielded item is a (node, key, value) tuple. - /// - /// The iterator order is unspecified. - pub fn upper_leaves(&self) -> impl Iterator + '_ { - self.nodes.upper_leaves().map(|(index, node)| { - let key_prefix = index_to_prefix(index); - let (key, value) = self.values.get_first(key_prefix).expect("upper leaf not found"); - debug_assert_eq!(*index, LeafNodeIndex::from_key(key, index.depth()).into()); - (*node, *key, *value) - }) - } - - /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt] - /// where each yielded item is a (node_index, value) tuple. - pub fn upper_leaf_nodes(&self) -> impl Iterator { - self.nodes.upper_leaves() - } - - /// Returns an iterator over bottom leaves (i.e., depth = 64) of this [TieredSmt]. - /// - /// Each yielded item consists of the hash of the leaf and its contents, where contents is - /// a vector containing key-value pairs of entries storied in this leaf. - /// - /// The iterator order is unspecified. - pub fn bottom_leaves(&self) -> impl Iterator)> + '_ { - self.nodes.bottom_leaves().map(|(&prefix, node)| { - let values = self.values.get_all(prefix).expect("bottom leaf not found"); - (*node, values) - }) - } - - // HELPER METHODS - // -------------------------------------------------------------------------------------------- - - /// Removes the node holding the key-value pair for the specified key from this tree, and - /// returns the value associated with the specified key. - /// - /// If no value was associated with the specified key, [ZERO; 4] is returned. - fn remove_leaf_node(&mut self, key: RpoDigest) -> Word { - // remove the key-value pair from the value store; if no value was associated with the - // specified key, return. - let old_value = match self.values.remove(&key) { - Some(old_value) => old_value, - None => return Self::EMPTY_VALUE, - }; - - // determine the location of the leaf holding the key-value pair to be removed - let (index, leaf_exists) = self.nodes.get_leaf_index(&key); - debug_assert!(leaf_exists); - - // if the leaf is at the bottom tier and after removing the key-value pair from it, the - // leaf is still not empty, we either just update it, or move it up to a higher tier (if - // the leaf doesn't have siblings at lower tiers) - if index.depth() == Self::MAX_DEPTH { - if let Some(entries) = self.values.get_all(index.value()) { - // if there is only one key-value pair left at the bottom leaf, and it can be - // moved up to a higher tier, truncate the branch and return - if entries.len() == 1 { - let new_depth = self.nodes.get_last_single_child_parent_depth(index.value()); - if new_depth != Self::MAX_DEPTH { - let node = hash_upper_leaf(entries[0].0, entries[0].1, new_depth); - self.root = self.nodes.truncate_branch(index.value(), new_depth, node); - return old_value; - } - } - - // otherwise just recompute the leaf hash and update the leaf node - let node = hash_bottom_leaf(&entries); - self.root = self.nodes.update_leaf_node(index, node); - return old_value; - }; - } - - // if the removed key-value pair has a lone sibling at the current tier with a root at - // higher tier, we need to move the sibling to a higher tier - if let Some((sib_key, sib_val, new_sib_index)) = self.values.get_lone_sibling(index) { - // determine the current index of the sibling node - let sib_index = LeafNodeIndex::from_key(sib_key, index.depth()); - debug_assert!(sib_index.depth() > new_sib_index.depth()); - - // compute node value for the new location of the sibling leaf and replace the subtree - // with this leaf node - let node = self.build_leaf_node(new_sib_index, *sib_key, *sib_val); - let new_sib_depth = new_sib_index.depth(); - self.root = self.nodes.replace_subtree_with_leaf(index, sib_index, new_sib_depth, node); - } else { - // if the removed key-value pair did not have a sibling at the current tier with a - // root at higher tiers, just clear the leaf node - self.root = self.nodes.clear_leaf_node(index); - } - - old_value - } - - /// Builds and returns a leaf node value for the node located as the specified index. - /// - /// This method assumes that the key-value pair for the node has already been inserted into - /// the value store, however, for depths 16, 32, and 48, the node is computed directly from - /// the passed-in values (for depth 64, the value store is queried to get all the key-value - /// pairs located at the specified index). - fn build_leaf_node(&self, index: LeafNodeIndex, key: RpoDigest, value: Word) -> RpoDigest { - let depth = index.depth(); - - // insert the key into index-key map and compute the new value of the node - if index.depth() == Self::MAX_DEPTH { - // for the bottom tier, we add the key-value pair to the existing leaf, or create a - // new leaf with this key-value pair - let values = self.values.get_all(index.value()).unwrap(); - hash_bottom_leaf(&values) - } else { - debug_assert_eq!(self.values.get_first(index_to_prefix(&index)), Some(&(key, value))); - hash_upper_leaf(key, value, depth) - } - } -} - -impl Default for TieredSmt { - fn default() -> Self { - let root = EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0]; - Self { - root, - nodes: NodeStore::new(root), - values: ValueStore::default(), - } - } -} - -// LEAF NODE INDEX -// ================================================================================================ -/// A wrapper around [NodeIndex] to provide type-safe references to nodes at depths 16, 32, 48, and -/// 64. -#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] -pub struct LeafNodeIndex(NodeIndex); - -impl LeafNodeIndex { - /// Returns a new [LeafNodeIndex] instantiated from the provided [NodeIndex]. - /// - /// In debug mode, panics if index depth is not 16, 32, 48, or 64. - pub fn new(index: NodeIndex) -> Self { - // check if the depth is 16, 32, 48, or 64; this works because for a valid depth, - // depth - 16, can be 0, 16, 32, or 48 - i.e., the value is either 0 or any of the 4th - // or 5th bits are set. We can test for this by computing a bitwise AND with a value - // which has all but the 4th and 5th bits set (which is !48). - debug_assert_eq!(((index.depth() - 16) & !48), 0, "invalid tier depth {}", index.depth()); - Self(index) - } - - /// Returns a new [LeafNodeIndex] instantiated from the specified key inserted at the specified - /// depth. - /// - /// The value for the key is computed by taking n most significant bits from the most significant - /// element of the key, where n is the specified depth. - pub fn from_key(key: &RpoDigest, depth: u8) -> Self { - let mse = get_key_prefix(key); - Self::new(NodeIndex::new_unchecked(depth, mse >> (TieredSmt::MAX_DEPTH - depth))) - } - - /// Returns a new [LeafNodeIndex] instantiated for testing purposes. - #[cfg(test)] - pub fn make(depth: u8, value: u64) -> Self { - Self::new(NodeIndex::make(depth, value)) - } - - /// Traverses towards the root until the specified depth is reached. - /// - /// The new depth must be a valid tier depth - i.e., 16, 32, 48, or 64. - pub fn move_up_to(&mut self, depth: u8) { - debug_assert_eq!(((depth - 16) & !48), 0, "invalid tier depth: {depth}"); - self.0.move_up_to(depth); - } -} - -impl Deref for LeafNodeIndex { - type Target = NodeIndex; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl From for LeafNodeIndex { - fn from(value: NodeIndex) -> Self { - Self::new(value) - } -} - -impl From for NodeIndex { - fn from(value: LeafNodeIndex) -> Self { - value.0 - } -} - -// HELPER FUNCTIONS -// ================================================================================================ - -/// Returns the value representing the 64 most significant bits of the specified key. -fn get_key_prefix(key: &RpoDigest) -> u64 { - Word::from(key)[3].as_int() -} - -/// Returns the index value shifted to be in the most significant bit positions of the returned -/// u64 value. -fn index_to_prefix(index: &NodeIndex) -> u64 { - index.value() << (TieredSmt::MAX_DEPTH - index.depth()) -} - -/// Returns tiered common prefix length between the most significant elements of the provided keys. -/// -/// Specifically: -/// - returns 64 if the most significant elements are equal. -/// - returns 48 if the common prefix is between 48 and 63 bits. -/// - returns 32 if the common prefix is between 32 and 47 bits. -/// - returns 16 if the common prefix is between 16 and 31 bits. -/// - returns 0 if the common prefix is fewer than 16 bits. -fn get_common_prefix_tier_depth(key1: &RpoDigest, key2: &RpoDigest) -> u8 { - let e1 = get_key_prefix(key1); - let e2 = get_key_prefix(key2); - let ex = (e1 ^ e2).leading_zeros() as u8; - (ex / 16) * 16 -} - -/// Computes node value for leaves at tiers 16, 32, or 48. -/// -/// Node value is computed as: hash(key || value, domain = depth). -pub fn hash_upper_leaf(key: RpoDigest, value: Word, depth: u8) -> RpoDigest { - const NUM_UPPER_TIERS: usize = TieredSmt::TIER_DEPTHS.len() - 1; - debug_assert!(TieredSmt::TIER_DEPTHS[..NUM_UPPER_TIERS].contains(&depth)); - Rpo256::merge_in_domain(&[key, value.into()], depth.into()) -} - -/// Computes node value for leaves at the bottom tier (depth 64). -/// -/// Node value is computed as: hash([key_0, value_0, ..., key_n, value_n], domain=64). -/// -/// TODO: when hashing in domain is implemented for `hash_elements()`, combine this function with -/// `hash_upper_leaf()` function. -pub fn hash_bottom_leaf(values: &[(RpoDigest, Word)]) -> RpoDigest { - let mut elements = Vec::with_capacity(values.len() * 8); - for (key, val) in values.iter() { - elements.extend_from_slice(key.as_elements()); - elements.extend_from_slice(val.as_slice()); - } - // TODO: hash in domain - Rpo256::hash_elements(&elements) -} diff --git a/src/merkle/tiered_smt/nodes.rs b/src/merkle/tiered_smt/nodes.rs deleted file mode 100644 index 1bb34df..0000000 --- a/src/merkle/tiered_smt/nodes.rs +++ /dev/null @@ -1,419 +0,0 @@ -use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, LeafNodeIndex, MerkleError, MerklePath, - NodeIndex, Rpo256, RpoDigest, Vec, -}; - -// CONSTANTS -// ================================================================================================ - -/// The number of levels between tiers. -const TIER_SIZE: u8 = super::TieredSmt::TIER_SIZE; - -/// Depths at which leaves can exist in a tiered SMT. -const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS; - -/// Maximum node depth. This is also the bottom tier of the tree. -const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH; - -// NODE STORE -// ================================================================================================ - -/// A store of nodes for a Tiered Sparse Merkle tree. -/// -/// The store contains information about all nodes as well as information about which of the nodes -/// represent leaf nodes in a Tiered Sparse Merkle tree. In the current implementation, [BTreeSet]s -/// are used to determine the position of the leaves in the tree. -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub struct NodeStore { - nodes: BTreeMap, - upper_leaves: BTreeSet, - bottom_leaves: BTreeSet, -} - -impl NodeStore { - // CONSTRUCTOR - // -------------------------------------------------------------------------------------------- - /// Returns a new instance of [NodeStore] instantiated with the specified root node. - /// - /// Root node is assumed to be a root of an empty sparse Merkle tree. - pub fn new(root_node: RpoDigest) -> Self { - let mut nodes = BTreeMap::default(); - nodes.insert(NodeIndex::root(), root_node); - - Self { - nodes, - upper_leaves: BTreeSet::default(), - bottom_leaves: BTreeSet::default(), - } - } - - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- - - /// Returns a node at the specified index. - /// - /// # Errors - /// Returns an error if: - /// - The specified index depth is 0 or greater than 64. - /// - The node with the specified index does not exists in the Merkle tree. This is possible - /// when a leaf node with the same index prefix exists at a tier higher than the requested - /// node. - pub fn get_node(&self, index: NodeIndex) -> Result { - self.validate_node_access(index)?; - Ok(self.get_node_unchecked(&index)) - } - - /// Returns a Merkle path from the node at the specified index to the root. - /// - /// The node itself is not included in the path. - /// - /// # Errors - /// Returns an error if: - /// - The specified index depth is 0 or greater than 64. - /// - The node with the specified index does not exists in the Merkle tree. This is possible - /// when a leaf node with the same index prefix exists at a tier higher than the node to - /// which the path is requested. - pub fn get_path(&self, mut index: NodeIndex) -> Result { - self.validate_node_access(index)?; - - let mut path = Vec::with_capacity(index.depth() as usize); - for _ in 0..index.depth() { - let node = self.get_node_unchecked(&index.sibling()); - path.push(node); - index.move_up(); - } - - Ok(path.into()) - } - - /// Returns a Merkle path to the node specified by the key together with a flag indicating, - /// whether this node is a leaf at depths 16, 32, or 48. - pub fn get_proof(&self, key: &RpoDigest) -> (MerklePath, NodeIndex, bool) { - let (index, leaf_exists) = self.get_leaf_index(key); - let index: NodeIndex = index.into(); - let path = self.get_path(index).expect("failed to retrieve Merkle path for a node index"); - (path, index, leaf_exists) - } - - /// Returns an index at which a leaf node for the specified key should be inserted. - /// - /// The second value in the returned tuple is set to true if the node at the returned index - /// is already a leaf node. - pub fn get_leaf_index(&self, key: &RpoDigest) -> (LeafNodeIndex, bool) { - // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if - // a node at any of the tiers is either a leaf or a root of an empty subtree. - const NUM_UPPER_TIERS: usize = TIER_DEPTHS.len() - 1; - for &tier_depth in TIER_DEPTHS[..NUM_UPPER_TIERS].iter() { - let index = LeafNodeIndex::from_key(key, tier_depth); - if self.upper_leaves.contains(&index) { - return (index, true); - } else if !self.nodes.contains_key(&index) { - return (index, false); - } - } - - // if we got here, that means all of the nodes checked so far are internal nodes, and - // the new node would need to be inserted in the bottom tier. - let index = LeafNodeIndex::from_key(key, MAX_DEPTH); - (index, self.bottom_leaves.contains(&index.value())) - } - - /// Traverses the tree up from the bottom tier starting at the specified leaf index and - /// returns the depth of the first node which hash more than one child. The returned depth - /// is rounded up to the next tier. - pub fn get_last_single_child_parent_depth(&self, leaf_index: u64) -> u8 { - let mut index = NodeIndex::new_unchecked(MAX_DEPTH, leaf_index); - - for _ in (TIER_DEPTHS[0]..MAX_DEPTH).rev() { - let sibling_index = index.sibling(); - if self.nodes.contains_key(&sibling_index) { - break; - } - index.move_up(); - } - - let tier = (index.depth() - 1) / TIER_SIZE; - TIER_DEPTHS[tier as usize] - } - - // ITERATORS - // -------------------------------------------------------------------------------------------- - - /// Returns an iterator over all inner nodes of the Tiered Sparse Merkle tree (i.e., nodes not - /// at depths 16 32, 48, or 64). - /// - /// The iterator order is unspecified. - pub fn inner_nodes(&self) -> impl Iterator + '_ { - self.nodes.iter().filter_map(|(index, node)| { - if self.is_internal_node(index) { - Some(InnerNodeInfo { - value: *node, - left: self.get_node_unchecked(&index.left_child()), - right: self.get_node_unchecked(&index.right_child()), - }) - } else { - None - } - }) - } - - /// Returns an iterator over the upper leaves (i.e., leaves with depths 16, 32, 48) of the - /// Tiered Sparse Merkle tree. - pub fn upper_leaves(&self) -> impl Iterator { - self.upper_leaves.iter().map(|index| (index, &self.nodes[index])) - } - - /// Returns an iterator over the bottom leaves (i.e., leaves with depth 64) of the Tiered - /// Sparse Merkle tree. - pub fn bottom_leaves(&self) -> impl Iterator { - self.bottom_leaves.iter().map(|value| { - let index = NodeIndex::new_unchecked(MAX_DEPTH, *value); - (value, &self.nodes[&index]) - }) - } - - // STATE MUTATORS - // -------------------------------------------------------------------------------------------- - - /// Replaces the leaf node at the specified index with a tree consisting of two leaves located - /// at the specified indexes. Recomputes and returns the new root. - pub fn replace_leaf_with_subtree( - &mut self, - leaf_index: LeafNodeIndex, - subtree_leaves: [(LeafNodeIndex, RpoDigest); 2], - ) -> RpoDigest { - debug_assert!(self.is_non_empty_leaf(&leaf_index)); - debug_assert!(!is_empty_root(&subtree_leaves[0].1)); - debug_assert!(!is_empty_root(&subtree_leaves[1].1)); - debug_assert_eq!(subtree_leaves[0].0.depth(), subtree_leaves[1].0.depth()); - debug_assert!(leaf_index.depth() < subtree_leaves[0].0.depth()); - - self.upper_leaves.remove(&leaf_index); - - if subtree_leaves[0].0 == subtree_leaves[1].0 { - // if the subtree is for a single node at depth 64, we only need to insert one node - debug_assert_eq!(subtree_leaves[0].0.depth(), MAX_DEPTH); - debug_assert_eq!(subtree_leaves[0].1, subtree_leaves[1].1); - self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1) - } else { - self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1); - self.insert_leaf_node(subtree_leaves[1].0, subtree_leaves[1].1) - } - } - - /// Replaces a subtree containing the retained and the removed leaf nodes, with a leaf node - /// containing the retained leaf. - /// - /// This has the effect of deleting the the node at the `removed_leaf` index from the tree, - /// moving the node at the `retained_leaf` index up to the tier specified by `new_depth`. - pub fn replace_subtree_with_leaf( - &mut self, - removed_leaf: LeafNodeIndex, - retained_leaf: LeafNodeIndex, - new_depth: u8, - node: RpoDigest, - ) -> RpoDigest { - debug_assert!(!is_empty_root(&node)); - debug_assert!(self.is_non_empty_leaf(&removed_leaf)); - debug_assert!(self.is_non_empty_leaf(&retained_leaf)); - debug_assert_eq!(removed_leaf.depth(), retained_leaf.depth()); - debug_assert!(removed_leaf.depth() > new_depth); - - // remove the branches leading up to the tier to which the retained leaf is to be moved - self.remove_branch(removed_leaf, new_depth); - self.remove_branch(retained_leaf, new_depth); - - // compute the index of the common root for retained and removed leaves - let mut new_index = retained_leaf; - new_index.move_up_to(new_depth); - - // insert the node at the root index - self.insert_leaf_node(new_index, node) - } - - /// Inserts the specified node at the specified index; recomputes and returns the new root - /// of the Tiered Sparse Merkle tree. - /// - /// This method assumes that the provided node is a non-empty value, and that there is no node - /// at the specified index. - pub fn insert_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest { - debug_assert!(!is_empty_root(&node)); - debug_assert_eq!(self.nodes.get(&index), None); - - // mark the node as the leaf - if index.depth() == MAX_DEPTH { - self.bottom_leaves.insert(index.value()); - } else { - self.upper_leaves.insert(index.into()); - }; - - // insert the node and update the path from the node to the root - let mut index: NodeIndex = index.into(); - for _ in 0..index.depth() { - self.nodes.insert(index, node); - let sibling = self.get_node_unchecked(&index.sibling()); - node = Rpo256::merge(&index.build_node(node, sibling)); - index.move_up(); - } - - // update the root - self.nodes.insert(NodeIndex::root(), node); - node - } - - /// Updates the node at the specified index with the specified node value; recomputes and - /// returns the new root of the Tiered Sparse Merkle tree. - /// - /// This method can accept `node` as either an empty or a non-empty value. - pub fn update_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest { - debug_assert!(self.is_non_empty_leaf(&index)); - - // if the value we are updating the node to is a root of an empty tree, clear the leaf - // flag for this node - if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] { - if index.depth() == MAX_DEPTH { - self.bottom_leaves.remove(&index.value()); - } else { - self.upper_leaves.remove(&index); - } - } else { - debug_assert!(!is_empty_root(&node)); - } - - // update the path from the node to the root - let mut index: NodeIndex = index.into(); - for _ in 0..index.depth() { - if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] { - self.nodes.remove(&index); - } else { - self.nodes.insert(index, node); - } - - let sibling = self.get_node_unchecked(&index.sibling()); - node = Rpo256::merge(&index.build_node(node, sibling)); - index.move_up(); - } - - // update the root - self.nodes.insert(NodeIndex::root(), node); - node - } - - /// Replaces the leaf node at the specified index with a root of an empty subtree; recomputes - /// and returns the new root of the Tiered Sparse Merkle tree. - pub fn clear_leaf_node(&mut self, index: LeafNodeIndex) -> RpoDigest { - debug_assert!(self.is_non_empty_leaf(&index)); - let node = EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize]; - self.update_leaf_node(index, node) - } - - /// Truncates a branch starting with specified leaf at the bottom tier to new depth. - /// - /// This involves removing the part of the branch below the new depth, and then inserting a new - /// // node at the new depth. - pub fn truncate_branch( - &mut self, - leaf_index: u64, - new_depth: u8, - node: RpoDigest, - ) -> RpoDigest { - debug_assert!(self.bottom_leaves.contains(&leaf_index)); - - let mut leaf_index = LeafNodeIndex::new(NodeIndex::new_unchecked(MAX_DEPTH, leaf_index)); - self.remove_branch(leaf_index, new_depth); - - leaf_index.move_up_to(new_depth); - self.insert_leaf_node(leaf_index, node) - } - - // HELPER METHODS - // -------------------------------------------------------------------------------------------- - - /// Returns true if the node at the specified index is a leaf node. - fn is_non_empty_leaf(&self, index: &LeafNodeIndex) -> bool { - if index.depth() == MAX_DEPTH { - self.bottom_leaves.contains(&index.value()) - } else { - self.upper_leaves.contains(index) - } - } - - /// Returns true if the node at the specified index is an internal node - i.e., there is - /// no leaf at that node and the node does not belong to the bottom tier. - fn is_internal_node(&self, index: &NodeIndex) -> bool { - if index.depth() == MAX_DEPTH { - false - } else { - !self.upper_leaves.contains(index) - } - } - - /// Checks if the specified index is valid in the context of this Merkle tree. - /// - /// # Errors - /// Returns an error if: - /// - The specified index depth is 0 or greater than 64. - /// - The node for the specified index does not exists in the Merkle tree. This is possible - /// when an ancestors of the specified index is a leaf node. - fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> { - if index.is_root() { - return Err(MerkleError::DepthTooSmall(index.depth())); - } else if index.depth() > MAX_DEPTH { - return Err(MerkleError::DepthTooBig(index.depth() as u64)); - } else { - // make sure that there are no leaf nodes in the ancestors of the index; since leaf - // nodes can live at specific depth, we just need to check these depths. - let tier = ((index.depth() - 1) / TIER_SIZE) as usize; - let mut tier_index = index; - for &depth in TIER_DEPTHS[..tier].iter().rev() { - tier_index.move_up_to(depth); - if self.upper_leaves.contains(&tier_index) { - return Err(MerkleError::NodeNotInSet(index)); - } - } - } - - Ok(()) - } - - /// Returns a node at the specified index. If the node does not exist at this index, a root - /// for an empty subtree at the index's depth is returned. - /// - /// Unlike [NodeStore::get_node()] this does not perform any checks to verify that the - /// returned node is valid in the context of this tree. - fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest { - match self.nodes.get(index) { - Some(node) => *node, - None => EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize], - } - } - - /// Removes a sequence of nodes starting at the specified index and traversing the tree up to - /// the specified depth. The node at the `end_depth` is also removed, and the appropriate leaf - /// flag is cleared. - /// - /// This method does not update any other nodes and does not recompute the tree root. - fn remove_branch(&mut self, index: LeafNodeIndex, end_depth: u8) { - if index.depth() == MAX_DEPTH { - self.bottom_leaves.remove(&index.value()); - } else { - self.upper_leaves.remove(&index); - } - - let mut index: NodeIndex = index.into(); - assert!(index.depth() > end_depth); - for _ in 0..(index.depth() - end_depth + 1) { - self.nodes.remove(&index); - index.move_up() - } - } -} - -// HELPER FUNCTIONS -// ================================================================================================ - -/// Returns true if the specified node is a root of an empty tree or an empty value ([ZERO; 4]). -fn is_empty_root(node: &RpoDigest) -> bool { - EmptySubtreeRoots::empty_hashes(MAX_DEPTH).contains(node) -} diff --git a/src/merkle/tiered_smt/proof.rs b/src/merkle/tiered_smt/proof.rs deleted file mode 100644 index 1de17d5..0000000 --- a/src/merkle/tiered_smt/proof.rs +++ /dev/null @@ -1,170 +0,0 @@ -use super::{ - get_common_prefix_tier_depth, get_key_prefix, hash_bottom_leaf, hash_upper_leaf, - EmptySubtreeRoots, LeafNodeIndex, MerklePath, RpoDigest, TieredSmtProofError, Vec, Word, -}; - -// CONSTANTS -// ================================================================================================ - -/// Maximum node depth. This is also the bottom tier of the tree. -const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH; - -/// Value of an empty leaf. -pub const EMPTY_VALUE: Word = super::TieredSmt::EMPTY_VALUE; - -/// Depths at which leaves can exist in a tiered SMT. -pub const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS; - -// TIERED SPARSE MERKLE TREE PROOF -// ================================================================================================ - -/// A proof which can be used to assert membership (or non-membership) of key-value pairs in a -/// Tiered Sparse Merkle tree. -/// -/// The proof consists of a Merkle path and one or more key-value entries which describe the node -/// located at the base of the path. If the node at the base of the path resolves to [ZERO; 4], -/// the entries will contain a single item with value set to [ZERO; 4]. -#[derive(PartialEq, Eq, Debug)] -pub struct TieredSmtProof { - path: MerklePath, - entries: Vec<(RpoDigest, Word)>, -} - -impl TieredSmtProof { - // CONSTRUCTOR - // -------------------------------------------------------------------------------------------- - /// Returns a new instance of [TieredSmtProof] instantiated from the specified path and entries. - /// - /// # Panics - /// Panics if: - /// - The length of the path is greater than 64. - /// - Entries is an empty vector. - /// - Entries contains more than 1 item, but the length of the path is not 64. - /// - Entries contains more than 1 item, and one of the items has value set to [ZERO; 4]. - /// - Entries contains multiple items with keys which don't share the same 64-bit prefix. - pub fn new(path: MerklePath, entries: I) -> Result - where - I: IntoIterator, - { - let entries: Vec<(RpoDigest, Word)> = entries.into_iter().collect(); - - if !TIER_DEPTHS.into_iter().any(|e| e == path.depth()) { - return Err(TieredSmtProofError::NotATierPath(path.depth())); - } - - if entries.is_empty() { - return Err(TieredSmtProofError::EntriesEmpty); - } - - if entries.len() > 1 { - if path.depth() != MAX_DEPTH { - return Err(TieredSmtProofError::MultipleEntriesOutsideLastTier); - } - - let prefix = get_key_prefix(&entries[0].0); - for entry in entries.iter().skip(1) { - if entry.1 == EMPTY_VALUE { - return Err(TieredSmtProofError::EmptyValueNotAllowed); - } - let current = get_key_prefix(&entry.0); - if prefix != current { - return Err(TieredSmtProofError::MismatchedPrefixes(prefix, current)); - } - } - } - - Ok(Self { path, entries }) - } - - // PROOF VERIFIER - // -------------------------------------------------------------------------------------------- - - /// Returns true if a Tiered Sparse Merkle tree with the specified root contains the provided - /// key-value pair. - /// - /// Note: this method cannot be used to assert non-membership. That is, if false is returned, - /// it does not mean that the provided key-value pair is not in the tree. - pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool { - // Handles the following scenarios: - // - the value is set - // - empty leaf, there is an explicit entry for the key with the empty value - // - shared 64-bit prefix, the target key is not included in the entries list, the value is implicitly the empty word - let v = match self.entries.iter().find(|(k, _)| k == key) { - Some((_, v)) => v, - None => &EMPTY_VALUE, - }; - - // The value must match for the proof to be valid - if v != value { - return false; - } - - // If the proof is for an empty value, we can verify it against any key which has a common - // prefix with the key storied in entries, but the prefix must be greater than the path - // length - if self.is_value_empty() - && get_common_prefix_tier_depth(key, &self.entries[0].0) < self.path.depth() - { - return false; - } - - // make sure the Merkle path resolves to the correct root - root == &self.compute_root() - } - - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- - - /// Returns the value associated with the specific key according to this proof, or None if - /// this proof does not contain a value for the specified key. - /// - /// A key-value pair generated by using this method should pass the `verify_membership()` check. - pub fn get(&self, key: &RpoDigest) -> Option { - if self.is_value_empty() { - let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0); - if common_prefix_tier < self.path.depth() { - None - } else { - Some(EMPTY_VALUE) - } - } else { - self.entries.iter().find(|(k, _)| k == key).map(|(_, value)| *value) - } - } - - /// Computes the root of a Tiered Sparse Merkle tree to which this proof resolve. - pub fn compute_root(&self) -> RpoDigest { - let node = self.build_node(); - let index = LeafNodeIndex::from_key(&self.entries[0].0, self.path.depth()); - self.path - .compute_root(index.value(), node) - .expect("failed to compute Merkle path root") - } - - /// Consume the proof and returns its parts. - pub fn into_parts(self) -> (MerklePath, Vec<(RpoDigest, Word)>) { - (self.path, self.entries) - } - - // HELPER METHODS - // -------------------------------------------------------------------------------------------- - - /// Returns true if the proof is for an empty value. - fn is_value_empty(&self) -> bool { - self.entries[0].1 == EMPTY_VALUE - } - - /// Converts the entries contained in this proof into a node value for node at the base of the - /// path contained in this proof. - fn build_node(&self) -> RpoDigest { - let depth = self.path.depth(); - if self.is_value_empty() { - EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[depth as usize] - } else if depth == MAX_DEPTH { - hash_bottom_leaf(&self.entries) - } else { - let (key, value) = self.entries[0]; - hash_upper_leaf(key, value, depth) - } - } -} diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs deleted file mode 100644 index 788ba28..0000000 --- a/src/merkle/tiered_smt/tests.rs +++ /dev/null @@ -1,968 +0,0 @@ -use super::{ - super::{super::ONE, super::WORD_SIZE, Felt, MerkleStore, EMPTY_WORD, ZERO}, - EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word, -}; - -// INSERTION TESTS -// ================================================================================================ - -#[test] -fn tsmt_insert_one() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); - let value = [ONE; WORD_SIZE]; - - // since the tree is empty, the first node will be inserted at depth 16 and the index will be - // 16 most significant bits of the key - let index = NodeIndex::make(16, raw >> 48); - let leaf_node = build_leaf_node(key, value, 16); - let tree_root = store.set_node(smt.root(), index, leaf_node).unwrap().root; - - smt.insert(key, value); - - assert_eq!(smt.root(), tree_root); - - // make sure the value was inserted, and the node is at the expected index - assert_eq!(smt.get_value(key), value); - assert_eq!(smt.get_node(index).unwrap(), leaf_node); - - // make sure the paths we get from the store and the tree match - let expected_path = store.get_path(tree_root, index).unwrap(); - assert_eq!(smt.get_path(index).unwrap(), expected_path.path); - - // make sure inner nodes match - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - assert_eq!(actual_nodes.len(), expected_nodes.len()); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); - - // make sure leaves are returned correctly - let mut leaves = smt.upper_leaves(); - assert_eq!(leaves.next(), Some((leaf_node, key, value))); - assert_eq!(leaves.next(), None); -} - -#[test] -fn tsmt_insert_two_16() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - // --- insert the first value --------------------------------------------- - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let val_a = [ONE; WORD_SIZE]; - smt.insert(key_a, val_a); - - // --- insert the second value -------------------------------------------- - // the key for this value has the same 16-bit prefix as the key for the first value, - // thus, on insertions, both values should be pushed to depth 32 tier - let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let val_b = [Felt::new(2); WORD_SIZE]; - smt.insert(key_b, val_b); - - // --- build Merkle store with equivalent data ---------------------------- - let mut tree_root = get_init_root(); - let index_a = NodeIndex::make(32, raw_a >> 32); - let leaf_node_a = build_leaf_node(key_a, val_a, 32); - tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; - - let index_b = NodeIndex::make(32, raw_b >> 32); - let leaf_node_b = build_leaf_node(key_b, val_b, 32); - tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; - - // --- verify that data is consistent between store and tree -------------- - - assert_eq!(smt.root(), tree_root); - - assert_eq!(smt.get_value(key_a), val_a); - assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; - assert_eq!(smt.get_path(index_a).unwrap(), expected_path); - - assert_eq!(smt.get_value(key_b), val_b); - assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; - assert_eq!(smt.get_path(index_b).unwrap(), expected_path); - - // make sure inner nodes match - the store contains more entries because it keeps track of - // all prior state - so, we don't check that the number of inner nodes is the same in both - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); - - // make sure leaves are returned correctly - let mut leaves = smt.upper_leaves(); - assert_eq!(leaves.next(), Some((leaf_node_a, key_a, val_a))); - assert_eq!(leaves.next(), Some((leaf_node_b, key_b, val_b))); - assert_eq!(leaves.next(), None); -} - -#[test] -fn tsmt_insert_two_32() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - // --- insert the first value --------------------------------------------- - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let val_a = [ONE; WORD_SIZE]; - smt.insert(key_a, val_a); - - // --- insert the second value -------------------------------------------- - // the key for this value has the same 32-bit prefix as the key for the first value, - // thus, on insertions, both values should be pushed to depth 48 tier - let raw_b = 0b_10101010_10101010_00011111_11111111_00010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let val_b = [Felt::new(2); WORD_SIZE]; - smt.insert(key_b, val_b); - - // --- build Merkle store with equivalent data ---------------------------- - let mut tree_root = get_init_root(); - let index_a = NodeIndex::make(48, raw_a >> 16); - let leaf_node_a = build_leaf_node(key_a, val_a, 48); - tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; - - let index_b = NodeIndex::make(48, raw_b >> 16); - let leaf_node_b = build_leaf_node(key_b, val_b, 48); - tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; - - // --- verify that data is consistent between store and tree -------------- - - assert_eq!(smt.root(), tree_root); - - assert_eq!(smt.get_value(key_a), val_a); - assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; - assert_eq!(smt.get_path(index_a).unwrap(), expected_path); - - assert_eq!(smt.get_value(key_b), val_b); - assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; - assert_eq!(smt.get_path(index_b).unwrap(), expected_path); - - // make sure inner nodes match - the store contains more entries because it keeps track of - // all prior state - so, we don't check that the number of inner nodes is the same in both - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); -} - -#[test] -fn tsmt_insert_three() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - // --- insert the first value --------------------------------------------- - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let val_a = [ONE; WORD_SIZE]; - smt.insert(key_a, val_a); - - // --- insert the second value -------------------------------------------- - // the key for this value has the same 16-bit prefix as the key for the first value, - // thus, on insertions, both values should be pushed to depth 32 tier - let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let val_b = [Felt::new(2); WORD_SIZE]; - smt.insert(key_b, val_b); - - // --- insert the third value --------------------------------------------- - // the key for this value has the same 16-bit prefix as the keys for the first two, - // values; thus, on insertions, it will be inserted into depth 32 tier, but will not - // affect locations of the other two values - let raw_c = 0b_10101010_10101010_11011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let val_c = [Felt::new(3); WORD_SIZE]; - smt.insert(key_c, val_c); - - // --- build Merkle store with equivalent data ---------------------------- - let mut tree_root = get_init_root(); - let index_a = NodeIndex::make(32, raw_a >> 32); - let leaf_node_a = build_leaf_node(key_a, val_a, 32); - tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; - - let index_b = NodeIndex::make(32, raw_b >> 32); - let leaf_node_b = build_leaf_node(key_b, val_b, 32); - tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; - - let index_c = NodeIndex::make(32, raw_c >> 32); - let leaf_node_c = build_leaf_node(key_c, val_c, 32); - tree_root = store.set_node(tree_root, index_c, leaf_node_c).unwrap().root; - - // --- verify that data is consistent between store and tree -------------- - - assert_eq!(smt.root(), tree_root); - - assert_eq!(smt.get_value(key_a), val_a); - assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; - assert_eq!(smt.get_path(index_a).unwrap(), expected_path); - - assert_eq!(smt.get_value(key_b), val_b); - assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; - assert_eq!(smt.get_path(index_b).unwrap(), expected_path); - - assert_eq!(smt.get_value(key_c), val_c); - assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c); - let expected_path = store.get_path(tree_root, index_c).unwrap().path; - assert_eq!(smt.get_path(index_c).unwrap(), expected_path); - - // make sure inner nodes match - the store contains more entries because it keeps track of - // all prior state - so, we don't check that the number of inner nodes is the same in both - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); -} - -// UPDATE TESTS -// ================================================================================================ - -#[test] -fn tsmt_update() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - // --- insert a value into the tree --------------------------------------- - let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); - let value_a = [ONE; WORD_SIZE]; - smt.insert(key, value_a); - - // --- update the value --------------------------------------------------- - let value_b = [Felt::new(2); WORD_SIZE]; - smt.insert(key, value_b); - - // --- verify consistency ------------------------------------------------- - let mut tree_root = get_init_root(); - let index = NodeIndex::make(16, raw >> 48); - let leaf_node = build_leaf_node(key, value_b, 16); - tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root; - - assert_eq!(smt.root(), tree_root); - - assert_eq!(smt.get_value(key), value_b); - assert_eq!(smt.get_node(index).unwrap(), leaf_node); - let expected_path = store.get_path(tree_root, index).unwrap().path; - assert_eq!(smt.get_path(index).unwrap(), expected_path); - - // make sure inner nodes match - the store contains more entries because it keeps track of - // all prior state - so, we don't check that the number of inner nodes is the same in both - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); -} - -// DELETION TESTS -// ================================================================================================ - -#[test] -fn tsmt_delete_16() { - let mut smt = TieredSmt::default(); - - // --- insert a value into the tree --------------------------------------- - let smt0 = smt.clone(); - let raw_a = 0b_01010101_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // --- insert another value into the tree --------------------------------- - let smt1 = smt.clone(); - let raw_b = 0b_01011111_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - // --- delete the last inserted value ------------------------------------- - assert_eq!(smt.insert(key_b, EMPTY_WORD), value_b); - assert_eq!(smt, smt1); - - // --- delete the first inserted value ------------------------------------ - assert_eq!(smt.insert(key_a, EMPTY_WORD), value_a); - assert_eq!(smt, smt0); -} - -#[test] -fn tsmt_delete_32() { - let mut smt = TieredSmt::default(); - - // --- insert a value into the tree --------------------------------------- - let smt0 = smt.clone(); - let raw_a = 0b_01010101_01101100_01111111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // --- insert another with the same 16-bit prefix into the tree ----------- - let smt1 = smt.clone(); - let raw_b = 0b_01010101_01101100_00111111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - // --- insert the 3rd value with the same 16-bit prefix into the tree ----- - let smt2 = smt.clone(); - let raw_c = 0b_01010101_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - smt.insert(key_c, value_c); - - // --- delete the last inserted value ------------------------------------- - assert_eq!(smt.insert(key_c, EMPTY_WORD), value_c); - assert_eq!(smt, smt2); - - // --- delete the last inserted value ------------------------------------- - assert_eq!(smt.insert(key_b, EMPTY_WORD), value_b); - assert_eq!(smt, smt1); - - // --- delete the first inserted value ------------------------------------ - assert_eq!(smt.insert(key_a, EMPTY_WORD), value_a); - assert_eq!(smt, smt0); -} - -#[test] -fn tsmt_delete_48_same_32_bit_prefix() { - let mut smt = TieredSmt::default(); - - // test the case when all values share the same 32-bit prefix - - // --- insert a value into the tree --------------------------------------- - let smt0 = smt.clone(); - let raw_a = 0b_01010101_01010101_11111111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // --- insert another with the same 32-bit prefix into the tree ----------- - let smt1 = smt.clone(); - let raw_b = 0b_01010101_01010101_11111111_11111111_11010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - // --- insert the 3rd value with the same 32-bit prefix into the tree ----- - let smt2 = smt.clone(); - let raw_c = 0b_01010101_01010101_11111111_11111111_11110110_10010011_11100000_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - smt.insert(key_c, value_c); - - // --- delete the last inserted value ------------------------------------- - assert_eq!(smt.insert(key_c, EMPTY_WORD), value_c); - assert_eq!(smt, smt2); - - // --- delete the last inserted value ------------------------------------- - assert_eq!(smt.insert(key_b, EMPTY_WORD), value_b); - assert_eq!(smt, smt1); - - // --- delete the first inserted value ------------------------------------ - assert_eq!(smt.insert(key_a, EMPTY_WORD), value_a); - assert_eq!(smt, smt0); -} - -#[test] -fn tsmt_delete_48_mixed_prefix() { - let mut smt = TieredSmt::default(); - - // test the case when some values share a 32-bit prefix and others share a 16-bit prefix - - // --- insert a value into the tree --------------------------------------- - let smt0 = smt.clone(); - let raw_a = 0b_01010101_01010101_11111111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // --- insert another with the same 16-bit prefix into the tree ----------- - let smt1 = smt.clone(); - let raw_b = 0b_01010101_01010101_01111111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - // --- insert a value with the same 32-bit prefix as the first value ----- - let smt2 = smt.clone(); - let raw_c = 0b_01010101_01010101_11111111_11111111_11010110_10010011_11100000_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - smt.insert(key_c, value_c); - - // --- insert another value with the same 32-bit prefix as the first value - let smt3 = smt.clone(); - let raw_d = 0b_01010101_01010101_11111111_11111111_11110110_10010011_11100000_00000000_u64; - let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_d)]); - let value_d = [ONE, ZERO, ZERO, ZERO]; - smt.insert(key_d, value_d); - - // --- delete the inserted values one-by-one ------------------------------ - assert_eq!(smt.insert(key_d, EMPTY_WORD), value_d); - assert_eq!(smt, smt3); - - assert_eq!(smt.insert(key_c, EMPTY_WORD), value_c); - assert_eq!(smt, smt2); - - assert_eq!(smt.insert(key_b, EMPTY_WORD), value_b); - assert_eq!(smt, smt1); - - assert_eq!(smt.insert(key_a, EMPTY_WORD), value_a); - assert_eq!(smt, smt0); -} - -#[test] -fn tsmt_delete_64() { - let mut smt = TieredSmt::default(); - - // test the case when all values share the same 48-bit prefix - - // --- insert a value into the tree --------------------------------------- - let smt0 = smt.clone(); - let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // --- insert a value with the same 48-bit prefix into the tree ----------- - let smt1 = smt.clone(); - let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - // --- insert a value with the same 32-bit prefix into the tree ----------- - let smt2 = smt.clone(); - let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - smt.insert(key_c, value_c); - - let smt3 = smt.clone(); - let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64; - let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]); - let value_d = [ONE, ZERO, ZERO, ZERO]; - smt.insert(key_d, value_d); - - // --- delete the last inserted value ------------------------------------- - assert_eq!(smt.insert(key_d, EMPTY_WORD), value_d); - assert_eq!(smt, smt3); - - assert_eq!(smt.insert(key_c, EMPTY_WORD), value_c); - assert_eq!(smt, smt2); - - assert_eq!(smt.insert(key_b, EMPTY_WORD), value_b); - assert_eq!(smt, smt1); - - assert_eq!(smt.insert(key_a, EMPTY_WORD), value_a); - assert_eq!(smt, smt0); -} - -#[test] -fn tsmt_delete_64_leaf_promotion() { - let mut smt = TieredSmt::default(); - - // --- delete from bottom tier (no promotion to upper tiers) -------------- - - // insert a value into the tree - let raw_a = 0b_01010101_01010101_11111111_11111111_10101010_10101010_11111111_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // insert another value with a key having the same 64-bit prefix - let key_b = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - // insert a value with a key which shared the same 48-bit prefix - let raw_c = 0b_01010101_01010101_11111111_11111111_10101010_10101010_00111111_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - smt.insert(key_c, value_c); - - // delete entry A and compare to the tree which was built from B and C - smt.insert(key_a, EMPTY_WORD); - - let mut expected_smt = TieredSmt::default(); - expected_smt.insert(key_b, value_b); - expected_smt.insert(key_c, value_c); - assert_eq!(smt, expected_smt); - - // entries B and C should stay at depth 64 - assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 64); - assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 64); - - // --- delete from bottom tier (promotion to depth 48) -------------------- - - let mut smt = TieredSmt::default(); - smt.insert(key_a, value_a); - smt.insert(key_b, value_b); - - // insert a value with a key which shared the same 32-bit prefix - let raw_c = 0b_01010101_01010101_11111111_11111111_11101010_10101010_11111111_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - smt.insert(key_c, value_c); - - // delete entry A and compare to the tree which was built from B and C - smt.insert(key_a, EMPTY_WORD); - - let mut expected_smt = TieredSmt::default(); - expected_smt.insert(key_b, value_b); - expected_smt.insert(key_c, value_c); - assert_eq!(smt, expected_smt); - - // entry B moves to depth 48, entry C stays at depth 48 - assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 48); - assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 48); - - // --- delete from bottom tier (promotion to depth 32) -------------------- - - let mut smt = TieredSmt::default(); - smt.insert(key_a, value_a); - smt.insert(key_b, value_b); - - // insert a value with a key which shared the same 16-bit prefix - let raw_c = 0b_01010101_01010101_01111111_11111111_10101010_10101010_11111111_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - smt.insert(key_c, value_c); - - // delete entry A and compare to the tree which was built from B and C - smt.insert(key_a, EMPTY_WORD); - - let mut expected_smt = TieredSmt::default(); - expected_smt.insert(key_b, value_b); - expected_smt.insert(key_c, value_c); - assert_eq!(smt, expected_smt); - - // entry B moves to depth 32, entry C stays at depth 32 - assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 32); - assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 32); - - // --- delete from bottom tier (promotion to depth 16) -------------------- - - let mut smt = TieredSmt::default(); - smt.insert(key_a, value_a); - smt.insert(key_b, value_b); - - // insert a value with a key which shared prefix < 16 bits - let raw_c = 0b_01010101_01010100_11111111_11111111_10101010_10101010_11111111_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - smt.insert(key_c, value_c); - - // delete entry A and compare to the tree which was built from B and C - smt.insert(key_a, EMPTY_WORD); - - let mut expected_smt = TieredSmt::default(); - expected_smt.insert(key_b, value_b); - expected_smt.insert(key_c, value_c); - assert_eq!(smt, expected_smt); - - // entry B moves to depth 16, entry C stays at depth 16 - assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 16); - assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 16); -} - -#[test] -fn test_order_sensitivity() { - let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000001_u64; - let value = [ONE; WORD_SIZE]; - - let key_1 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); - let key_2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]); - - let mut smt_1 = TieredSmt::default(); - - smt_1.insert(key_1, value); - smt_1.insert(key_2, value); - smt_1.insert(key_2, EMPTY_WORD); - - let mut smt_2 = TieredSmt::default(); - smt_2.insert(key_1, value); - - assert_eq!(smt_1.root(), smt_2.root()); -} - -// BOTTOM TIER TESTS -// ================================================================================================ - -#[test] -fn tsmt_bottom_tier() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - // common prefix for the keys - let prefix = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - - // --- insert the first value --------------------------------------------- - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(prefix)]); - let val_a = [ONE; WORD_SIZE]; - smt.insert(key_a, val_a); - - // --- insert the second value -------------------------------------------- - // this key has the same 64-bit prefix and thus both values should end up in the same - // node at depth 64 - let key_b = RpoDigest::from([ZERO, ONE, ONE, Felt::new(prefix)]); - let val_b = [Felt::new(2); WORD_SIZE]; - smt.insert(key_b, val_b); - - // --- build Merkle store with equivalent data ---------------------------- - let index = NodeIndex::make(64, prefix); - // to build bottom leaf we sort by key starting with the least significant element, thus - // key_b is smaller than key_a. - let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]); - let mut tree_root = get_init_root(); - tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root; - - // --- verify that data is consistent between store and tree -------------- - - assert_eq!(smt.root(), tree_root); - - assert_eq!(smt.get_value(key_a), val_a); - assert_eq!(smt.get_value(key_b), val_b); - - assert_eq!(smt.get_node(index).unwrap(), leaf_node); - let expected_path = store.get_path(tree_root, index).unwrap().path; - assert_eq!(smt.get_path(index).unwrap(), expected_path); - - // make sure inner nodes match - the store contains more entries because it keeps track of - // all prior state - so, we don't check that the number of inner nodes is the same in both - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); - - // make sure leaves are returned correctly - let smt_clone = smt.clone(); - let mut leaves = smt_clone.bottom_leaves(); - assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)]))); - assert_eq!(leaves.next(), None); - - // --- update a leaf at the bottom tier ------------------------------------------------------- - - let val_a2 = [Felt::new(3); WORD_SIZE]; - assert_eq!(smt.insert(key_a, val_a2), val_a); - - let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a2]); - store.set_node(tree_root, index, leaf_node).unwrap(); - - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); - - let mut leaves = smt.bottom_leaves(); - assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a2)]))); - assert_eq!(leaves.next(), None); -} - -#[test] -fn tsmt_bottom_tier_two() { - let mut smt = TieredSmt::default(); - let mut store = MerkleStore::default(); - - // --- insert the first value --------------------------------------------- - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let val_a = [ONE; WORD_SIZE]; - smt.insert(key_a, val_a); - - // --- insert the second value -------------------------------------------- - // the key for this value has the same 48-bit prefix as the key for the first value, - // thus, on insertions, both should end up in different nodes at depth 64 - let raw_b = 0b_10101010_10101010_00011111_11111111_10010110_10010011_01100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let val_b = [Felt::new(2); WORD_SIZE]; - smt.insert(key_b, val_b); - - // --- build Merkle store with equivalent data ---------------------------- - let mut tree_root = get_init_root(); - let index_a = NodeIndex::make(64, raw_a); - let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]); - tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root; - - let index_b = NodeIndex::make(64, raw_b); - let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]); - tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root; - - // --- verify that data is consistent between store and tree -------------- - - assert_eq!(smt.root(), tree_root); - - assert_eq!(smt.get_value(key_a), val_a); - assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a); - let expected_path = store.get_path(tree_root, index_a).unwrap().path; - assert_eq!(smt.get_path(index_a).unwrap(), expected_path); - - assert_eq!(smt.get_value(key_b), val_b); - assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); - let expected_path = store.get_path(tree_root, index_b).unwrap().path; - assert_eq!(smt.get_path(index_b).unwrap(), expected_path); - - // make sure inner nodes match - the store contains more entries because it keeps track of - // all prior state - so, we don't check that the number of inner nodes is the same in both - let expected_nodes = get_non_empty_nodes(&store); - let actual_nodes = smt.inner_nodes().collect::>(); - actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); - - // make sure leaves are returned correctly - let mut leaves = smt.bottom_leaves(); - assert_eq!(leaves.next(), Some((leaf_node_b, vec![(key_b, val_b)]))); - assert_eq!(leaves.next(), Some((leaf_node_a, vec![(key_a, val_a)]))); - assert_eq!(leaves.next(), None); -} - -// GET PROOF TESTS -// ================================================================================================ - -/// Tests the membership and non-membership proof for a single at depth 64 -#[test] -fn tsmt_get_proof_single_element_64() { - let mut smt = TieredSmt::default(); - - let raw_a = 0b_00000000_00000001_00000000_00000001_00000000_00000001_00000000_00000001_u64; - let key_a = [ONE, ONE, ONE, raw_a.into()].into(); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // push element `a` to depth 64, by inserting another value that shares the 48-bit prefix - let raw_b = 0b_00000000_00000001_00000000_00000001_00000000_00000001_00000000_00000000_u64; - let key_b = [ONE, ONE, ONE, raw_b.into()].into(); - smt.insert(key_b, [ONE, ONE, ONE, ONE]); - - // verify the proof for element `a` - let proof = smt.prove(key_a); - assert!(proof.verify_membership(&key_a, &value_a, &smt.root())); - - // check that a value that is not inserted in the tree produces a valid membership proof for the - // empty word - let key = [ZERO, ZERO, ZERO, ZERO].into(); - let proof = smt.prove(key); - assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root())); - - // check that a key that shared the 64-bit prefix with `a`, but is not inserted, also has a - // valid membership proof for the empty word - let key = [ONE, ONE, ZERO, raw_a.into()].into(); - let proof = smt.prove(key); - assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root())); -} - -#[test] -fn tsmt_get_proof() { - let mut smt = TieredSmt::default(); - - // --- insert a value into the tree --------------------------------------- - let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64; - let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE, ONE, ONE, ONE]; - smt.insert(key_a, value_a); - - // --- insert a value with the same 48-bit prefix into the tree ----------- - let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ONE, ONE, ZERO]; - smt.insert(key_b, value_b); - - let smt_alt = smt.clone(); - - // --- insert a value with the same 32-bit prefix into the tree ----------- - let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_00000000_u64; - let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - smt.insert(key_c, value_c); - - // --- insert a value with the same 64-bit prefix as A into the tree ------ - let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64; - let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]); - let value_d = [ONE, ZERO, ZERO, ZERO]; - smt.insert(key_d, value_d); - - // at this point the tree looks as follows: - // - A and D are located in the same node at depth 64. - // - B is located at depth 64 and shares the same 48-bit prefix with A and D. - // - C is located at depth 48 and shares the same 32-bit prefix with A, B, and D. - - // --- generate proof for key A and test that it verifies correctly ------- - let proof = smt.prove(key_a); - assert!(proof.verify_membership(&key_a, &value_a, &smt.root())); - - assert!(!proof.verify_membership(&key_a, &value_b, &smt.root())); - assert!(!proof.verify_membership(&key_a, &EMPTY_WORD, &smt.root())); - assert!(!proof.verify_membership(&key_b, &value_a, &smt.root())); - assert!(!proof.verify_membership(&key_a, &value_a, &smt_alt.root())); - - assert_eq!(proof.get(&key_a), Some(value_a)); - assert_eq!(proof.get(&key_b), None); - - // since A and D are stored in the same node, we should be able to use the proof to verify - // membership of D - assert!(proof.verify_membership(&key_d, &value_d, &smt.root())); - assert_eq!(proof.get(&key_d), Some(value_d)); - - // --- generate proof for key B and test that it verifies correctly ------- - let proof = smt.prove(key_b); - assert!(proof.verify_membership(&key_b, &value_b, &smt.root())); - - assert!(!proof.verify_membership(&key_b, &value_a, &smt.root())); - assert!(!proof.verify_membership(&key_b, &EMPTY_WORD, &smt.root())); - assert!(!proof.verify_membership(&key_a, &value_b, &smt.root())); - assert!(!proof.verify_membership(&key_b, &value_b, &smt_alt.root())); - - assert_eq!(proof.get(&key_b), Some(value_b)); - assert_eq!(proof.get(&key_a), None); - - // --- generate proof for key C and test that it verifies correctly ------- - let proof = smt.prove(key_c); - assert!(proof.verify_membership(&key_c, &value_c, &smt.root())); - - assert!(!proof.verify_membership(&key_c, &value_a, &smt.root())); - assert!(!proof.verify_membership(&key_c, &EMPTY_WORD, &smt.root())); - assert!(!proof.verify_membership(&key_a, &value_c, &smt.root())); - assert!(!proof.verify_membership(&key_c, &value_c, &smt_alt.root())); - - assert_eq!(proof.get(&key_c), Some(value_c)); - assert_eq!(proof.get(&key_b), None); - - // --- generate proof for key D and test that it verifies correctly ------- - let proof = smt.prove(key_d); - assert!(proof.verify_membership(&key_d, &value_d, &smt.root())); - - assert!(!proof.verify_membership(&key_d, &value_b, &smt.root())); - assert!(!proof.verify_membership(&key_d, &EMPTY_WORD, &smt.root())); - assert!(!proof.verify_membership(&key_b, &value_d, &smt.root())); - assert!(!proof.verify_membership(&key_d, &value_d, &smt_alt.root())); - - assert_eq!(proof.get(&key_d), Some(value_d)); - assert_eq!(proof.get(&key_b), None); - - // since A and D are stored in the same node, we should be able to use the proof to verify - // membership of A - assert!(proof.verify_membership(&key_a, &value_a, &smt.root())); - assert_eq!(proof.get(&key_a), Some(value_a)); - - // --- generate proof for an empty key at depth 64 ------------------------ - // this key has the same 48-bit prefix as A but is different from B - let raw = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000011_u64; - let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); - - let proof = smt.prove(key); - assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root())); - - assert!(!proof.verify_membership(&key, &value_a, &smt.root())); - assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root())); - - assert_eq!(proof.get(&key), Some(EMPTY_WORD)); - assert_eq!(proof.get(&key_b), None); - - // the same proof should verify against any key with the same 64-bit prefix - let key2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]); - assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root())); - assert_eq!(proof.get(&key2), Some(EMPTY_WORD)); - - // but verifying if against a key with the same 63-bit prefix (or smaller) should fail - let raw3 = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000010_u64; - let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]); - assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root())); - assert_eq!(proof.get(&key3), None); - - // --- generate proof for an empty key at depth 48 ------------------------ - // this key has the same 32-prefix as A, B, C, and D, but is different from C - let raw = 0b_01010101_01010101_11111111_11111111_00110101_10101010_11111100_00000000_u64; - let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); - - let proof = smt.prove(key); - assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root())); - - assert!(!proof.verify_membership(&key, &value_a, &smt.root())); - assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root())); - - assert_eq!(proof.get(&key), Some(EMPTY_WORD)); - assert_eq!(proof.get(&key_b), None); - - // the same proof should verify against any key with the same 48-bit prefix - let raw2 = 0b_01010101_01010101_11111111_11111111_00110101_10101010_01111100_00000000_u64; - let key2 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw2)]); - assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root())); - assert_eq!(proof.get(&key2), Some(EMPTY_WORD)); - - // but verifying against a key with the same 47-bit prefix (or smaller) should fail - let raw3 = 0b_01010101_01010101_11111111_11111111_00110101_10101011_11111100_00000000_u64; - let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]); - assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root())); - assert_eq!(proof.get(&key3), None); -} - -// ERROR TESTS -// ================================================================================================ - -#[test] -fn tsmt_node_not_available() { - let mut smt = TieredSmt::default(); - - let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); - let value = [ONE; WORD_SIZE]; - - // build an index which is just below the inserted leaf node - let index = NodeIndex::make(17, raw >> 47); - - // since we haven't inserted the node yet, we should be able to get node and path to this index - assert!(smt.get_node(index).is_ok()); - assert!(smt.get_path(index).is_ok()); - - smt.insert(key, value); - - // but once the node is inserted, everything under it should be unavailable - assert!(smt.get_node(index).is_err()); - assert!(smt.get_path(index).is_err()); - - let index = NodeIndex::make(32, raw >> 32); - assert!(smt.get_node(index).is_err()); - assert!(smt.get_path(index).is_err()); - - let index = NodeIndex::make(34, raw >> 30); - assert!(smt.get_node(index).is_err()); - assert!(smt.get_path(index).is_err()); - - let index = NodeIndex::make(50, raw >> 14); - assert!(smt.get_node(index).is_err()); - assert!(smt.get_path(index).is_err()); - - let index = NodeIndex::make(64, raw); - assert!(smt.get_node(index).is_err()); - assert!(smt.get_path(index).is_err()); -} - -// HELPER FUNCTIONS -// ================================================================================================ - -fn get_init_root() -> RpoDigest { - EmptySubtreeRoots::empty_hashes(64)[0] -} - -fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest { - Rpo256::merge_in_domain(&[key, value.into()], depth.into()) -} - -fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { - assert_eq!(keys.len(), values.len()); - - let mut elements = Vec::with_capacity(keys.len()); - for (key, val) in keys.iter().zip(values.iter()) { - elements.extend_from_slice(key.as_elements()); - elements.extend_from_slice(val.as_slice()); - } - - Rpo256::hash_elements(&elements) -} - -fn get_non_empty_nodes(store: &MerkleStore) -> Vec { - store - .inner_nodes() - .filter(|node| !is_empty_subtree(&node.value)) - .collect::>() -} - -fn is_empty_subtree(node: &RpoDigest) -> bool { - EmptySubtreeRoots::empty_hashes(255).contains(node) -} diff --git a/src/merkle/tiered_smt/values.rs b/src/merkle/tiered_smt/values.rs deleted file mode 100644 index 4fac76a..0000000 --- a/src/merkle/tiered_smt/values.rs +++ /dev/null @@ -1,584 +0,0 @@ -use super::{get_key_prefix, BTreeMap, LeafNodeIndex, RpoDigest, StarkField, Vec, Word}; -use crate::utils::vec; -use core::{ - cmp::{Ord, Ordering}, - ops::RangeBounds, -}; -use winter_utils::collections::btree_map::Entry; - -// CONSTANTS -// ================================================================================================ - -/// Depths at which leaves can exist in a tiered SMT. -const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS; - -/// Maximum node depth. This is also the bottom tier of the tree. -const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH; - -// VALUE STORE -// ================================================================================================ -/// A store for key-value pairs for a Tiered Sparse Merkle tree. -/// -/// The store is organized in a [BTreeMap] where keys are 64 most significant bits of a key, and -/// the values are the corresponding key-value pairs (or a list of key-value pairs if more that -/// a single key-value pair shares the same 64-bit prefix). -/// -/// The store supports lookup by the full key (i.e. [RpoDigest]) as well as by the 64-bit key -/// prefix. -#[derive(Debug, Default, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub struct ValueStore { - values: BTreeMap, -} - -impl ValueStore { - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- - - /// Returns a reference to the value stored under the specified key, or None if there is no - /// value associated with the specified key. - pub fn get(&self, key: &RpoDigest) -> Option<&Word> { - let prefix = get_key_prefix(key); - self.values.get(&prefix).and_then(|entry| entry.get(key)) - } - - /// Returns the first key-value pair such that the key prefix is greater than or equal to the - /// specified prefix. - pub fn get_first(&self, prefix: u64) -> Option<&(RpoDigest, Word)> { - self.range(prefix..).next() - } - - /// Returns the first key-value pair such that the key prefix is greater than or equal to the - /// specified prefix and the key value is not equal to the exclude_key value. - pub fn get_first_filtered( - &self, - prefix: u64, - exclude_key: &RpoDigest, - ) -> Option<&(RpoDigest, Word)> { - self.range(prefix..).find(|(key, _)| key != exclude_key) - } - - /// Returns a vector with key-value pairs for all keys with the specified 64-bit prefix, or - /// None if no keys with the specified prefix are present in this store. - pub fn get_all(&self, prefix: u64) -> Option> { - self.values.get(&prefix).map(|entry| match entry { - StoreEntry::Single(kv_pair) => vec![*kv_pair], - StoreEntry::List(kv_pairs) => kv_pairs.clone(), - }) - } - - /// Returns information about a sibling of a leaf node with the specified index, but only if - /// this is the only sibling the leaf has in some subtree starting at the first tier. - /// - /// For example, if `index` is an index at depth 32, and there is a leaf node at depth 32 with - /// the same root at depth 16 as `index`, we say that this leaf is a lone sibling. - /// - /// The returned tuple contains: they key-value pair of the sibling as well as the index of - /// the node for the root of the common subtree in which both nodes are leaves. - /// - /// This method assumes that the key-value pair for the specified index has already been - /// removed from the store. - pub fn get_lone_sibling( - &self, - index: LeafNodeIndex, - ) -> Option<(&RpoDigest, &Word, LeafNodeIndex)> { - // iterate over tiers from top to bottom, looking at the tiers which are strictly above - // the depth of the index. This implies that only tiers at depth 32 and 48 will be - // considered. For each tier, check if the parent of the index at the higher tier - // contains a single node. The fist tier (depth 16) is excluded because we cannot move - // nodes at depth 16 to a higher tier. This implies that nodes at the first tier will - // never have "lone siblings". - for &tier_depth in TIER_DEPTHS.iter().filter(|&t| index.depth() > *t) { - // compute the index of the root at a higher tier - let mut parent_index = index; - parent_index.move_up_to(tier_depth); - - // find the lone sibling, if any; we need to handle the "last node" at a given tier - // separately specify the bounds for the search correctly. - let start_prefix = parent_index.value() << (MAX_DEPTH - tier_depth); - let sibling = if start_prefix.leading_ones() as u8 == tier_depth { - let mut iter = self.range(start_prefix..); - iter.next().filter(|_| iter.next().is_none()) - } else { - let end_prefix = (parent_index.value() + 1) << (MAX_DEPTH - tier_depth); - let mut iter = self.range(start_prefix..end_prefix); - iter.next().filter(|_| iter.next().is_none()) - }; - - if let Some((key, value)) = sibling { - return Some((key, value, parent_index)); - } - } - - None - } - - /// Returns an iterator over all key-value pairs in this store. - pub fn iter(&self) -> impl Iterator { - self.values.iter().flat_map(|(_, entry)| entry.iter()) - } - - // STATE MUTATORS - // -------------------------------------------------------------------------------------------- - - /// Inserts the specified key-value pair into this store and returns the value previously - /// associated with the specified key. - /// - /// If no value was previously associated with the specified key, None is returned. - pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option { - let prefix = get_key_prefix(&key); - match self.values.entry(prefix) { - Entry::Occupied(mut entry) => entry.get_mut().insert(key, value), - Entry::Vacant(entry) => { - entry.insert(StoreEntry::new(key, value)); - None - } - } - } - - /// Removes the key-value pair for the specified key from this store and returns the value - /// associated with this key. - /// - /// If no value was associated with the specified key, None is returned. - pub fn remove(&mut self, key: &RpoDigest) -> Option { - let prefix = get_key_prefix(key); - match self.values.entry(prefix) { - Entry::Occupied(mut entry) => { - let (value, remove_entry) = entry.get_mut().remove(key); - if remove_entry { - entry.remove_entry(); - } - value - } - Entry::Vacant(_) => None, - } - } - - // HELPER METHODS - // -------------------------------------------------------------------------------------------- - - /// Returns an iterator over all key-value pairs contained in this store such that the most - /// significant 64 bits of the key lay within the specified bounds. - /// - /// The order of iteration is from the smallest to the largest key. - fn range>(&self, bounds: R) -> impl Iterator { - self.values.range(bounds).flat_map(|(_, entry)| entry.iter()) - } -} - -// VALUE NODE -// ================================================================================================ - -/// An entry in the [ValueStore]. -/// -/// An entry can contain either a single key-value pair or a vector of key-value pairs sorted by -/// key. -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub enum StoreEntry { - Single((RpoDigest, Word)), - List(Vec<(RpoDigest, Word)>), -} - -impl StoreEntry { - // CONSTRUCTOR - // -------------------------------------------------------------------------------------------- - /// Returns a new [StoreEntry] instantiated with a single key-value pair. - pub fn new(key: RpoDigest, value: Word) -> Self { - Self::Single((key, value)) - } - - // PUBLIC ACCESSORS - // -------------------------------------------------------------------------------------------- - - /// Returns the value associated with the specified key, or None if this entry does not contain - /// a value associated with the specified key. - pub fn get(&self, key: &RpoDigest) -> Option<&Word> { - match self { - StoreEntry::Single(kv_pair) => { - if kv_pair.0 == *key { - Some(&kv_pair.1) - } else { - None - } - } - StoreEntry::List(kv_pairs) => { - match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) { - Ok(pos) => Some(&kv_pairs[pos].1), - Err(_) => None, - } - } - } - } - - /// Returns an iterator over all key-value pairs in this entry. - pub fn iter(&self) -> impl Iterator { - EntryIterator { entry: self, pos: 0 } - } - - // STATE MUTATORS - // -------------------------------------------------------------------------------------------- - - /// Inserts the specified key-value pair into this entry and returns the value previously - /// associated with the specified key, or None if no value was associated with the specified - /// key. - /// - /// If a new key is inserted, this will also transform a `SingleEntry` into a `ListEntry`. - pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option { - match self { - StoreEntry::Single(kv_pair) => { - // if the key is already in this entry, update the value and return - if kv_pair.0 == key { - let old_value = kv_pair.1; - kv_pair.1 = value; - return Some(old_value); - } - - // 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(|a, b| cmp_digests(&a.0, &b.0)); - - *self = StoreEntry::List(pairs); - None - } - StoreEntry::List(pairs) => { - match pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, &key)) { - Ok(pos) => { - let old_value = pairs[pos].1; - pairs[pos].1 = value; - Some(old_value) - } - Err(pos) => { - pairs.insert(pos, (key, value)); - None - } - } - } - } - } - - /// Removes the key-value pair with the specified key from this entry, and returns the value - /// of the removed pair. If the entry did not contain a key-value pair for the specified key, - /// None is returned. - /// - /// If the last last key-value pair was removed from the entry, the second tuple value will - /// be set to true. - pub fn remove(&mut self, key: &RpoDigest) -> (Option, bool) { - match self { - StoreEntry::Single(kv_pair) => { - if kv_pair.0 == *key { - (Some(kv_pair.1), true) - } else { - (None, false) - } - } - StoreEntry::List(kv_pairs) => { - match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) { - Ok(pos) => { - let kv_pair = kv_pairs.remove(pos); - if kv_pairs.len() == 1 { - *self = StoreEntry::Single(kv_pairs[0]); - } - (Some(kv_pair.1), false) - } - Err(_) => (None, false), - } - } - } - } -} - -/// A custom iterator over key-value pairs of a [StoreEntry]. -/// -/// For a `SingleEntry` this returns only one value, but for `ListEntry`, this iterates over the -/// entire list of key-value pairs. -pub struct EntryIterator<'a> { - entry: &'a StoreEntry, - pos: usize, -} - -impl<'a> Iterator for EntryIterator<'a> { - type Item = &'a (RpoDigest, Word); - - fn next(&mut self) -> Option { - match self.entry { - StoreEntry::Single(kv_pair) => { - if self.pos == 0 { - self.pos = 1; - Some(kv_pair) - } else { - None - } - } - StoreEntry::List(kv_pairs) => { - if self.pos >= kv_pairs.len() { - None - } else { - let kv_pair = &kv_pairs[self.pos]; - self.pos += 1; - Some(kv_pair) - } - } - } - } -} - -// HELPER FUNCTIONS -// ================================================================================================ - -/// Compares two digests element-by-element using their integer representations starting with the -/// most significant element. -fn cmp_digests(d1: &RpoDigest, d2: &RpoDigest) -> Ordering { - let d1 = Word::from(d1); - let d2 = Word::from(d2); - - for (v1, v2) in d1.iter().zip(d2.iter()).rev() { - let v1 = v1.as_int(); - let v2 = v2.as_int(); - if v1 != v2 { - return v1.cmp(&v2); - } - } - - Ordering::Equal -} - -// TESTS -// ================================================================================================ - -#[cfg(test)] -mod tests { - use super::{LeafNodeIndex, RpoDigest, StoreEntry, ValueStore}; - use crate::{Felt, ONE, WORD_SIZE, ZERO}; - - #[test] - fn test_insert() { - let mut store = ValueStore::default(); - - // insert the first key-value pair into the store - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE; WORD_SIZE]; - - assert!(store.insert(key_a, value_a).is_none()); - assert_eq!(store.values.len(), 1); - - let entry = store.values.get(&raw_a).unwrap(); - let expected_entry = StoreEntry::Single((key_a, value_a)); - assert_eq!(entry, &expected_entry); - - // insert a key-value pair with a different key into the store; since the keys are - // different, another entry is added to the values map - let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ZERO, ONE, ZERO]; - - assert!(store.insert(key_b, value_b).is_none()); - assert_eq!(store.values.len(), 2); - - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = StoreEntry::Single((key_a, value_a)); - assert_eq!(entry1, &expected_entry1); - - let entry2 = store.values.get(&raw_b).unwrap(); - let expected_entry2 = StoreEntry::Single((key_b, value_b)); - assert_eq!(entry2, &expected_entry2); - - // insert a key-value pair with the same 64-bit key prefix as the first key; this should - // transform the first entry into a List entry - let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - - assert!(store.insert(key_c, value_c).is_none()); - assert_eq!(store.values.len(), 2); - - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a)]); - assert_eq!(entry1, &expected_entry1); - - let entry2 = store.values.get(&raw_b).unwrap(); - let expected_entry2 = StoreEntry::Single((key_b, value_b)); - assert_eq!(entry2, &expected_entry2); - - // replace values for keys a and b - let value_a2 = [ONE, ONE, ONE, ZERO]; - let value_b2 = [ZERO, ZERO, ZERO, ONE]; - - assert_eq!(store.insert(key_a, value_a2), Some(value_a)); - assert_eq!(store.values.len(), 2); - - assert_eq!(store.insert(key_b, value_b2), Some(value_b)); - assert_eq!(store.values.len(), 2); - - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2)]); - assert_eq!(entry1, &expected_entry1); - - let entry2 = store.values.get(&raw_b).unwrap(); - let expected_entry2 = StoreEntry::Single((key_b, value_b2)); - assert_eq!(entry2, &expected_entry2); - - // insert one more key-value pair with the same 64-bit key-prefix as the first key - let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_d = [ZERO, ONE, ZERO, ZERO]; - - assert!(store.insert(key_d, value_d).is_none()); - assert_eq!(store.values.len(), 2); - - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = - StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2), (key_d, value_d)]); - assert_eq!(entry1, &expected_entry1); - - let entry2 = store.values.get(&raw_b).unwrap(); - let expected_entry2 = StoreEntry::Single((key_b, value_b2)); - assert_eq!(entry2, &expected_entry2); - } - - #[test] - fn test_remove() { - // populate the value store - let mut store = ValueStore::default(); - - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE; WORD_SIZE]; - store.insert(key_a, value_a); - - let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ZERO, ONE, ZERO]; - store.insert(key_b, value_b); - - let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - store.insert(key_c, value_c); - - let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_d = [ZERO, ONE, ZERO, ZERO]; - store.insert(key_d, value_d); - - assert_eq!(store.values.len(), 2); - - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = - StoreEntry::List(vec![(key_c, value_c), (key_a, value_a), (key_d, value_d)]); - assert_eq!(entry1, &expected_entry1); - - let entry2 = store.values.get(&raw_b).unwrap(); - let expected_entry2 = StoreEntry::Single((key_b, value_b)); - assert_eq!(entry2, &expected_entry2); - - // remove non-existent keys - let key_e = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_a)]); - assert!(store.remove(&key_e).is_none()); - - let raw_f = 0b_11111110_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_f = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_f)]); - assert!(store.remove(&key_f).is_none()); - - // remove keys from the list entry - assert_eq!(store.remove(&key_c).unwrap(), value_c); - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = StoreEntry::List(vec![(key_a, value_a), (key_d, value_d)]); - assert_eq!(entry1, &expected_entry1); - - assert_eq!(store.remove(&key_a).unwrap(), value_a); - let entry1 = store.values.get(&raw_a).unwrap(); - let expected_entry1 = StoreEntry::Single((key_d, value_d)); - assert_eq!(entry1, &expected_entry1); - - assert_eq!(store.remove(&key_d).unwrap(), value_d); - assert!(store.values.get(&raw_a).is_none()); - assert_eq!(store.values.len(), 1); - - // remove a key from a single entry - assert_eq!(store.remove(&key_b).unwrap(), value_b); - assert!(store.values.get(&raw_b).is_none()); - assert_eq!(store.values.len(), 0); - } - - #[test] - fn test_range() { - // populate the value store - let mut store = ValueStore::default(); - - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE; WORD_SIZE]; - store.insert(key_a, value_a); - - let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ZERO, ONE, ZERO]; - store.insert(key_b, value_b); - - let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); - let value_c = [ONE, ONE, ZERO, ZERO]; - store.insert(key_c, value_c); - - let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); - let value_d = [ZERO, ONE, ZERO, ZERO]; - store.insert(key_d, value_d); - - let raw_e = 0b_10101000_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_e = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_e)]); - let value_e = [ZERO, ZERO, ZERO, ONE]; - store.insert(key_e, value_e); - - // check the entire range - let mut iter = store.range(..u64::MAX); - assert_eq!(iter.next(), Some(&(key_e, value_e))); - assert_eq!(iter.next(), Some(&(key_c, value_c))); - assert_eq!(iter.next(), Some(&(key_a, value_a))); - assert_eq!(iter.next(), Some(&(key_d, value_d))); - assert_eq!(iter.next(), Some(&(key_b, value_b))); - assert_eq!(iter.next(), None); - - // check all but e - let mut iter = store.range(raw_a..u64::MAX); - assert_eq!(iter.next(), Some(&(key_c, value_c))); - assert_eq!(iter.next(), Some(&(key_a, value_a))); - assert_eq!(iter.next(), Some(&(key_d, value_d))); - assert_eq!(iter.next(), Some(&(key_b, value_b))); - assert_eq!(iter.next(), None); - - // check all but e and b - let mut iter = store.range(raw_a..raw_b); - assert_eq!(iter.next(), Some(&(key_c, value_c))); - assert_eq!(iter.next(), Some(&(key_a, value_a))); - assert_eq!(iter.next(), Some(&(key_d, value_d))); - assert_eq!(iter.next(), None); - } - - #[test] - fn test_get_lone_sibling() { - // populate the value store - let mut store = ValueStore::default(); - - let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); - let value_a = [ONE; WORD_SIZE]; - store.insert(key_a, value_a); - - let raw_b = 0b_11111111_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64; - let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); - let value_b = [ONE, ZERO, ONE, ZERO]; - store.insert(key_b, value_b); - - // check sibling node for `a` - let index = LeafNodeIndex::make(32, 0b_10101010_10101010_00011111_11111110); - let parent_index = LeafNodeIndex::make(16, 0b_10101010_10101010); - assert_eq!(store.get_lone_sibling(index), Some((&key_a, &value_a, parent_index))); - - // check sibling node for `b` - let index = LeafNodeIndex::make(32, 0b_11111111_11111111_00011111_11111111); - let parent_index = LeafNodeIndex::make(16, 0b_11111111_11111111); - assert_eq!(store.get_lone_sibling(index), Some((&key_b, &value_b, parent_index))); - - // check some other sibling for some other index - let index = LeafNodeIndex::make(32, 0b_11101010_10101010); - assert_eq!(store.get_lone_sibling(index), None); - } -}