From 463e90b75c87e6317872c436afdca746d61b6dd8 Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Thu, 8 Dec 2022 17:13:17 +0100 Subject: [PATCH] feat: add simple sparse merkle tree This commit introduces the `SparseMerkleTree` and `TreeStorage`. The sparse merkle tree is a concrete structure that will provide the functionalities of a keyed naive SMT with inclusion proof. It introduces the `MerklePath` structure to fully encapsulate a merkle opening verification - however, it can be just converted into a vector of digests so it will be compatible with the opening prior to this commit. It also introduces the tree storage trait, a backend definition that will be the responsible to manage the state of a tree. This brings the advantage of decoupling the implementation logic of the merkle tree with its backend requirements, and will use `alloc::borrow::Cow` instead of either references or owned values to give further flexibility to the storage. Some storages will lock its own state to fetch a value, and they might provide values to multi-threaded applications - hence, extending the lock might not be desirable for fetched values. In such cases, the storage will have the option to return the value as owned, freeing it from its lifetime bound. There is a memory storage added that will provide in-memory tree storage, backed by `BTreeMap`. closes #21 --- src/hash/mod.rs | 3 +- src/lib.rs | 14 +- src/merkle/merkle_tree.rs | 13 +- src/merkle/mod.rs | 50 +++++- src/merkle/simple_smt/mod.rs | 274 +++++++++++++++++++++++++++++++ src/merkle/simple_smt/path.rs | 94 +++++++++++ src/merkle/simple_smt/storage.rs | 103 ++++++++++++ src/merkle/simple_smt/tests.rs | 168 +++++++++++++++++++ 8 files changed, 705 insertions(+), 14 deletions(-) create mode 100644 src/merkle/simple_smt/mod.rs create mode 100644 src/merkle/simple_smt/path.rs create mode 100644 src/merkle/simple_smt/storage.rs create mode 100644 src/merkle/simple_smt/tests.rs diff --git a/src/hash/mod.rs b/src/hash/mod.rs index 9508754c..ee87895b 100644 --- a/src/hash/mod.rs +++ b/src/hash/mod.rs @@ -1,5 +1,4 @@ -use super::{Felt, FieldElement, StarkField, ONE, ZERO}; -use winter_crypto::{Digest, ElementHasher, Hasher}; +use super::{Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField, ONE, ZERO}; pub mod blake; pub mod rpo; diff --git a/src/lib.rs b/src/lib.rs index ff2eb43f..2885d900 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ #![cfg_attr(not(feature = "std"), no_std)] #[cfg(not(feature = "std"))] -#[cfg_attr(test, macro_use)] +#[macro_use] extern crate alloc; pub mod hash; @@ -10,6 +10,7 @@ pub mod merkle; // RE-EXPORTS // ================================================================================================ +pub use winter_crypto::{Digest, ElementHasher, Hasher}; pub use winter_math::{fields::f64::BaseElement as Felt, FieldElement, StarkField}; pub mod utils { @@ -17,17 +18,26 @@ pub mod utils { collections, string, uninit_vector, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable, SliceReader, }; + + #[cfg(not(feature = "std"))] + pub use alloc::borrow::Cow; + + #[cfg(feature = "std")] + pub use std::borrow::Cow; } // TYPE ALIASES // ================================================================================================ /// A group of four field elements in the Miden base field. -pub type Word = [Felt; 4]; +pub type Word = [Felt; WORD_SIZE]; // CONSTANTS // ================================================================================================ +/// Number of field elements in a word. +pub const WORD_SIZE: usize = 4; + /// Field element representing ZERO in the Miden base filed. pub const ZERO: Felt = Felt::ZERO; diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 8763d96b..c0c381cd 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -1,4 +1,4 @@ -use super::{Digest, Felt, MerkleError, Rpo256, Vec, Word}; +use super::{Felt, MerkleError, Rpo256, RpoDigest, Vec, Word}; use crate::{utils::uninit_vector, FieldElement}; use core::slice; use winter_math::log2; @@ -22,7 +22,7 @@ impl MerkleTree { pub fn new(leaves: Vec) -> Result { let n = leaves.len(); if n <= 1 { - return Err(MerkleError::DepthTooSmall); + return Err(MerkleError::DepthTooSmall(n as u32)); } else if !n.is_power_of_two() { return Err(MerkleError::NumLeavesNotPowerOfTwo(n)); } @@ -35,7 +35,8 @@ impl MerkleTree { nodes[n..].copy_from_slice(&leaves); // re-interpret nodes as an array of two nodes fused together - let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [Digest; 2], n) }; + let two_nodes = + unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [RpoDigest; 2], n) }; // calculate all internal tree nodes for i in (1..n).rev() { @@ -68,7 +69,7 @@ impl MerkleTree { /// * The specified index not valid for the specified depth. pub fn get_node(&self, depth: u32, index: u64) -> Result { if depth == 0 { - return Err(MerkleError::DepthTooSmall); + return Err(MerkleError::DepthTooSmall(depth)); } else if depth > self.depth() { return Err(MerkleError::DepthTooBig(depth)); } @@ -89,7 +90,7 @@ impl MerkleTree { /// * The specified index not valid for the specified depth. pub fn get_path(&self, depth: u32, index: u64) -> Result, MerkleError> { if depth == 0 { - return Err(MerkleError::DepthTooSmall); + return Err(MerkleError::DepthTooSmall(depth)); } else if depth > self.depth() { return Err(MerkleError::DepthTooBig(depth)); } @@ -123,7 +124,7 @@ impl MerkleTree { let n = self.nodes.len() / 2; let two_nodes = - unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [Digest; 2], n) }; + unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [RpoDigest; 2], n) }; for _ in 0..depth { index /= 2; diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 1b137b9b..a04b0932 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -1,8 +1,13 @@ use super::{ - hash::rpo::{Rpo256, RpoDigest as Digest}, - utils::collections::{BTreeMap, Vec}, - Felt, Word, ZERO, + hash::rpo::{Rpo256, RpoDigest}, + utils::{ + collections::{BTreeMap, Vec}, + string::String, + Cow, + }, + Digest, Felt, Hasher, Word, ZERO, }; +use core::{convert::Infallible, fmt}; mod merkle_tree; pub use merkle_tree::MerkleTree; @@ -10,20 +15,57 @@ pub use merkle_tree::MerkleTree; mod merkle_path_set; pub use merkle_path_set::MerklePathSet; +mod simple_smt; +pub use simple_smt::{DepthDigest, MemoryTreeStorage, MerklePath, SparseMerkleTree, TreeStorage}; + // ERRORS // ================================================================================================ #[derive(Clone, Debug)] pub enum MerkleError { - DepthTooSmall, + DepthTooSmall(u32), DepthTooBig(u32), NumLeavesNotPowerOfTwo(usize), InvalidIndex(u32, u64), InvalidDepth(u32, u32), InvalidPath(Vec), NodeNotInSet(u64), + StorageInconsistency(String), +} + +impl From for MerkleError { + fn from(_e: Infallible) -> Self { + unreachable!() + } } +impl fmt::Display for MerkleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + use MerkleError::*; + match self { + DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"), + DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"), + NumLeavesNotPowerOfTwo(leaves) => { + write!(f, "the leaves count {leaves} is not a power of 2") + } + InvalidIndex(depth, index) => write!( + f, + "the leaf index {index} is not valid for the depth {depth}" + ), + InvalidDepth(expected, provided) => write!( + f, + "the provided depth {provided} is not valid for {expected}" + ), + InvalidPath(_path) => write!(f, "the provided path is not valid"), + NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"), + StorageInconsistency(message) => write!(f, "storage inconsistency: {message}"), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for MerkleError {} + // HELPER FUNCTIONS // ================================================================================================ diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs new file mode 100644 index 00000000..9596ae5a --- /dev/null +++ b/src/merkle/simple_smt/mod.rs @@ -0,0 +1,274 @@ +use super::{BTreeMap, Cow, Digest, Hasher, MerkleError, Vec}; +use core::{marker::PhantomData, mem::replace}; + +mod path; +pub use path::MerklePath; + +mod storage; +pub use storage::{DepthDigest, MemoryTreeStorage, TreeStorage}; + +#[cfg(test)] +mod tests; + +// SPARSE MERKLE TREE +// ================================================================================================ + +/// A simple SMT backed by a hasher and tree storage. +/// +/// This structure will not store any data, except the maximum length of the tree. All the data +/// will be sent to its storage implementation. +pub struct SparseMerkleTree +where + H: Hasher, + S: TreeStorage, +{ + hasher: PhantomData, + storage: S, + depth: u32, +} + +impl SparseMerkleTree +where + H: Hasher, + S: TreeStorage, +{ + // CONSTANTS + // -------------------------------------------------------------------------------------------- + + /// Minimum supported depth. + pub const MIN_DEPTH: u32 = 1; + + /// Maximum supported depth. + pub const MAX_DEPTH: u32 = 63; + + // CONSTRUCTORS + // -------------------------------------------------------------------------------------------- + + /// Create a new tree with a fixed depth. + /// + /// # Errors + /// + /// This function will fail with either a depth out of range (check [`Self::MIN_DEPTH`] and + /// [`Self::MAX_DEPTH`], or an internal storage error. + pub fn new(depth: u32) -> Result { + // validate the range of the depth. + if depth < Self::MIN_DEPTH { + return Err(MerkleError::DepthTooSmall(depth)); + } else if Self::MAX_DEPTH < depth { + return Err(MerkleError::DepthTooBig(depth)); + } + + // with an internal state of the current depth empty digest, iterate from root to leaf. + let empty_digests = (0..=depth).scan(H::Digest::default(), |state, i| { + Some(DepthDigest { + // the scan starts with the leaves and ends with the root. + depth: depth - i, + // the next state is the merge of the duplicate of the current state. + digest: replace(state, H::merge(&[*state; 2])), + }) + }); + let storage = S::new(depth, empty_digests).map_err(|e| e.into())?; + + Ok(Self { + hasher: PhantomData, + storage, + depth, + }) + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Insert a new keyed leaf. + /// + /// # Errors + /// + /// This function will fail if the provided key is not within the range of the depth, that is + /// `[0, 2^{depth})`, or if there is an internal storage error. + pub fn insert(&mut self, mut key: u64, leaf: H::Digest) -> Result<(), MerkleError> { + // check the key range against the depth. + if 1 << self.depth as u64 <= key { + return Err(MerkleError::InvalidIndex(self.depth, key)); + } + + // insert the leaf into the target depth. + self.storage + .insert(self.depth, key, leaf) + .map_err(|e| e.into())?; + + // traverse the tree from leaf to root, updating left or right node depending on the parity + // of the key + for n in (1..=self.depth).rev() { + let digest = merge_with_sibling::(n, key, &self.storage)?; + key /= 2; + self.storage + .insert(n - 1, key, digest) + .map_err(|e| e.into())?; + } + + Ok(()) + } + + // GETTERS + // -------------------------------------------------------------------------------------------- + + /// Fetch the current root commitment of the tree. + /// + /// # Errors + /// + /// Unless there is an internal storage failure, this function is infallible. + pub fn root(&self) -> Result { + // attempt to fetch a note for depth 0 (root) and key 0 (first element) + if let Some(root) = self.storage.get(0, 0).map_err(|e| e.into())? { + return Ok(root.into_owned()); + } + // if no node is found for (0,0), then it hasn't been calculated yet; should fetch empty + // digest instead. + empty_digest(0, &self.storage).map(|digest| digest.into_owned()) + } + + /// Check if the provided key exists in the storage, returning `true` if it does. + pub fn leaf_exists(&self, key: u64) -> Result { + self.storage.exists(self.depth, key).map_err(|e| e.into()) + } + + /// Attempt to fetch the keyed leaf from the storage, returning `Ok(None)` if it was not + /// previously inserted. + pub fn get(&self, key: u64) -> Result, MerkleError> { + Ok(self + .storage + .get(self.depth, key) + .map_err(|e| e.into())? + .map(Cow::into_owned)) + } + + /// Create a new opening path to the commitment root from the keyed leaf. Will return + /// `Ok(None)` if the keyed leaf doesn't exist in the storage. + pub fn get_path(&self, key: u64) -> Result>, MerkleError> { + // fetch the leaf + let leaf = match self.storage.get(self.depth, key).map_err(|e| e.into())? { + Some(l) => l.into_owned(), + None => return Ok(None), + }; + + // fetch the root. + // + // if a leaf exists, then the storage logically will also have its root. a failure here is + // an inconsistency of the storage. + let root = match self.storage.get(0, 0).map_err(|e| e.into())? { + Some(root) => root.into_owned(), + None => { + return Err(MerkleError::StorageInconsistency(format!( + "the leaf {key} exists in the storage, but it couldn't provide its root" + ))) + } + }; + + // iterate from leaf to root with an internal state equivalent to either `[node, sibling]` + // or `[sibling, node]`, depending on the parity of the key index. + let path = (1..=self.depth) + .rev() + .scan(key, |key, n| { + match fetch_sibling::(n, *key, &self.storage) { + Ok(sibling) => { + *key /= 2; + Some(Ok(sibling)) + } + Err(e) => Some(Err(e)), + } + }) + .collect::>()?; + + Ok(Some(MerklePath::new(leaf, root, key, path))) + } +} + +// SPARSE MERKLE TREE HELPERS +// ================================================================================================ + +/// Fetch the empty digest for a given depth in the provided storage. +fn empty_digest(depth: u32, storage: &S) -> Result, MerkleError> +where + S: TreeStorage, +{ + storage + .empty_digest(depth) + .map_err(|e| e.into())? + .ok_or_else(|| { + MerkleError::StorageInconsistency( + format!("the storage was initialized for the depth {depth}, but failed to provide an empty digest") + ) + }) +} + +/// Fetch the sibling node of a given node from the provided storage. +/// +/// Will be inconsistent if `depth == 0` since the root node has no sibling. +/// +/// # Example +/// +/// a +/// / \ +/// b c +/// +/// fetch_sibling(1, 1, _) |-> b +fn fetch_sibling(depth: u32, key: u64, storage: &S) -> Result +where + H: Hasher, + S: TreeStorage, +{ + debug_assert!(depth != 0); + + let is_right = key & 1 == 1; + let sibling = if is_right { key - 1 } else { key + 1 }; + + Ok(storage + .get(depth, sibling) + .map_err(|e| e.into()) + .transpose() + .unwrap_or_else(|| empty_digest(depth, storage))? + .into_owned()) +} + +/// Merge a node with its sibling with the provided hash. +/// +/// Will be inconsistent if `depth == 0` since the root node has no sibling. +/// +/// # Example +/// +/// a +/// / \ +/// b c +/// +/// merge_with_sibling(1, 1, _) |-> H([b, c]) +fn merge_with_sibling(depth: u32, key: u64, storage: &S) -> Result +where + H: Hasher, + S: TreeStorage, +{ + debug_assert!(depth != 0); + + let node = storage + .get(depth, key) + .map_err(|e| e.into()) + .transpose() + .unwrap_or_else(|| empty_digest(depth, storage))? + .into_owned(); + + let is_right = key & 1 == 1; + let sibling = if is_right { key - 1 } else { key + 1 }; + let sibling = storage + .get(depth, sibling) + .map_err(|e| e.into()) + .transpose() + .unwrap_or_else(|| empty_digest(depth, storage))? + .into_owned(); + + let input = if is_right { + [sibling, node] + } else { + [node, sibling] + }; + + Ok(H::merge(&input)) +} diff --git a/src/merkle/simple_smt/path.rs b/src/merkle/simple_smt/path.rs new file mode 100644 index 00000000..1d2a27de --- /dev/null +++ b/src/merkle/simple_smt/path.rs @@ -0,0 +1,94 @@ +use super::{Hasher, Vec}; +use core::marker::PhantomData; + +// SPARSE MERKLE TREE PATH +// ================================================================================================ + +/// A self-contained merkle opening proof that will traverse from a leaf with its siblings until +/// the root. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MerklePath +where + H: Hasher, +{ + hasher: PhantomData, + key: u64, + leaf: H::Digest, + root: H::Digest, + /// Path, not inclusive (won't contain leaf nor root), from the leaf to the root + path: Vec, +} + +impl MerklePath +where + H: Hasher, +{ + // CONSTRUCTORS + // -------------------------------------------------------------------------------------------- + + /// Create a new merkle opening path. The path will be a vector starting with the sibling of + /// the leaf until the sibling that composes the root. + pub fn new(leaf: H::Digest, root: H::Digest, key: u64, path: Vec) -> Self { + Self { + hasher: PhantomData, + key, + leaf, + root, + path, + } + } + + /// Depth of the tree that generated this opening proof. + pub fn depth(&self) -> usize { + self.path.len() + } + + /// Key of the leaf that was fetched from the tree that generated this opening proof. + pub const fn key(&self) -> u64 { + self.key + } + + /// Leaf of the opening proof. + pub const fn leaf(&self) -> &H::Digest { + &self.leaf + } + + /// Commitment root that leaf is a member of. + pub const fn root(&self) -> &H::Digest { + &self.root + } + + /// Siblings path from leaf to root. + pub fn path(&self) -> &[H::Digest] { + &self.path + } + + /// Takes siblings path as owned. + pub fn into_path(self) -> Vec { + self.path + } + + /// Check the correctness of the opening proof. + pub fn check(&self) -> bool { + let mut key = self.key; + + let computed = self.path.iter().fold(self.leaf, |state, node| { + // check the parity of the key and shl + let is_odd = key & 1 == 1; + key /= 2; + + // define the input of the merge depending on the parity + let input = if is_odd { + [*node, state] + } else { + [state, *node] + }; + + // compute the merged nodes + H::merge(&input) + }); + + // if the final computed node is the root, then the proof is correct + computed == self.root + } +} diff --git a/src/merkle/simple_smt/storage.rs b/src/merkle/simple_smt/storage.rs new file mode 100644 index 00000000..fd5320e5 --- /dev/null +++ b/src/merkle/simple_smt/storage.rs @@ -0,0 +1,103 @@ +use super::{BTreeMap, Cow, Digest, MerkleError}; +use core::convert::Infallible; + +// MERKLE TREE STORAGE +// ================================================================================================ + +/// A backend storage for merkle tree implementations. +/// +/// Will have the minimal set of functionalities in order to provide the required resources to +/// different merkle tree designs. +/// +/// It uses [`Cow`] for any fetched data so the storage will have design freedom to decide on +/// whether or not it will return owned values. +pub trait TreeStorage: Sized { + /// The digest will represent the nodes and leaves of the tree. + type Digest: Digest; + /// The error type of the storage. Will always be represented as + /// [`MerkleError::StorageInconsistency`]. + type Error: Into; + + /// Create a new tree with a given max depth and a set of empty digests. The empty digests + /// should be stored as immutable and be retrieved via [`Self::empty_digest`]. + /// + /// A consistent storage will have a count of empty digests equal to the max depth + 1. + fn new(max_depth: u32, empty_digests: I) -> Result + where + I: Iterator>; + + /// An empty digest of a given depth. + fn empty_digest(&self, depth: u32) -> Result>, Self::Error>; + + /// Check if the provided key exists in the set. + fn exists(&self, depth: u32, key: u64) -> Result { + self.get(depth, key).map(|l| l.is_some()) + } + + /// Fetch a node or leaf from the storage. + fn get(&self, depth: u32, key: u64) -> Result>, Self::Error>; + + /// Insert a node or leaf into the storage. + fn insert(&mut self, depth: u32, key: u64, value: Self::Digest) -> Result<(), Self::Error>; +} + +// MERKLE TREE STORAGE HELPER STRUCTURES +// ================================================================================================ + +/// A depth + digest tuple representation. +/// +/// Will be used to initiate a storage with its empty digest set. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct DepthDigest { + /// Tree depth linked to the digest. + pub depth: u32, + /// Provided digest. + pub digest: D, +} + +// MEMORY MERKLE TREE STORAGE +// ================================================================================================ + +/// A concrete in-memory, infallible merkle storage provider. +/// +/// The infallibility relies on the properties of [`BTreeMap`], so no previously inserted node +/// should fail to be fetched. Unless there is memory exhaustion of the host, this provider will +/// not fail. +pub struct MemoryTreeStorage +where + D: Digest, +{ + empty_digests: BTreeMap, + nodes: BTreeMap<(u32, u64), D>, +} + +impl TreeStorage for MemoryTreeStorage +where + D: Digest, +{ + type Digest = D; + type Error = Infallible; + + fn new(_max_depth: u32, empty_digests: I) -> Result + where + I: Iterator>, + { + Ok(Self { + empty_digests: empty_digests.map(|d| (d.depth, d.digest)).collect(), + nodes: BTreeMap::default(), + }) + } + + fn empty_digest(&self, depth: u32) -> Result>, Self::Error> { + Ok(self.empty_digests.get(&depth).map(Cow::Borrowed)) + } + + fn get(&self, depth: u32, key: u64) -> Result>, Self::Error> { + Ok(self.nodes.get(&(depth, key)).map(Cow::Borrowed)) + } + + fn insert(&mut self, depth: u32, key: u64, value: Self::Digest) -> Result<(), Self::Error> { + self.nodes.insert((depth, key), value); + Ok(()) + } +} diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs new file mode 100644 index 00000000..63d6e39f --- /dev/null +++ b/src/merkle/simple_smt/tests.rs @@ -0,0 +1,168 @@ +use super::*; +use crate::hash::{blake::Blake3_256, rpo::Rpo256}; +use proptest::prelude::*; +use rand_utils::prng_array; + +type Digest = ::Digest; +type MemorySmt = SparseMerkleTree>; + +type RpoDigest = ::Digest; +type RpoMemorySmt = SparseMerkleTree>; + +#[test] +fn empty_digests_are_consistent() { + let depth = 5; + let root = MemorySmt::new(depth).unwrap().root().unwrap(); + let computed = (0..depth).fold([Default::default(); 2], |state, _| { + let digest = Blake3_256::merge(&state); + [digest; 2] + }); + + assert_eq!(computed[0], root); +} + +#[test] +fn small_tree_opening_is_consistent() { + // ____k____ + // / \ + // _i_ _j_ + // / \ / \ + // e f g h + // / \ / \ / \ / \ + // a b 0 0 c 0 0 d + + let z = Default::default(); + + let a = Blake3_256::merge(&[z; 2]); + let b = Blake3_256::merge(&[a; 2]); + let c = Blake3_256::merge(&[b; 2]); + let d = Blake3_256::merge(&[c; 2]); + + let e = Blake3_256::merge(&[a, b]); + let f = Blake3_256::merge(&[z; 2]); + let g = Blake3_256::merge(&[c, z]); + let h = Blake3_256::merge(&[z, d]); + + let i = Blake3_256::merge(&[e, f]); + let j = Blake3_256::merge(&[g, h]); + + let k = Blake3_256::merge(&[i, j]); + + let depth = 3; + let mut tree = MemorySmt::new(depth).unwrap(); + + let cases = vec![ + (0, a, [b, f, j]), + (1, b, [a, f, j]), + (4, c, [z, h, i]), + (7, d, [z, g, i]), + ]; + + for (i, leaf, _) in cases.clone() { + tree.insert(i, leaf).unwrap(); + } + + assert_eq!(k, tree.root().unwrap()); + + for (i, leaf, path) in cases { + let opening = tree.get_path(i).unwrap().unwrap(); + + assert_eq!(depth, opening.depth() as u32); + assert_eq!(i, opening.key()); + assert_eq!(&leaf, opening.leaf()); + assert_eq!(&k, opening.root()); + assert_eq!(opening.path(), &path); + assert!(opening.check()); + } +} + +proptest! { + #[test] + fn arbitrary_openings_single_leaf( + depth in MemorySmt::MIN_DEPTH..MemorySmt::MAX_DEPTH, + key in prop::num::u64::ANY, + ref bytes in any::>() + ) { + let mut tree = MemorySmt::new(depth).unwrap(); + + let key = key % (1 << depth as u64); + let leaf = Blake3_256::hash(&bytes); + tree.insert(key, leaf).unwrap(); + + let opening = tree.get_path(key).unwrap().unwrap(); + assert!(opening.check()); + } + + #[test] + fn arbitrary_openings_single_leaf_rpo( + depth in RpoMemorySmt::MIN_DEPTH..RpoMemorySmt::MAX_DEPTH, + key in prop::num::u64::ANY, + ref bytes in any::<[u8; 32]>() + ) { + let mut tree = RpoMemorySmt::new(depth).unwrap(); + + let key = key % (1 << depth as u64); + let leaf = Rpo256::hash(&bytes[..]); + tree.insert(key, leaf).unwrap(); + + let opening = tree.get_path(key).unwrap().unwrap(); + assert!(opening.check()); + } + + #[test] + fn arbitrary_openings_multiple_leaves( + depth in MemorySmt::MIN_DEPTH..MemorySmt::MAX_DEPTH, + count in prop::num::u8::ANY, + ref seed in any::<[u8; 32]>() + ) { + let mut tree = MemorySmt::new(depth).unwrap(); + let mut seed = *seed; + + for _ in 0..count { + seed = prng_array(seed); + + let mut key = [0u8; 8]; + let mut leaf = [0u8; 24]; + + key.copy_from_slice(&seed[..8]); + leaf.copy_from_slice(&seed[8..]); + + let key = u64::from_le_bytes(key); + let key = key % (1 << depth as u64); + let leaf = Blake3_256::hash(&leaf); + tree.insert(key, leaf).unwrap(); + + let opening = tree.get_path(key).unwrap().unwrap(); + assert!(opening.check()); + } + } + + #[test] + fn arbitrary_openings_multiple_leaves_rpo( + depth in MemorySmt::MIN_DEPTH..MemorySmt::MAX_DEPTH, + // RPO is more expensive than blake; trim the count + count in 0u8..10u8, + ref seed in any::<[u8; 32]>() + ) { + let mut tree = RpoMemorySmt::new(depth).unwrap(); + let mut seed = *seed; + + for _ in 0..count { + seed = prng_array(seed); + + let mut key = [0u8; 8]; + let mut leaf = [0u8; 24]; + + key.copy_from_slice(&seed[..8]); + leaf.copy_from_slice(&seed[8..]); + + let key = u64::from_le_bytes(key); + let key = key % (1 << depth as u64); + let leaf = Rpo256::hash(&leaf); + tree.insert(key, leaf).unwrap(); + + let opening = tree.get_path(key).unwrap().unwrap(); + assert!(opening.check()); + } + } +}