Merge pull request #64 from nuttycom/typed_merkle_path

Use `MerklePath` when returning witnesses.

incrementalmerkletree/src/lib.rs

@@ -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());
+    }
 }

incrementalmerkletree/src/witness.rs

@@ -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 {
+    pub fn position(&self) -> Position {
-        self.tree.size() - 1
+        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),
+                Some(node) => node.clone(),
-                None => (filler.next(Level::from((i + 1) as u8)), false),
+                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()
     }
 }

shardtree/src/lib.rs

@@ -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 {