diff --git a/Cargo.lock b/Cargo.lock index ddd40fbd..c8150b9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -116,6 +116,7 @@ dependencies = [ "either", "proptest", "rand", + "rand_chacha", "rand_core", ] diff --git a/incrementalmerkletree/Cargo.toml b/incrementalmerkletree/Cargo.toml index 4239dfa5..067da673 100644 --- a/incrementalmerkletree/Cargo.toml +++ b/incrementalmerkletree/Cargo.toml @@ -27,6 +27,7 @@ rand_core = { version = "0.6", optional = true } proptest = "1.0.0" rand = "0.8" rand_core = "0.6" +rand_chacha = "0.3" [features] # The legacy-api feature guards types and functions that were previously diff --git a/incrementalmerkletree/src/frontier.rs b/incrementalmerkletree/src/frontier.rs index 09e7ef97..854f26b1 100644 --- a/incrementalmerkletree/src/frontier.rs +++ b/incrementalmerkletree/src/frontier.rs @@ -1,5 +1,6 @@ -use std::convert::TryFrom; use std::mem::size_of; +use std::num::NonZeroU64; +use std::{convert::TryFrom, num::NonZeroU8}; use crate::{Address, Hashable, Level, MerklePath, Position, Source}; @@ -182,6 +183,92 @@ impl NonEmptyFrontier { } } +#[cfg(any(test, feature = "test-dependencies"))] +impl NonEmptyFrontier +where + Standard: Distribution, +{ + /// Generates a random frontier of a Merkle tree having the specified nonzero size. + pub fn random_of_size(rng: &mut R, tree_size: NonZeroU64) -> Self { + let position = (u64::from(tree_size) - 1).into(); + NonEmptyFrontier::from_parts( + position, + rng.gen(), + std::iter::repeat_with(|| rng.gen()) + .take(position.past_ommer_count().into()) + .collect(), + ) + .unwrap() + } + + pub fn random_with_prior_subtree_roots( + rng: &mut R, + tree_size: NonZeroU64, + subtree_depth: NonZeroU8, + ) -> (Vec, Self) { + let prior_subtree_count: u64 = + u64::from(tree_size) / 2u64.pow(u8::from(subtree_depth).into()); + if prior_subtree_count > 0 { + let prior_roots: Vec = std::iter::repeat_with(|| rng.gen()) + .take(prior_subtree_count as usize) + .collect(); + + let subtree_root_level = Level::from(u8::from(subtree_depth)); + + // Generate replacement ommers for the random frontier from the prior subtree roots. + let mut replacement_ommers: Vec<(Level, H)> = vec![]; + let mut roots_iter = prior_roots.iter(); + loop { + if let Some(top) = replacement_ommers.pop() { + if let Some(prev) = replacement_ommers.pop() { + if top.0 == prev.0 { + // Combine, then continue the outer loop so that we eagerly combine as + // many values from the stack as we can before pushing more on. + replacement_ommers + .push((top.0 + 1, H::combine(top.0, &prev.1, &top.1))); + continue; + } else { + // We can't combine yet, so push `prev` back on. `top` will get pushed + // back on or consumed below. + replacement_ommers.push(prev); + } + } + + if let Some(root) = roots_iter.next() { + if top.0 == subtree_root_level { + replacement_ommers.push(( + subtree_root_level + 1, + H::combine(subtree_root_level, &top.1, root), + )); + } else { + replacement_ommers.push(top); + replacement_ommers.push((subtree_root_level, root.clone())); + } + } else { + // No more roots, so we just push `top` back on and break. + replacement_ommers.push(top); + break; + } + } else if let Some(root) = roots_iter.next() { + replacement_ommers.push((subtree_root_level, root.clone())); + } else { + break; + } + } + + let mut result = Self::random_of_size(rng, tree_size); + let olen = result.ommers.len(); + for (idx, (_, ommer)) in replacement_ommers.into_iter().enumerate() { + result.ommers[olen - (idx + 1)] = ommer; + } + + (prior_roots, result) + } else { + (vec![], Self::random_of_size(rng, tree_size)) + } + } +} + /// A possibly-empty Merkle frontier. #[derive(Debug, Clone, PartialEq, Eq)] pub struct Frontier { @@ -300,20 +387,30 @@ where { /// Generates a random frontier of a Merkle tree having the specified size. pub fn random_of_size(rng: &mut R, tree_size: u64) -> Self { - if tree_size == 0 { - Frontier::empty() - } else { - let position = (tree_size - 1).into(); - Frontier::from_parts( - position, - rng.gen(), - std::iter::repeat_with(|| rng.gen()) - .take(position.past_ommer_count().into()) - .collect(), - ) - .unwrap() + assert!(tree_size <= 2u64.checked_pow(DEPTH.into()).unwrap()); + Frontier { + frontier: NonZeroU64::new(tree_size) + .map(|sz| NonEmptyFrontier::random_of_size(rng, sz)), } } + + pub fn random_with_prior_subtree_roots( + rng: &mut R, + tree_size: u64, + subtree_depth: NonZeroU8, + ) -> (Vec, Self) { + assert!(tree_size <= 2u64.checked_pow(DEPTH.into()).unwrap()); + NonZeroU64::new(tree_size).map_or((vec![], Frontier::empty()), |tree_size| { + let (prior_roots, frontier) = + NonEmptyFrontier::random_with_prior_subtree_roots(rng, tree_size, subtree_depth); + ( + prior_roots, + Frontier { + frontier: Some(frontier), + }, + ) + }) + } } #[cfg(feature = "legacy-api")] @@ -577,11 +674,12 @@ impl CommitmentTree { } } -#[cfg(feature = "test-dependencies")] +#[cfg(any(test, feature = "test-dependencies"))] pub mod testing { use core::fmt::Debug; use proptest::collection::vec; use proptest::prelude::*; + use rand::{distributions::Standard, prelude::Distribution}; use std::collections::hash_map::DefaultHasher; use std::hash::Hasher; @@ -616,6 +714,12 @@ pub mod testing { } } + impl Distribution for Standard { + fn sample(&self, rng: &mut R) -> TestNode { + TestNode(rng.gen()) + } + } + pub fn arb_test_node() -> impl Strategy + Clone { any::().prop_map(TestNode) } @@ -662,11 +766,14 @@ pub mod testing { #[cfg(test)] mod tests { - use super::*; + use rand::SeedableRng; + use rand_chacha::ChaChaRng; + + use super::{testing::TestNode, *}; #[cfg(feature = "legacy-api")] use { - super::testing::{arb_commitment_tree, arb_test_node, TestNode}, + super::testing::{arb_commitment_tree, arb_test_node}, proptest::prelude::*, }; @@ -792,6 +899,34 @@ mod tests { assert_eq!(frontier, frontier0); } + #[test] + fn test_random_frontier_structure() { + let tree_size = (2u64.pow(4)) * 3 + 5; + + let mut f: Frontier = Frontier::empty(); + for i in 0..tree_size { + f.append(TestNode(i)); + } + let f = f.frontier.expect("Frontier should not be empty."); + + let mut rng = ChaChaRng::seed_from_u64(0); + let (prior_roots, f0) = Frontier::::random_with_prior_subtree_roots( + &mut rng, + tree_size, + NonZeroU8::new(4).unwrap(), + ); + let f0 = f0.frontier.expect("Frontier should not be empty."); + + assert_eq!(prior_roots.len(), 3); + assert_eq!(f.position, f0.position); + assert_eq!(f.ommers.len(), f0.ommers.len()); + + let expected_largest_ommer = + TestNode::combine(Level::from(4), &prior_roots[0], &prior_roots[1]); + assert_eq!(f0.ommers[f0.ommers.len() - 1], expected_largest_ommer); + assert_eq!(f0.ommers[f0.ommers.len() - 2], prior_roots[2]); + } + #[cfg(feature = "legacy-api")] proptest! { #[test] diff --git a/incrementalmerkletree/src/lib.rs b/incrementalmerkletree/src/lib.rs index eab0b91c..2391dbad 100644 --- a/incrementalmerkletree/src/lib.rs +++ b/incrementalmerkletree/src/lib.rs @@ -58,7 +58,7 @@ pub mod frontier; #[cfg_attr(docsrs, doc(cfg(feature = "legacy-api")))] pub mod witness; -#[cfg(feature = "test-dependencies")] +#[cfg(any(test, feature = "test-dependencies"))] #[cfg_attr(docsrs, doc(cfg(feature = "test-dependencies")))] pub mod testing; @@ -611,8 +611,14 @@ impl MerklePath { pub trait Hashable: fmt::Debug { fn empty_leaf() -> Self; + /// Combines two provided nodes that both exist at the specified level of the tree, + /// producing a new node at level `level + 1`. fn combine(level: Level, a: &Self, b: &Self) -> Self; + /// Produces an empty root at the specified level of the tree by combining empty leaf values. + /// + /// At each successive level, the value is produced by combining the value at the level below + /// with a copy of itself. fn empty_root(level: Level) -> Self where Self: Sized,