diff --git a/zcash_primitives/src/merkle_tree.rs b/zcash_primitives/src/merkle_tree.rs index 24590c097..e0b5861fa 100644 --- a/zcash_primitives/src/merkle_tree.rs +++ b/zcash_primitives/src/merkle_tree.rs @@ -1,13 +1,17 @@ //! Implementation of a Merkle tree of commitments used to prove the existence of notes. use byteorder::{LittleEndian, ReadBytesExt}; -use incrementalmerkletree::{self, bridgetree, Altitude}; +use incrementalmerkletree::{ + self, + bridgetree::{self, Leaf}, + Altitude, +}; use std::collections::VecDeque; use std::convert::TryFrom; use std::io::{self, Read, Write}; use zcash_encoding::{Optional, Vector}; -use crate::sapling::{SAPLING_COMMITMENT_TREE_DEPTH, SAPLING_COMMITMENT_TREE_DEPTH_U8}; +use crate::sapling::SAPLING_COMMITMENT_TREE_DEPTH; pub mod incremental; @@ -100,7 +104,7 @@ impl PathFiller { /// /// The depth of the Merkle tree is fixed at 32, equal to the depth of the Sapling /// commitment tree. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct CommitmentTree { pub(crate) left: Option, pub(crate) right: Option, @@ -117,7 +121,35 @@ impl CommitmentTree { } } - pub fn to_frontier(&self) -> bridgetree::Frontier + pub fn from_frontier(frontier: &bridgetree::Frontier) -> Self + where + Node: Clone, + { + frontier.value().map_or_else(Self::empty, |f| { + let (left, right) = match f.leaf() { + Leaf::Left(v) => (Some(v.clone()), None), + Leaf::Right(l, r) => (Some(l.clone()), Some(r.clone())), + }; + let mut ommers_iter = f.ommers().iter().cloned(); + let upos: usize = f.position().into(); + Self { + left, + right, + parents: (1..DEPTH) + .into_iter() + .map(|i| { + if upos & (1 << i) == 0 { + None + } else { + ommers_iter.next() + } + }) + .collect(), + } + }) + } + + pub fn to_frontier(&self) -> bridgetree::Frontier where Node: incrementalmerkletree::Hashable + Clone, { @@ -576,12 +608,18 @@ impl MerklePath { #[cfg(test)] mod tests { - use super::{CommitmentTree, Hashable, IncrementalWitness, MerklePath, PathFiller}; - use crate::sapling::Node; - + use incrementalmerkletree::bridgetree::Frontier; + use proptest::prelude::*; use std::convert::TryInto; use std::io::{self, Read, Write}; + use crate::sapling::{testing::arb_node, Node}; + + use super::{ + testing::arb_commitment_tree, CommitmentTree, Hashable, IncrementalWitness, MerklePath, + PathFiller, + }; + const HEX_EMPTY_ROOTS: [&str; 33] = [ "0100000000000000000000000000000000000000000000000000000000000000", "817de36ab2d57feb077634bca77819c8e0bd298c04f6fed0e6a83cc1356ca155", @@ -1137,6 +1175,17 @@ mod tests { assert!(witness.append(node).is_err()); } } + + proptest! { + #[test] + fn prop_commitment_tree_roundtrip(ct in arb_commitment_tree(32, arb_node(), 8)) { + let frontier: Frontier = ct.to_frontier(); + let ct0 = CommitmentTree::from_frontier(&frontier); + assert_eq!(ct, ct0); + let frontier0: Frontier = ct0.to_frontier(); + assert_eq!(frontier, frontier0); + } + } } #[cfg(any(test, feature = "test-dependencies"))] @@ -1150,12 +1199,15 @@ pub mod testing { pub fn arb_commitment_tree>( min_size: usize, arb_node: T, + depth: u8, ) -> impl Strategy> { - vec(arb_node, min_size..(min_size + 100)).prop_map(|v| { + assert!((1 << depth) >= min_size + 100); + vec(arb_node, min_size..(min_size + 100)).prop_map(move |v| { let mut tree = CommitmentTree::empty(); for node in v.into_iter() { tree.append(node).unwrap(); } + tree.parents.resize_with((depth - 1).into(), || None); tree }) } diff --git a/zcash_primitives/src/merkle_tree/incremental.rs b/zcash_primitives/src/merkle_tree/incremental.rs index c6cc09d58..6bc397ff5 100644 --- a/zcash_primitives/src/merkle_tree/incremental.rs +++ b/zcash_primitives/src/merkle_tree/incremental.rs @@ -226,7 +226,7 @@ mod tests { proptest! { #[test] - fn frontier_serialization_v0(t in arb_commitment_tree(0, sapling::arb_node())) + fn frontier_serialization_v0(t in arb_commitment_tree(0, sapling::arb_node(), 32)) { let mut buffer = vec![]; t.write(&mut buffer).unwrap(); @@ -237,7 +237,7 @@ mod tests { } #[test] - fn frontier_serialization_v1(t in arb_commitment_tree(1, sapling::arb_node())) + fn frontier_serialization_v1(t in arb_commitment_tree(1, sapling::arb_node(), 32)) { let original: Frontier = t.to_frontier(); diff --git a/zcash_primitives/src/transaction/components/sapling/builder.rs b/zcash_primitives/src/transaction/components/sapling/builder.rs index c56dd0f02..bbdd0cec8 100644 --- a/zcash_primitives/src/transaction/components/sapling/builder.rs +++ b/zcash_primitives/src/transaction/components/sapling/builder.rs @@ -581,7 +581,7 @@ pub mod testing { n_notes ), commitment_trees in vec( - arb_commitment_tree(n_notes, arb_node()).prop_map( + arb_commitment_tree(n_notes, arb_node(), 32).prop_map( |t| IncrementalWitness::from_tree(&t).path().unwrap() ), n_notes