zec
/
incrementalmerkletree
mirror of https://github.com/zcash/incrementalmerkletree.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #64 from nuttycom/typed_merkle_path 

Use `MerklePath` when returning witnesses.
This commit is contained in:
Kris Nuttycombe 2023-04-11 09:40:27 -06:00 committed by GitHub
parent be81d67cef 69d444c439
commit ea1686e8f8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 84 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
incrementalmerkletree/src/lib.rs
View File

@ -464,6 +464,53 @@ impl<'a> From<&'a Address> for Option<Position> {
}
}
/// A path from a position in a particular commitment tree to the root of that tree.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MerklePath<H, const DEPTH: u8> {
path_elems: Vec<H>,
position: Position,
}
impl<H, const DEPTH: u8> MerklePath<H, DEPTH> {
/// Constructs a Merkle path directly from a path and position.
#[allow(clippy::result_unit_err)]
pub fn from_parts(path_elems: Vec<H>, position: Position) -> Result<Self, ()> {
if path_elems.len() == usize::from(DEPTH) {
Ok(MerklePath {
path_elems,
position,
})
} else {
Err(())
}
}
pub fn path_elems(&self) -> &[H] {
&self.path_elems
}
pub fn position(&self) -> Position {
self.position
}
}
impl<H: Hashable, const DEPTH: u8> MerklePath<H, DEPTH> {
/// Returns the root of the tree corresponding to this path applied to `leaf`.
pub fn root(&self, leaf: H) -> H {
self.path_elems
.iter()
.enumerate()
.fold(leaf, |root, (i, h)| {
let level = Level(i as u8);
if (self.position.0 >> i) & 0x1 == 0 {
H::combine(level, &root, h)
} else {
H::combine(level, h, &root)
}
})
}
}
/// A trait describing the operations that make a type suitable for use as
/// a leaf or node value in a merkle tree.
pub trait Hashable: Sized + core::fmt::Debug {
@ -480,6 +527,8 @@ pub trait Hashable: Sized + core::fmt::Debug {
#[cfg(test)]
pub(crate) mod tests {
use crate::MerklePath;
use super::{Address, Level, Position, Source};
use core::ops::Range;
use either::Either;
@ -663,4 +712,23 @@ pub(crate) mod tests {
Either::Left(Address::from_parts(Level(5), 1))
);
}
#[test]
fn merkle_path_root() {
let path: MerklePath<String, 3> = MerklePath::from_parts(
vec!["a".to_string(), "cd".to_string(), "efgh".to_string()],
Position(1),
)
.unwrap();
assert_eq!(path.root("b".to_string()), "abcdefgh".to_string());
let path: MerklePath<String, 3> = MerklePath::from_parts(
vec!["d".to_string(), "ab".to_string(), "efgh".to_string()],
Position(2),
)
.unwrap();
assert_eq!(path.root("c".to_string()), "abcdefgh".to_string());
}
}

67
incrementalmerkletree/src/witness.rs
View File

@ -3,57 +3,9 @@ use std::iter::repeat;
use crate::{
frontier::{CommitmentTree, PathFiller},
Hashable, Level,
Hashable, Level, MerklePath, Position,
};
/// A path from a position in a particular commitment tree to the root of that tree.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MerklePath<H, const DEPTH: u8> {
auth_path: Vec<(H, bool)>,
position: u64,
}
impl<H, const DEPTH: u8> MerklePath<H, DEPTH> {
/// Constructs a Merkle path directly from a path and position.
#[allow(clippy::result_unit_err)]
pub fn from_parts(auth_path: Vec<(H, bool)>, position: u64) -> Result<Self, ()> {
if auth_path.len() == usize::from(DEPTH) {
Ok(MerklePath {
auth_path,
position,
})
} else {
Err(())
}
}
pub fn auth_path(&self) -> &[(H, bool)] {
&self.auth_path
}
pub fn position(&self) -> u64 {
self.position
}
}
impl<H: Hashable, const DEPTH: u8> MerklePath<H, DEPTH> {
/// 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
@ -66,6 +18,7 @@ impl<H: Hashable, const DEPTH: u8> MerklePath<H, DEPTH> {
/// use incrementalmerkletree::{
/// frontier::{CommitmentTree, testing::TestNode},
/// witness::IncrementalWitness,
/// Position
/// };
///
/// let mut tree = CommitmentTree::<TestNode, 8>::empty();
@ -73,7 +26,7 @@ impl<H: Hashable, const DEPTH: u8> MerklePath<H, DEPTH> {
/// tree.append(TestNode(0));
/// tree.append(TestNode(1));
/// let mut witness = IncrementalWitness::from_tree(tree.clone());
/// assert_eq!(witness.position(), 1);
/// assert_eq!(witness.position(), Position::from(1));
/// assert_eq!(tree.root(), witness.root());
///
/// let next = TestNode(2);
@ -131,8 +84,8 @@ impl<H, const DEPTH: u8> IncrementalWitness<H, DEPTH> {
}
/// Returns the position of the witnessed leaf node in the commitment tree.
pub fn position(&self) -> usize {
self.tree.size() - 1
pub fn position(&self) -> Position {
Position::from(self.tree.size() - 1)
}
/// Finds the next "depth" of an unfilled subtree.
@ -233,9 +186,9 @@ impl<H: Hashable + Clone, const DEPTH: u8> IncrementalWitness<H, DEPTH> {
if let Some(node) = &self.tree.left {
if self.tree.right.is_some() {
auth_path.push((node.clone(), true));
auth_path.push(node.clone());
} else {
auth_path.push((filler.next(0.into()), false));
auth_path.push(filler.next(0.into()));
}
} else {
// Can't create an authentication path for the beginning of the tree
@ -251,13 +204,13 @@ impl<H: Hashable + Clone, const DEPTH: u8> IncrementalWitness<H, DEPTH> {
.enumerate()
{
auth_path.push(match p {
Some(node) => (node.clone(), true),
None => (filler.next(Level::from((i + 1) as u8)), false),
Some(node) => node.clone(),
None => filler.next(Level::from((i + 1) as u8)),
});
}
assert_eq!(auth_path.len(), usize::from(depth));
MerklePath::from_parts(auth_path, self.position() as u64).ok()
MerklePath::from_parts(auth_path, self.position()).ok()
}
}

10
shardtree/src/lib.rs
View File

@ -7,7 +7,7 @@ use either::Either;
use std::collections::{BTreeMap, BTreeSet};
use std::rc::Rc;
use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
use incrementalmerkletree::{Address, Hashable, Level, MerklePath, Position, Retention};
bitflags! {
pub struct RetentionFlags: u8 {
@ -2078,7 +2078,7 @@ impl<
&self,
position: Position,
checkpoint_depth: usize,
) -> Result<Vec<H>, QueryError> {
) -> Result<MerklePath<H, DEPTH>, QueryError> {
let max_leaf_position = self
.max_leaf_position(checkpoint_depth)
.and_then(|v| v.ok_or_else(|| QueryError::TreeIncomplete(vec![Self::root_addr()])))?;
@ -2105,7 +2105,7 @@ impl<
cur_addr = cur_addr.parent();
}
Ok(witness)
Ok(MerklePath::from_parts(witness, position).unwrap())
}
}
@ -2652,7 +2652,9 @@ mod tests {
}
fn witness(&self, position: Position, checkpoint_depth: usize) -> Option<Vec<H>> {
ShardTree::witness(self, position, checkpoint_depth).ok()
ShardTree::witness(self, position, checkpoint_depth)
.ok()
.map(|p| p.path_elems().to_vec())
}
fn remove_mark(&mut self, position: Position) -> bool {