refactor: update SimpleSmt interfaces
This commit is contained in:
parent
0ddd0db89b
commit
fcf03478ba
7 changed files with 180 additions and 244 deletions
|
@ -18,8 +18,8 @@ fn smt_rpo(c: &mut Criterion) {
|
|||
(i, word)
|
||||
})
|
||||
.collect();
|
||||
let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap();
|
||||
trees.push(tree);
|
||||
let tree = SimpleSmt::with_leaves(depth, entries).unwrap();
|
||||
trees.push((tree, count));
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -29,10 +29,9 @@ fn smt_rpo(c: &mut Criterion) {
|
|||
|
||||
let mut insert = c.benchmark_group(format!("smt update_leaf"));
|
||||
|
||||
for tree in trees.iter_mut() {
|
||||
for (tree, count) in trees.iter_mut() {
|
||||
let depth = tree.depth();
|
||||
let count = tree.leaves_count() as u64;
|
||||
let key = count >> 2;
|
||||
let key = *count >> 2;
|
||||
insert.bench_with_input(
|
||||
format!("simple smt(depth:{depth},count:{count})"),
|
||||
&(key, leaf),
|
||||
|
@ -48,10 +47,9 @@ fn smt_rpo(c: &mut Criterion) {
|
|||
|
||||
let mut path = c.benchmark_group(format!("smt get_leaf_path"));
|
||||
|
||||
for tree in trees.iter_mut() {
|
||||
for (tree, count) in trees.iter_mut() {
|
||||
let depth = tree.depth();
|
||||
let count = tree.leaves_count() as u64;
|
||||
let key = count >> 2;
|
||||
let key = *count >> 2;
|
||||
path.bench_with_input(
|
||||
format!("simple smt(depth:{depth},count:{count})"),
|
||||
&key,
|
||||
|
|
|
@ -104,10 +104,7 @@ fn get_leaf_simplesmt(c: &mut Criterion) {
|
|||
.enumerate()
|
||||
.map(|(c, v)| (c.try_into().unwrap(), v.into()))
|
||||
.collect::<Vec<(u64, Word)>>();
|
||||
let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
|
||||
.unwrap()
|
||||
.with_leaves(smt_leaves.clone())
|
||||
.unwrap();
|
||||
let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
|
||||
let store = MerkleStore::from(&smt);
|
||||
let depth = smt.depth();
|
||||
let root = smt.root();
|
||||
|
@ -215,10 +212,7 @@ fn get_node_simplesmt(c: &mut Criterion) {
|
|||
.enumerate()
|
||||
.map(|(c, v)| (c.try_into().unwrap(), v.into()))
|
||||
.collect::<Vec<(u64, Word)>>();
|
||||
let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
|
||||
.unwrap()
|
||||
.with_leaves(smt_leaves.clone())
|
||||
.unwrap();
|
||||
let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
|
||||
let store = MerkleStore::from(&smt);
|
||||
let root = smt.root();
|
||||
let half_depth = smt.depth() / 2;
|
||||
|
@ -292,10 +286,7 @@ fn get_leaf_path_simplesmt(c: &mut Criterion) {
|
|||
.enumerate()
|
||||
.map(|(c, v)| (c.try_into().unwrap(), v.into()))
|
||||
.collect::<Vec<(u64, Word)>>();
|
||||
let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
|
||||
.unwrap()
|
||||
.with_leaves(smt_leaves.clone())
|
||||
.unwrap();
|
||||
let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
|
||||
let store = MerkleStore::from(&smt);
|
||||
let depth = smt.depth();
|
||||
let root = smt.root();
|
||||
|
@ -361,7 +352,7 @@ fn new(c: &mut Criterion) {
|
|||
.map(|(c, v)| (c.try_into().unwrap(), v.into()))
|
||||
.collect::<Vec<(u64, Word)>>()
|
||||
},
|
||||
|l| black_box(SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l)),
|
||||
|l| black_box(SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, l)),
|
||||
BatchSize::SmallInput,
|
||||
)
|
||||
});
|
||||
|
@ -376,7 +367,7 @@ fn new(c: &mut Criterion) {
|
|||
.collect::<Vec<(u64, Word)>>()
|
||||
},
|
||||
|l| {
|
||||
let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l).unwrap();
|
||||
let smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, l).unwrap();
|
||||
black_box(MerkleStore::from(&smt));
|
||||
},
|
||||
BatchSize::SmallInput,
|
||||
|
@ -442,10 +433,7 @@ fn update_leaf_simplesmt(c: &mut Criterion) {
|
|||
.enumerate()
|
||||
.map(|(c, v)| (c.try_into().unwrap(), v.into()))
|
||||
.collect::<Vec<(u64, Word)>>();
|
||||
let mut smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
|
||||
.unwrap()
|
||||
.with_leaves(smt_leaves.clone())
|
||||
.unwrap();
|
||||
let mut smt = SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, smt_leaves.clone()).unwrap();
|
||||
let mut store = MerkleStore::from(&smt);
|
||||
let depth = smt.depth();
|
||||
let root = smt.root();
|
||||
|
|
|
@ -1,6 +1,12 @@
|
|||
use super::{Felt, RpoDigest, WORD_SIZE, ZERO};
|
||||
use super::{Felt, RpoDigest, Word, WORD_SIZE, ZERO};
|
||||
use core::slice;
|
||||
|
||||
// CONSTANTS
|
||||
// ================================================================================================
|
||||
|
||||
/// A word consisting of 4 ZERO elements.
|
||||
pub const EMPTY_WORD: Word = [ZERO; WORD_SIZE];
|
||||
|
||||
// EMPTY NODES SUBTREES
|
||||
// ================================================================================================
|
||||
|
||||
|
@ -1570,7 +1576,7 @@ fn all_depths_opens_to_zero() {
|
|||
assert_eq!(depth as usize + 1, subtree.len());
|
||||
|
||||
// assert the opening is zero
|
||||
let initial = RpoDigest::new([ZERO; WORD_SIZE]);
|
||||
let initial = RpoDigest::new(EMPTY_WORD);
|
||||
assert_eq!(initial, subtree.remove(0));
|
||||
|
||||
// compute every node of the path manually and compare with the output
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use super::{
|
||||
hash::rpo::{Rpo256, RpoDigest},
|
||||
utils::collections::{vec, BTreeMap, Vec},
|
||||
utils::collections::{vec, BTreeMap, BTreeSet, Vec},
|
||||
Felt, StarkField, Word, WORD_SIZE, ZERO,
|
||||
};
|
||||
use core::fmt;
|
||||
|
@ -10,6 +10,7 @@ use core::fmt;
|
|||
|
||||
mod empty_roots;
|
||||
pub use empty_roots::EmptySubtreeRoots;
|
||||
use empty_roots::EMPTY_WORD;
|
||||
|
||||
mod index;
|
||||
pub use index::NodeIndex;
|
||||
|
@ -43,12 +44,13 @@ pub enum MerkleError {
|
|||
ConflictingRoots(Vec<Word>),
|
||||
DepthTooSmall(u8),
|
||||
DepthTooBig(u64),
|
||||
DuplicateValuesForKey(u64),
|
||||
NodeNotInStore(Word, NodeIndex),
|
||||
NumLeavesNotPowerOfTwo(usize),
|
||||
InvalidIndex { depth: u8, value: u64 },
|
||||
InvalidDepth { expected: u8, provided: u8 },
|
||||
InvalidPath(MerklePath),
|
||||
InvalidEntriesCount(usize, usize),
|
||||
InvalidNumEntries(usize, usize),
|
||||
NodeNotInSet(u64),
|
||||
RootNotInStore(Word),
|
||||
}
|
||||
|
@ -60,6 +62,7 @@ impl fmt::Display for MerkleError {
|
|||
ConflictingRoots(roots) => write!(f, "the merkle paths roots do not match {roots:?}"),
|
||||
DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"),
|
||||
DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"),
|
||||
DuplicateValuesForKey(key) => write!(f, "multiple values provided for key {key}"),
|
||||
NumLeavesNotPowerOfTwo(leaves) => {
|
||||
write!(f, "the leaves count {leaves} is not a power of 2")
|
||||
}
|
||||
|
@ -72,7 +75,7 @@ impl fmt::Display for MerkleError {
|
|||
"the provided depth {provided} is not valid for {expected}"
|
||||
),
|
||||
InvalidPath(_path) => write!(f, "the provided path is not valid"),
|
||||
InvalidEntriesCount(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
|
||||
InvalidNumEntries(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
|
||||
NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"),
|
||||
NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the store", hash, index.value(), index.depth(),),
|
||||
RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root),
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use super::{
|
||||
BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256,
|
||||
RpoDigest, Vec, Word,
|
||||
BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
|
||||
Rpo256, RpoDigest, Vec, Word, EMPTY_WORD,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
|
@ -10,6 +10,7 @@ mod tests;
|
|||
// ================================================================================================
|
||||
|
||||
/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
|
||||
///
|
||||
/// The root of the tree is recomputed on each new leaf update.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SimpleSmt {
|
||||
|
@ -20,18 +21,6 @@ pub struct SimpleSmt {
|
|||
empty_hashes: Vec<RpoDigest>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
struct BranchNode {
|
||||
left: RpoDigest,
|
||||
right: RpoDigest,
|
||||
}
|
||||
|
||||
impl BranchNode {
|
||||
fn parent(&self) -> RpoDigest {
|
||||
Rpo256::merge(&[self.left, self.right])
|
||||
}
|
||||
}
|
||||
|
||||
impl SimpleSmt {
|
||||
// CONSTANTS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
@ -45,7 +34,12 @@ impl SimpleSmt {
|
|||
// CONSTRUCTORS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
/// Creates a new simple SMT with the provided depth.
|
||||
/// Returns a new [SimpleSmt] instantiated with the specified depth.
|
||||
///
|
||||
/// All leaves in the returned tree are set to [ZERO; 4].
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if the depth is 0 or is greater than 64.
|
||||
pub fn new(depth: u8) -> Result<Self, MerkleError> {
|
||||
// validate the range of the depth.
|
||||
if depth < Self::MIN_DEPTH {
|
||||
|
@ -66,36 +60,47 @@ impl SimpleSmt {
|
|||
})
|
||||
}
|
||||
|
||||
/// Appends the provided entries as leaves of the tree.
|
||||
/// Returns a new [SimpleSmt] instantiated with the specified depth and with leaves
|
||||
/// set as specified by the provided entries.
|
||||
///
|
||||
/// All leaves omitted from the entries list are set to [ZERO; 4].
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// The function will fail if the provided entries count exceed the maximum tree capacity, that
|
||||
/// is `2^{depth}`.
|
||||
pub fn with_leaves<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
|
||||
/// Returns an error if:
|
||||
/// - If the depth is 0 or is greater than 64.
|
||||
/// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
|
||||
/// - The provided entries contain multiple values for the same key.
|
||||
pub fn with_leaves<R, I>(depth: u8, entries: R) -> Result<Self, MerkleError>
|
||||
where
|
||||
R: IntoIterator<IntoIter = I>,
|
||||
I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
|
||||
{
|
||||
// check if the leaves count will fit the depth setup
|
||||
let mut entries = entries.into_iter();
|
||||
let max = 1 << self.depth.min(63);
|
||||
// create an empty tree
|
||||
let mut tree = Self::new(depth)?;
|
||||
|
||||
// check if the number of leaves can be accommodated by the tree's depth; we use a min
|
||||
// depth of 63 because we consider passing in a vector of size 2^64 infeasible.
|
||||
let entries = entries.into_iter();
|
||||
let max = 1 << tree.depth.min(63);
|
||||
if entries.len() > max {
|
||||
return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
|
||||
return Err(MerkleError::InvalidNumEntries(max, entries.len()));
|
||||
}
|
||||
|
||||
// append leaves and return
|
||||
entries.try_for_each(|(key, leaf)| self.insert_leaf(key, leaf))?;
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Replaces the internal empty digests used when a given depth doesn't contain a node.
|
||||
pub fn with_empty_subtrees<I>(mut self, hashes: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = RpoDigest>,
|
||||
{
|
||||
self.replace_empty_subtrees(hashes.into_iter().collect());
|
||||
self
|
||||
// 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.update_leaf(key, value)?;
|
||||
if old_value != EMPTY_WORD || 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 == EMPTY_WORD && !empty_entries.insert(key) {
|
||||
return Err(MerkleError::DuplicateValuesForKey(key));
|
||||
}
|
||||
}
|
||||
Ok(tree)
|
||||
}
|
||||
|
||||
// PUBLIC ACCESSORS
|
||||
|
@ -111,40 +116,43 @@ impl SimpleSmt {
|
|||
self.depth
|
||||
}
|
||||
|
||||
// PROVIDERS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
/// Returns the set count of the keys of the leaves.
|
||||
pub fn leaves_count(&self) -> usize {
|
||||
self.leaves.len()
|
||||
}
|
||||
|
||||
/// Returns a node at the specified index.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if:
|
||||
/// * The specified depth is greater than the depth of the tree.
|
||||
/// Returns an error if the specified index has depth set to 0 or the depth is greater than
|
||||
/// the depth of this Merkle tree.
|
||||
pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
|
||||
if index.is_root() {
|
||||
Err(MerkleError::DepthTooSmall(index.depth()))
|
||||
} else if index.depth() > self.depth() {
|
||||
Err(MerkleError::DepthTooBig(index.depth() as u64))
|
||||
} else if index.depth() == self.depth() {
|
||||
self.get_leaf_node(index.value())
|
||||
.or_else(|| self.empty_hashes.get(index.depth() as usize).copied().map(Word::from))
|
||||
.ok_or(MerkleError::NodeNotInSet(index.value()))
|
||||
// the lookup in empty_hashes could fail only if empty_hashes were not built correctly
|
||||
// by the constructor as we check the depth of the lookup above.
|
||||
Ok(self
|
||||
.get_leaf_node(index.value())
|
||||
.unwrap_or_else(|| self.empty_hashes[index.depth() as usize].into()))
|
||||
} else {
|
||||
let branch_node = self.get_branch_node(&index);
|
||||
Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
|
||||
Ok(self.get_branch_node(&index).parent().into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a Merkle path from the node at the specified key to the root. The node itself is
|
||||
/// not included in the path.
|
||||
/// Returns a value of the leaf at the specified index.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if:
|
||||
/// * The specified depth is greater than the depth of the tree.
|
||||
/// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
|
||||
pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
|
||||
let index = NodeIndex::new(self.depth, index)?;
|
||||
self.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 has depth set to 0 or the depth is greater than
|
||||
/// the depth of this Merkle tree.
|
||||
pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
|
||||
if index.is_root() {
|
||||
return Err(MerkleError::DepthTooSmall(index.depth()));
|
||||
|
@ -163,18 +171,18 @@ impl SimpleSmt {
|
|||
Ok(path.into())
|
||||
}
|
||||
|
||||
/// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
|
||||
/// included in the path.
|
||||
/// Return a Merkle path from the leaf at the specified index to the root.
|
||||
///
|
||||
/// The leaf itself is not included in the path.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if:
|
||||
/// * The specified key does not exist as a leaf node.
|
||||
pub fn get_leaf_path(&self, key: u64) -> Result<MerklePath, MerkleError> {
|
||||
let index = NodeIndex::new(self.depth(), key)?;
|
||||
/// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
|
||||
pub fn get_leaf_path(&self, index: u64) -> Result<MerklePath, MerkleError> {
|
||||
let index = NodeIndex::new(self.depth(), index)?;
|
||||
self.get_path(index)
|
||||
}
|
||||
|
||||
/// Iterator over the inner nodes of the [SimpleSmt].
|
||||
/// Returns an iterator over the inner nodes of this Merkle tree.
|
||||
pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
|
||||
self.branches.values().map(|e| InnerNodeInfo {
|
||||
value: e.parent().into(),
|
||||
|
@ -186,27 +194,21 @@ impl SimpleSmt {
|
|||
// STATE MUTATORS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
/// Replaces the leaf located at the specified key, and recomputes hashes by walking up the
|
||||
/// tree.
|
||||
/// Updates value of the leaf at the specified index returning the old leaf value.
|
||||
///
|
||||
/// This also recomputes all hashes between the leaf and the root, updating the root itself.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if the specified key is not a valid leaf index for this tree.
|
||||
pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
|
||||
let index = NodeIndex::new(self.depth(), key)?;
|
||||
if !self.check_leaf_node_exists(key) {
|
||||
return Err(MerkleError::NodeNotInSet(index.value()));
|
||||
/// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}.
|
||||
pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<Word, MerkleError> {
|
||||
let old_value = self.insert_leaf_node(index, value).unwrap_or(EMPTY_WORD);
|
||||
|
||||
// if the old value and new value are the same, there is nothing to update
|
||||
if value == old_value {
|
||||
return Ok(value);
|
||||
}
|
||||
self.insert_leaf(key, value)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree
|
||||
pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
|
||||
self.insert_leaf_node(key, value);
|
||||
|
||||
// TODO consider using a map `index |-> word` instead of `index |-> (word, word)`
|
||||
let mut index = NodeIndex::new(self.depth(), key)?;
|
||||
let mut index = NodeIndex::new(self.depth(), index)?;
|
||||
let mut value = RpoDigest::from(value);
|
||||
for _ in 0..index.depth() {
|
||||
let is_right = index.is_value_odd();
|
||||
|
@ -217,26 +219,18 @@ impl SimpleSmt {
|
|||
value = Rpo256::merge(&[left, right]);
|
||||
}
|
||||
self.root = value.into();
|
||||
Ok(())
|
||||
Ok(old_value)
|
||||
}
|
||||
|
||||
// HELPER METHODS
|
||||
// --------------------------------------------------------------------------------------------
|
||||
|
||||
fn replace_empty_subtrees(&mut self, hashes: Vec<RpoDigest>) {
|
||||
self.empty_hashes = hashes;
|
||||
}
|
||||
|
||||
fn check_leaf_node_exists(&self, key: u64) -> bool {
|
||||
self.leaves.contains_key(&key)
|
||||
}
|
||||
|
||||
fn get_leaf_node(&self, key: u64) -> Option<Word> {
|
||||
self.leaves.get(&key).copied()
|
||||
}
|
||||
|
||||
fn insert_leaf_node(&mut self, key: u64, node: Word) {
|
||||
self.leaves.insert(key, node);
|
||||
fn insert_leaf_node(&mut self, key: u64, node: Word) -> Option<Word> {
|
||||
self.leaves.insert(key, node)
|
||||
}
|
||||
|
||||
fn get_branch_node(&self, index: &NodeIndex) -> BranchNode {
|
||||
|
@ -254,3 +248,18 @@ impl SimpleSmt {
|
|||
self.branches.insert(index, branch);
|
||||
}
|
||||
}
|
||||
|
||||
// BRANCH NODE
|
||||
// ================================================================================================
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
struct BranchNode {
|
||||
left: RpoDigest,
|
||||
right: RpoDigest,
|
||||
}
|
||||
|
||||
impl BranchNode {
|
||||
fn parent(&self) -> RpoDigest {
|
||||
Rpo256::merge(&[self.left, self.right])
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,9 +1,10 @@
|
|||
use super::{
|
||||
super::{int_to_node, InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt},
|
||||
NodeIndex, Rpo256, Vec, Word,
|
||||
NodeIndex, Rpo256, Vec, Word, EMPTY_WORD,
|
||||
};
|
||||
use proptest::prelude::*;
|
||||
use rand_utils::prng_array;
|
||||
|
||||
// TEST DATA
|
||||
// ================================================================================================
|
||||
|
||||
const KEYS4: [u64; 4] = [0, 1, 2, 3];
|
||||
const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
|
||||
|
@ -23,25 +24,17 @@ const VALUES8: [Word; 8] = [
|
|||
|
||||
const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8];
|
||||
|
||||
// TESTS
|
||||
// ================================================================================================
|
||||
|
||||
#[test]
|
||||
fn build_empty_tree() {
|
||||
// tree of depth 3
|
||||
let smt = SimpleSmt::new(3).unwrap();
|
||||
let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap();
|
||||
assert_eq!(mt.root(), smt.root());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_digests_are_consistent() {
|
||||
let depth = 5;
|
||||
let root = SimpleSmt::new(depth).unwrap().root();
|
||||
let computed: [RpoDigest; 2] = (0..depth).fold([Default::default(); 2], |state, _| {
|
||||
let digest = Rpo256::merge(&state);
|
||||
[digest; 2]
|
||||
});
|
||||
|
||||
assert_eq!(Word::from(computed[0]), root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_sparse_tree() {
|
||||
let mut smt = SimpleSmt::new(3).unwrap();
|
||||
|
@ -51,80 +44,59 @@ fn build_sparse_tree() {
|
|||
let key = 6;
|
||||
let new_node = int_to_node(7);
|
||||
values[key as usize] = new_node;
|
||||
smt.insert_leaf(key, new_node).expect("Failed to insert leaf");
|
||||
let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf");
|
||||
let mt2 = MerkleTree::new(values.clone()).unwrap();
|
||||
assert_eq!(mt2.root(), smt.root());
|
||||
assert_eq!(
|
||||
mt2.get_path(NodeIndex::make(3, 6)).unwrap(),
|
||||
smt.get_path(NodeIndex::make(3, 6)).unwrap()
|
||||
);
|
||||
assert_eq!(old_value, EMPTY_WORD);
|
||||
|
||||
// insert second value at distinct leaf branch
|
||||
let key = 2;
|
||||
let new_node = int_to_node(3);
|
||||
values[key as usize] = new_node;
|
||||
smt.insert_leaf(key, new_node).expect("Failed to insert leaf");
|
||||
let old_value = smt.update_leaf(key, new_node).expect("Failed to update leaf");
|
||||
let mt3 = MerkleTree::new(values).unwrap();
|
||||
assert_eq!(mt3.root(), smt.root());
|
||||
assert_eq!(
|
||||
mt3.get_path(NodeIndex::make(3, 2)).unwrap(),
|
||||
smt.get_path(NodeIndex::make(3, 2)).unwrap()
|
||||
);
|
||||
assert_eq!(old_value, EMPTY_WORD);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_full_tree() {
|
||||
let tree = SimpleSmt::new(2)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
|
||||
.unwrap();
|
||||
fn test_depth2_tree() {
|
||||
let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap();
|
||||
|
||||
// check internal structure
|
||||
let (root, node2, node3) = compute_internal_nodes();
|
||||
assert_eq!(root, tree.root());
|
||||
assert_eq!(node2, tree.get_node(NodeIndex::make(1, 0)).unwrap());
|
||||
assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_values() {
|
||||
let tree = SimpleSmt::new(2)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
|
||||
.unwrap();
|
||||
|
||||
// check depth 2
|
||||
// check get_node()
|
||||
assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap());
|
||||
assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap());
|
||||
assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap());
|
||||
assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_path() {
|
||||
let tree = SimpleSmt::new(2)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
|
||||
.unwrap();
|
||||
|
||||
let (_, node2, node3) = compute_internal_nodes();
|
||||
|
||||
// check depth 2
|
||||
// check get_path(): depth 2
|
||||
assert_eq!(vec![VALUES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap());
|
||||
assert_eq!(vec![VALUES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap());
|
||||
assert_eq!(vec![VALUES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap());
|
||||
assert_eq!(vec![VALUES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap());
|
||||
|
||||
// check depth 1
|
||||
// check get_path(): depth 1
|
||||
assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap());
|
||||
assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parent_node_iterator() -> Result<(), MerkleError> {
|
||||
let tree = SimpleSmt::new(2)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
|
||||
.unwrap();
|
||||
fn test_inner_node_iterator() -> Result<(), MerkleError> {
|
||||
let tree = SimpleSmt::with_leaves(2, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap();
|
||||
|
||||
// check depth 2
|
||||
assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap());
|
||||
|
@ -166,35 +138,28 @@ fn test_parent_node_iterator() -> Result<(), MerkleError> {
|
|||
|
||||
#[test]
|
||||
fn update_leaf() {
|
||||
let mut tree = SimpleSmt::new(3)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS8.into_iter().zip(VALUES8.into_iter()))
|
||||
.unwrap();
|
||||
let mut tree = SimpleSmt::with_leaves(3, KEYS8.into_iter().zip(VALUES8.into_iter())).unwrap();
|
||||
|
||||
// update one value
|
||||
let key = 3;
|
||||
let new_node = int_to_node(9);
|
||||
let mut expected_values = VALUES8.to_vec();
|
||||
expected_values[key] = new_node;
|
||||
let expected_tree = SimpleSmt::new(3)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS8.into_iter().zip(expected_values.clone().into_iter()))
|
||||
.unwrap();
|
||||
let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
|
||||
|
||||
tree.update_leaf(key as u64, new_node).unwrap();
|
||||
assert_eq!(expected_tree.root, tree.root);
|
||||
let old_leaf = tree.update_leaf(key as u64, new_node).unwrap();
|
||||
assert_eq!(expected_tree.root(), tree.root);
|
||||
assert_eq!(old_leaf, VALUES8[key]);
|
||||
|
||||
// update another value
|
||||
let key = 6;
|
||||
let new_node = int_to_node(10);
|
||||
expected_values[key] = new_node;
|
||||
let expected_tree = SimpleSmt::new(3)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS8.into_iter().zip(expected_values.into_iter()))
|
||||
.unwrap();
|
||||
let expected_tree = MerkleTree::new(expected_values.clone()).unwrap();
|
||||
|
||||
tree.update_leaf(key as u64, new_node).unwrap();
|
||||
assert_eq!(expected_tree.root, tree.root);
|
||||
let old_leaf = tree.update_leaf(key as u64, new_node).unwrap();
|
||||
assert_eq!(expected_tree.root(), tree.root);
|
||||
assert_eq!(old_leaf, VALUES8[key]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
@ -226,7 +191,7 @@ fn small_tree_opening_is_consistent() {
|
|||
|
||||
let depth = 3;
|
||||
let entries = vec![(0, a), (1, b), (4, c), (7, d)];
|
||||
let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap();
|
||||
let tree = SimpleSmt::with_leaves(depth, entries).unwrap();
|
||||
|
||||
assert_eq!(tree.root(), Word::from(k));
|
||||
|
||||
|
@ -250,56 +215,30 @@ fn small_tree_opening_is_consistent() {
|
|||
}
|
||||
}
|
||||
|
||||
proptest! {
|
||||
#[test]
|
||||
fn arbitrary_openings_single_leaf(
|
||||
depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
|
||||
key in prop::num::u64::ANY,
|
||||
leaf in prop::num::u64::ANY,
|
||||
) {
|
||||
let mut tree = SimpleSmt::new(depth).unwrap();
|
||||
#[test]
|
||||
fn fail_on_duplicates() {
|
||||
let entries = [(1_u64, int_to_node(1)), (5, int_to_node(2)), (1_u64, int_to_node(3))];
|
||||
let smt = SimpleSmt::with_leaves(64, entries);
|
||||
assert!(smt.is_err());
|
||||
|
||||
let key = key % (1 << depth as u64);
|
||||
let leaf = int_to_node(leaf);
|
||||
let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2)), (1_u64, int_to_node(0))];
|
||||
let smt = SimpleSmt::with_leaves(64, entries);
|
||||
assert!(smt.is_err());
|
||||
|
||||
tree.insert_leaf(key, leaf.into()).unwrap();
|
||||
tree.get_leaf_path(key).unwrap();
|
||||
let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2)), (1_u64, int_to_node(1))];
|
||||
let smt = SimpleSmt::with_leaves(64, entries);
|
||||
assert!(smt.is_err());
|
||||
|
||||
// traverse to root, fetching all paths
|
||||
for d in 1..depth {
|
||||
let k = key >> (depth - d);
|
||||
tree.get_path(NodeIndex::make(d, k)).unwrap();
|
||||
}
|
||||
}
|
||||
let entries = [(1_u64, int_to_node(1)), (5, int_to_node(2)), (1_u64, int_to_node(0))];
|
||||
let smt = SimpleSmt::with_leaves(64, entries);
|
||||
assert!(smt.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn arbitrary_openings_multiple_leaves(
|
||||
depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
|
||||
count in 2u8..10u8,
|
||||
ref seed in any::<[u8; 32]>()
|
||||
) {
|
||||
let mut tree = SimpleSmt::new(depth).unwrap();
|
||||
let mut seed = *seed;
|
||||
let leaves = (1 << depth) - 1;
|
||||
|
||||
for _ in 0..count {
|
||||
seed = prng_array(seed);
|
||||
|
||||
let mut key = [0u8; 8];
|
||||
let mut leaf = [0u8; 8];
|
||||
|
||||
key.copy_from_slice(&seed[..8]);
|
||||
leaf.copy_from_slice(&seed[8..16]);
|
||||
|
||||
let key = u64::from_le_bytes(key);
|
||||
let key = key % leaves;
|
||||
let leaf = u64::from_le_bytes(leaf);
|
||||
let leaf = int_to_node(leaf);
|
||||
|
||||
tree.insert_leaf(key, leaf).unwrap();
|
||||
tree.get_leaf_path(key).unwrap();
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn with_no_duplicates_empty_node() {
|
||||
let entries = [(1_u64, int_to_node(0)), (5, int_to_node(2))];
|
||||
let smt = SimpleSmt::with_leaves(64, entries);
|
||||
assert!(smt.is_ok());
|
||||
}
|
||||
|
||||
// HELPER FUNCTIONS
|
||||
|
|
|
@ -183,10 +183,7 @@ fn test_get_invalid_node() {
|
|||
fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
|
||||
let keys2: [u64; 2] = [0, 1];
|
||||
let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
|
||||
let smt = SimpleSmt::new(1)
|
||||
.unwrap()
|
||||
.with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
|
||||
.unwrap();
|
||||
let smt = SimpleSmt::with_leaves(1, keys2.into_iter().zip(leaves2.into_iter())).unwrap();
|
||||
let store = MerkleStore::from(&smt);
|
||||
|
||||
let idx = NodeIndex::make(1, 0);
|
||||
|
@ -202,10 +199,9 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
|
|||
|
||||
#[test]
|
||||
fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
|
||||
let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
|
||||
.unwrap();
|
||||
let smt =
|
||||
SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, KEYS4.into_iter().zip(LEAVES4.into_iter()))
|
||||
.unwrap();
|
||||
|
||||
let store = MerkleStore::from(&smt);
|
||||
|
||||
|
@ -526,10 +522,7 @@ fn test_constructors() -> Result<(), MerkleError> {
|
|||
}
|
||||
|
||||
let depth = 32;
|
||||
let smt = SimpleSmt::new(depth)
|
||||
.unwrap()
|
||||
.with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
|
||||
.unwrap();
|
||||
let smt = SimpleSmt::with_leaves(depth, KEYS4.into_iter().zip(LEAVES4.into_iter())).unwrap();
|
||||
let store = MerkleStore::from(&smt);
|
||||
let depth = smt.depth();
|
||||
|
||||
|
|
Loading…
Add table
Reference in a new issue