diff --git a/incrementalmerkletree/CHANGELOG.md b/incrementalmerkletree/CHANGELOG.md index 39135aa..08281ff 100644 --- a/incrementalmerkletree/CHANGELOG.md +++ b/incrementalmerkletree/CHANGELOG.md @@ -11,7 +11,12 @@ and this project adheres to Rust's notion of - `incrementalmerkletree::frontier` Types that model the state at the rightmost node of a Merkle tree that is filled sequentially from the left. These have been migrated here from the `bridgetree` crate as they are useful outside - of the context of the `bridgetree` data structures. + of the context of the `bridgetree` data structures. Additional legacy types + used for this modeling have been moved here from the `librustzcash` crate; + these migrated types are available under a `legacy-api` feature flag. +- `incrementalmerkletree::witness` Types migrated from `librustzcash` under + the `legacy-api` feature flag related to constructing witnesses for leaves + of a Merkle tree. ## [0.3.1] - 2023-02-28 diff --git a/incrementalmerkletree/Cargo.toml b/incrementalmerkletree/Cargo.toml index 92c709f..45a7c39 100644 --- a/incrementalmerkletree/Cargo.toml +++ b/incrementalmerkletree/Cargo.toml @@ -21,4 +21,5 @@ proptest = { version = "1.0.0", optional = true } proptest = "1.0.0" [features] +legacy-api = [] test-dependencies = ["proptest"] diff --git a/incrementalmerkletree/src/frontier.rs b/incrementalmerkletree/src/frontier.rs index 6d64a48..e60e21e 100644 --- a/incrementalmerkletree/src/frontier.rs +++ b/incrementalmerkletree/src/frontier.rs @@ -3,6 +3,9 @@ use std::mem::size_of; use crate::{Address, Hashable, Level, Position, Source}; +#[cfg(feature = "legacy-api")] +use {std::collections::VecDeque, std::iter::repeat}; + /// Validation errors that can occur during reconstruction of a Merkle frontier from /// its constituent parts. #[derive(Clone, Debug, PartialEq, Eq)] @@ -244,9 +247,227 @@ impl Frontier { } } +#[cfg(feature = "legacy-api")] +pub(crate) struct PathFiller { + pub(crate) queue: VecDeque, +} + +#[cfg(feature = "legacy-api")] +impl PathFiller { + pub(crate) fn empty() -> Self { + PathFiller { + queue: VecDeque::new(), + } + } + + pub(crate) fn next(&mut self, level: Level) -> H { + self.queue + .pop_front() + .unwrap_or_else(|| H::empty_root(level)) + } +} + +/// A Merkle tree of note commitments. +#[derive(Clone, Debug, PartialEq, Eq)] +#[cfg(feature = "legacy-api")] +pub struct CommitmentTree { + pub(crate) left: Option, + pub(crate) right: Option, + pub(crate) parents: Vec>, +} + +#[cfg(feature = "legacy-api")] +impl CommitmentTree { + /// Creates an empty tree. + pub fn empty() -> Self { + CommitmentTree { + left: None, + right: None, + parents: vec![], + } + } + + pub fn from_frontier(frontier: &Frontier) -> Self + where + H: Clone, + { + frontier.value().map_or_else(Self::empty, |f| { + let mut ommers_iter = f.ommers().iter().cloned(); + let (left, right) = if f.position().is_odd() { + ( + ommers_iter + .next() + .expect("An ommer must exist if the frontier position is odd"), + Some(f.leaf().clone()), + ) + } else { + (f.leaf().clone(), None) + }; + + let upos: usize = f.position().into(); + Self { + left: Some(left), + right, + parents: (1u8..DEPTH) + .into_iter() + .map(|i| { + if upos & (1 << i) == 0 { + None + } else { + ommers_iter.next() + } + }) + .collect(), + } + }) + } + + pub fn to_frontier(&self) -> Frontier + where + H: Hashable + Clone, + { + if self.size() == 0 { + Frontier::empty() + } else { + let ommers_iter = self.parents.iter().filter_map(|v| v.as_ref()).cloned(); + let (leaf, ommers) = match (self.left.as_ref(), self.right.as_ref()) { + (Some(a), None) => (a.clone(), ommers_iter.collect()), + (Some(a), Some(b)) => ( + b.clone(), + Some(a.clone()).into_iter().chain(ommers_iter).collect(), + ), + _ => unreachable!(), + }; + + // If a frontier cannot be successfully constructed from the + // parts of a commitment tree, it is a programming error. + Frontier::from_parts((self.size() - 1).into(), leaf, ommers) + .expect("Frontier should be constructable from CommitmentTree.") + } + } + + /// Returns the number of leaf nodes in the tree. + pub fn size(&self) -> usize { + self.parents.iter().enumerate().fold( + match (self.left.as_ref(), self.right.as_ref()) { + (None, None) => 0, + (Some(_), None) => 1, + (Some(_), Some(_)) => 2, + (None, Some(_)) => unreachable!(), + }, + |acc, (i, p)| { + // Treat occupation of parents array as a binary number + // (right-shifted by 1) + acc + if p.is_some() { 1 << (i + 1) } else { 0 } + }, + ) + } + + pub(crate) fn is_complete(&self, depth: u8) -> bool { + if depth == 0 { + self.left.is_some() && self.right.is_none() && self.parents.is_empty() + } else { + self.left.is_some() + && self.right.is_some() + && self + .parents + .iter() + .chain(repeat(&None)) + .take((depth - 1).into()) + .all(|p| p.is_some()) + } + } +} + +#[cfg(feature = "legacy-api")] +impl CommitmentTree { + /// Adds a leaf node to the tree. + /// + /// Returns an error if the tree is full. + #[allow(clippy::result_unit_err)] + pub fn append(&mut self, node: H) -> Result<(), ()> { + if self.is_complete(DEPTH) { + // Tree is full + return Err(()); + } + + match (&self.left, &self.right) { + (None, _) => self.left = Some(node), + (_, None) => self.right = Some(node), + (Some(l), Some(r)) => { + let mut combined = H::combine(0.into(), l, r); + self.left = Some(node); + self.right = None; + + for i in 0..DEPTH { + let i_usize = usize::from(i); + if i_usize < self.parents.len() { + if let Some(p) = &self.parents[i_usize] { + combined = H::combine((i + 1).into(), p, &combined); + self.parents[i_usize] = None; + } else { + self.parents[i_usize] = Some(combined); + break; + } + } else { + self.parents.push(Some(combined)); + break; + } + } + } + } + + Ok(()) + } + + /// Returns the current root of the tree. + pub fn root(&self) -> H { + self.root_inner(DEPTH, PathFiller::empty()) + } + + pub(crate) fn root_inner(&self, depth: u8, mut filler: PathFiller) -> H { + assert!(depth > 0); + + // 1) Hash left and right leaves together. + // - Empty leaves are used as needed. + // - Note that `filler.next` is side-effecting and so cannot be factored out. + let leaf_root = H::combine( + 0.into(), + &self + .left + .as_ref() + .map_or_else(|| filler.next(0.into()), |n| n.clone()), + &self + .right + .as_ref() + .map_or_else(|| filler.next(0.into()), |n| n.clone()), + ); + + // 2) Extend the parents to the desired depth with None values, then hash from leaf to + // root. Roots of the empty subtrees are used as needed. + self.parents + .iter() + .chain(repeat(&None)) + .take((depth - 1).into()) + .enumerate() + .fold(leaf_root, |root, (i, p)| { + let level = Level::from(i as u8 + 1); + match p { + Some(node) => H::combine(level, node, &root), + None => H::combine(level, &root, &filler.next(level)), + } + }) + } +} + #[cfg(feature = "test-dependencies")] pub mod testing { - use crate::Hashable; + use core::fmt::Debug; + use proptest::prelude::*; + use std::collections::hash_map::DefaultHasher; + use std::hash::Hasher; + + use crate::{Hashable, Level}; impl crate::testing::Frontier for super::Frontier @@ -259,12 +480,65 @@ pub mod testing { super::Frontier::root(self) } } + + #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] + pub struct TestNode(pub u64); + + impl Hashable for TestNode { + fn empty_leaf() -> Self { + Self(0) + } + + fn combine(level: Level, a: &Self, b: &Self) -> Self { + let mut hasher = DefaultHasher::new(); + hasher.write_u8(level.into()); + hasher.write_u64(a.0); + hasher.write_u64(b.0); + Self(hasher.finish()) + } + } + + prop_compose! { + pub fn arb_test_node()(i in any::()) -> TestNode { + TestNode(i) + } + } + + #[cfg(feature = "legacy-api")] + use {crate::frontier::CommitmentTree, proptest::collection::vec}; + + #[cfg(feature = "legacy-api")] + pub fn arb_commitment_tree< + H: Hashable + Clone + Debug, + T: Strategy, + const DEPTH: u8, + >( + min_size: usize, + arb_node: T, + ) -> impl Strategy> { + 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 + }) + } } #[cfg(test)] mod tests { + use super::*; + #[cfg(feature = "legacy-api")] + use { + super::testing::{arb_commitment_tree, arb_test_node, TestNode}, + proptest::prelude::*, + }; + #[test] fn nonempty_frontier_root() { let mut frontier = NonEmptyFrontier::new("a".to_string()); @@ -319,4 +593,73 @@ mod tests { frontier.witness(4, bridge_value_at) ); } + + #[test] + #[cfg(feature = "legacy-api")] + fn test_commitment_tree_complete() { + let mut t: CommitmentTree = CommitmentTree::empty(); + for n in 1u64..=32 { + t.append(TestNode(n)).unwrap(); + // every tree of a power-of-two height is complete + let is_complete = n.count_ones() == 1; + let level = usize::BITS - 1 - n.leading_zeros(); //log2 + assert_eq!( + is_complete, + t.is_complete(level.try_into().unwrap()), + "Tree {:?} {} complete at height {}", + t, + if is_complete { + "should be" + } else { + "should not be" + }, + n + ); + } + } + + #[test] + #[cfg(feature = "legacy-api")] + fn test_commitment_tree_roundtrip() { + let ct = CommitmentTree { + left: Some("a".to_string()), + right: Some("b".to_string()), + parents: vec![ + Some("c".to_string()), + Some("d".to_string()), + Some("e".to_string()), + Some("f".to_string()), + None, + None, + None, + ], + }; + + 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(feature = "legacy-api")] + proptest! { + #[test] + fn prop_commitment_tree_roundtrip(ct in arb_commitment_tree(32, arb_test_node())) { + 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); + } + + #[test] + fn prop_commitment_tree_roundtrip_str(ct in arb_commitment_tree::<_, _, 8>(32, any::().prop_map(|c| c.to_string()))) { + 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); + } + } } diff --git a/incrementalmerkletree/src/lib.rs b/incrementalmerkletree/src/lib.rs index 31f05b3..c88ddfc 100644 --- a/incrementalmerkletree/src/lib.rs +++ b/incrementalmerkletree/src/lib.rs @@ -8,6 +8,9 @@ use std::ops::{Add, AddAssign, Range, Sub}; pub mod frontier; +#[cfg(feature = "legacy-api")] +pub mod witness; + #[cfg(feature = "test-dependencies")] pub mod testing; @@ -219,6 +222,13 @@ impl From for usize { } } +impl TryFrom for Level { + type Error = TryFromIntError; + fn try_from(sz: usize) -> Result { + ::try_from(sz).map(Self) + } +} + impl Sub for Level { type Output = Self; fn sub(self, value: u8) -> Self { diff --git a/incrementalmerkletree/src/witness.rs b/incrementalmerkletree/src/witness.rs new file mode 100644 index 0000000..97147b4 --- /dev/null +++ b/incrementalmerkletree/src/witness.rs @@ -0,0 +1,239 @@ +use std::convert::TryInto; +use std::iter::repeat; + +use crate::{ + frontier::{CommitmentTree, PathFiller}, + Hashable, Level, +}; + +/// A path from a position in a particular commitment tree to the root of that tree. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MerklePath { + auth_path: Vec<(H, bool)>, + position: u64, +} + +impl MerklePath { + pub fn auth_path(&self) -> &[(H, bool)] { + &self.auth_path + } + + pub fn position(&self) -> u64 { + self.position + } + + /// Constructs a Merkle path directly from a path and position. + pub fn from_path(auth_path: Vec<(H, bool)>, position: u64) -> Result { + if auth_path.len() == usize::from(DEPTH) { + Ok(MerklePath { + auth_path, + position, + }) + } else { + Err(()) + } + } +} + +impl MerklePath { + /// Returns the root of the tree corresponding to this path applied to `leaf`. + pub fn root(&self, leaf: H) -> H { + self.auth_path + .iter() + .enumerate() + .fold(leaf, |root, (i, (p, leaf_is_on_right))| { + let level = u8::try_from(i) + .expect("Parents list length may not exceed what is representable by an u8") + .into(); + match leaf_is_on_right { + false => H::combine(level, &root, p), + true => H::combine(level, p, &root), + } + }) + } +} + +/// An updatable witness to a path from a position in a particular [`CommitmentTree`]. +/// +/// Appending the same commitments in the same order to both the original +/// [`CommitmentTree`] and this `IncrementalWitness` will result in a witness to the path +/// from the target position to the root of the updated tree. +/// +/// # Examples +/// +/// ``` +/// use incrementalmerkletree::{ +/// frontier::{CommitmentTree, testing::TestNode}, +/// witness::IncrementalWitness, +/// }; +/// +/// let mut tree = CommitmentTree::::empty(); +/// +/// tree.append(TestNode(0)); +/// tree.append(TestNode(1)); +/// let mut witness = IncrementalWitness::from_tree(tree.clone()); +/// assert_eq!(witness.position(), 1); +/// assert_eq!(tree.root(), witness.root()); +/// +/// let next = TestNode(2); +/// tree.append(next.clone()); +/// witness.append(next); +/// assert_eq!(tree.root(), witness.root()); +/// ``` +#[derive(Clone, Debug)] +pub struct IncrementalWitness { + tree: CommitmentTree, + filled: Vec, + cursor_depth: u8, + cursor: Option>, +} + +impl IncrementalWitness { + /// Creates an `IncrementalWitness` for the most recent commitment added to the given + /// [`CommitmentTree`]. + pub fn from_tree(tree: CommitmentTree) -> IncrementalWitness { + IncrementalWitness { + tree, + filled: vec![], + cursor_depth: 0, + cursor: None, + } + } + + /// Returns the position of the witnessed leaf node in the commitment tree. + pub fn position(&self) -> usize { + self.tree.size() - 1 + } + + /// Finds the next "depth" of an unfilled subtree. + fn next_depth(&self) -> u8 { + let mut skip: u8 = self + .filled + .len() + .try_into() + .expect("Merkle tree depths may not exceed the bounds of a u8"); + + if self.tree.left.is_none() { + if skip > 0 { + skip -= 1; + } else { + return 0; + } + } + + if self.tree.right.is_none() { + if skip > 0 { + skip -= 1; + } else { + return 0; + } + } + + let mut d = 1; + for p in &self.tree.parents { + if p.is_none() { + if skip > 0 { + skip -= 1; + } else { + return d; + } + } + d += 1; + } + + d + skip + } +} + +impl IncrementalWitness { + fn filler(&self) -> PathFiller { + let cursor_root = self + .cursor + .as_ref() + .map(|c| c.root_inner(self.cursor_depth, PathFiller::empty())); + + PathFiller { + queue: self.filled.iter().cloned().chain(cursor_root).collect(), + } + } + + /// Tracks a leaf node that has been added to the underlying tree. + /// + /// Returns an error if the tree is full. + #[allow(clippy::result_unit_err)] + pub fn append(&mut self, node: H) -> Result<(), ()> { + if let Some(mut cursor) = self.cursor.take() { + cursor.append(node).expect("cursor should not be full"); + if cursor.is_complete(self.cursor_depth) { + self.filled + .push(cursor.root_inner(self.cursor_depth, PathFiller::empty())); + } else { + self.cursor = Some(cursor); + } + } else { + self.cursor_depth = self.next_depth(); + if self.cursor_depth >= DEPTH { + // Tree is full + return Err(()); + } + + if self.cursor_depth == 0 { + self.filled.push(node); + } else { + let mut cursor = CommitmentTree::empty(); + cursor.append(node).expect("cursor should not be full"); + self.cursor = Some(cursor); + } + } + + Ok(()) + } + + /// Returns the current root of the tree corresponding to the witness. + pub fn root(&self) -> H { + self.root_inner(DEPTH) + } + + fn root_inner(&self, depth: u8) -> H { + self.tree.root_inner(depth, self.filler()) + } + + /// Returns the current witness, or None if the tree is empty. + pub fn path(&self) -> Option> { + self.path_inner(DEPTH) + } + + fn path_inner(&self, depth: u8) -> Option> { + let mut filler = self.filler(); + let mut auth_path = Vec::new(); + + if let Some(node) = &self.tree.left { + if self.tree.right.is_some() { + auth_path.push((node.clone(), true)); + } else { + auth_path.push((filler.next(0.into()), false)); + } + } else { + // Can't create an authentication path for the beginning of the tree + return None; + } + + for (i, p) in self + .tree + .parents + .iter() + .chain(repeat(&None)) + .take((depth - 1).into()) + .enumerate() + { + auth_path.push(match p { + Some(node) => (node.clone(), true), + None => (filler.next(Level::from((i + 1) as u8)), false), + }); + } + + assert_eq!(auth_path.len(), usize::from(depth)); + + MerklePath::from_path(auth_path, self.position() as u64).ok() + } +}