Merge pull request #53 from 0xPolygonMiden/vlopes11-36-feat-add-merkle-types

feat: add merkle path wrapper

src/merkle/merkle_tree.rs

@@ -1,4 +1,4 @@
-use super::{Felt, MerkleError, Rpo256, RpoDigest, Vec, Word};
+use super::{Felt, MerkleError, MerklePath, Rpo256, RpoDigest, Vec, Word};
 use crate::{utils::uninit_vector, FieldElement};
 use core::slice;
 use winter_math::log2;
@@ -88,7 +88,7 @@ impl MerkleTree {
     /// Returns an error if:
     /// * The specified depth is greater than the depth of the tree.
     /// * The specified index not valid for the specified depth.
-    pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
+    pub fn get_path(&self, depth: u32, index: u64) -> Result<MerklePath, MerkleError> {
         if depth == 0 {
             return Err(MerkleError::DepthTooSmall(depth));
         } else if depth > self.depth() {
@@ -106,7 +106,7 @@ impl MerkleTree {
             pos >>= 1;
         }
 
-        Ok(path)
+        Ok(path.into())
     }
 
     /// Replaces the leaf at the specified index with the provided value.
@@ -206,14 +206,14 @@ mod tests {
         let (_, node2, node3) = compute_internal_nodes();
 
         // check depth 2
-        assert_eq!(vec![LEAVES4[1], node3], tree.get_path(2, 0).unwrap());
-        assert_eq!(vec![LEAVES4[0], node3], tree.get_path(2, 1).unwrap());
-        assert_eq!(vec![LEAVES4[3], node2], tree.get_path(2, 2).unwrap());
-        assert_eq!(vec![LEAVES4[2], node2], tree.get_path(2, 3).unwrap());
+        assert_eq!(vec![LEAVES4[1], node3], *tree.get_path(2, 0).unwrap());
+        assert_eq!(vec![LEAVES4[0], node3], *tree.get_path(2, 1).unwrap());
+        assert_eq!(vec![LEAVES4[3], node2], *tree.get_path(2, 2).unwrap());
+        assert_eq!(vec![LEAVES4[2], node2], *tree.get_path(2, 3).unwrap());
 
         // check depth 1
-        assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
-        assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
+        assert_eq!(vec![node3], *tree.get_path(1, 0).unwrap());
+        assert_eq!(vec![node2], *tree.get_path(1, 1).unwrap());
     }
 
     #[test]

src/merkle/mod.rs

@@ -1,6 +1,6 @@
 use super::{
     hash::rpo::{Rpo256, RpoDigest},
-    utils::collections::{BTreeMap, Vec},
+    utils::collections::{vec, BTreeMap, Vec},
     Felt, Word, ZERO,
 };
 use core::fmt;
@@ -8,8 +8,11 @@ use core::fmt;
 mod merkle_tree;
 pub use merkle_tree::MerkleTree;
 
-mod merkle_path_set;
-pub use merkle_path_set::MerklePathSet;
+mod path;
+pub use path::MerklePath;
+
+mod path_set;
+pub use path_set::MerklePathSet;
 
 mod simple_smt;
 pub use simple_smt::SimpleSmt;
@@ -24,7 +27,7 @@ pub enum MerkleError {
     NumLeavesNotPowerOfTwo(usize),
     InvalidIndex(u32, u64),
     InvalidDepth(u32, u32),
-    InvalidPath(Vec<Word>),
+    InvalidPath(MerklePath),
     InvalidEntriesCount(usize, usize),
     NodeNotInSet(u64),
 }

src/merkle/path.rs

@@ -0,0 +1,89 @@
+use super::{vec, Rpo256, Vec, Word};
+use core::ops::{Deref, DerefMut};
+
+// MERKLE PATH
+// ================================================================================================
+
+/// A merkle path container, composed of a sequence of nodes of a Merkle tree.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct MerklePath {
+    nodes: Vec<Word>,
+}
+
+impl MerklePath {
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Creates a new Merkle path from a list of nodes.
+    pub fn new(nodes: Vec<Word>) -> Self {
+        Self { nodes }
+    }
+
+    // PROVIDERS
+    // --------------------------------------------------------------------------------------------
+
+    /// Computes the merkle root for this opening.
+    pub fn compute_root(&self, mut index: u64, node: Word) -> Word {
+        self.nodes.iter().copied().fold(node, |node, sibling| {
+            // build the input node, considering the parity of the current index.
+            let is_right_sibling = (index & 1) == 1;
+            let input = if is_right_sibling {
+                [sibling.into(), node.into()]
+            } else {
+                [node.into(), sibling.into()]
+            };
+            // compute the node and move to the next iteration.
+            index >>= 1;
+            Rpo256::merge(&input).into()
+        })
+    }
+
+    /// Returns the depth in which this Merkle path proof is valid.
+    pub fn depth(&self) -> u8 {
+        self.nodes.len() as u8
+    }
+
+    /// Verifies the Merkle opening proof towards the provided root.
+    ///
+    /// Returns `true` if `node` exists at `index` in a Merkle tree with `root`.
+    pub fn verify(&self, index: u64, node: Word, root: &Word) -> bool {
+        root == &self.compute_root(index, node)
+    }
+}
+
+impl From<Vec<Word>> for MerklePath {
+    fn from(path: Vec<Word>) -> Self {
+        Self::new(path)
+    }
+}
+
+impl Deref for MerklePath {
+    // we use `Vec` here instead of slice so we can call vector mutation methods directly from the
+    // merkle path (example: `Vec::remove`).
+    type Target = Vec<Word>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.nodes
+    }
+}
+
+impl DerefMut for MerklePath {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.nodes
+    }
+}
+
+impl FromIterator<Word> for MerklePath {
+    fn from_iter<T: IntoIterator<Item = Word>>(iter: T) -> Self {
+        Self::new(iter.into_iter().collect())
+    }
+}
+
+impl IntoIterator for MerklePath {
+    type Item = Word;
+    type IntoIter = vec::IntoIter<Word>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.nodes.into_iter()
+    }
+}

src/merkle/merkle_path_set.rs → src/merkle/path_set.rs

@@ -1,4 +1,4 @@
-use super::{BTreeMap, MerkleError, Rpo256, Vec, Word, ZERO};
+use super::{BTreeMap, MerkleError, MerklePath, Rpo256, Vec, Word, ZERO};
 
 // MERKLE PATH SET
 // ================================================================================================
@@ -8,7 +8,7 @@ use super::{BTreeMap, MerkleError, Rpo256, Vec, Word, ZERO};
 pub struct MerklePathSet {
     root: Word,
     total_depth: u32,
-    paths: BTreeMap<u64, Vec<Word>>,
+    paths: BTreeMap<u64, MerklePath>,
 }
 
 impl MerklePathSet {
@@ -18,7 +18,7 @@ impl MerklePathSet {
     /// Returns an empty MerklePathSet.
     pub fn new(depth: u32) -> Result<Self, MerkleError> {
         let root = [ZERO; 4];
-        let paths = BTreeMap::<u64, Vec<Word>>::new();
+        let paths = BTreeMap::new();
 
         Ok(Self {
             root,
@@ -30,59 +30,15 @@ impl MerklePathSet {
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
-    /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
-    /// specify the leaf node at which the path starts.
-    ///
-    /// # Errors
-    /// Returns an error if:
-    /// - The specified index is not valid in the context of this Merkle path set (i.e., the index
-    ///   implies a greater depth than is specified for this set).
-    /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
-    ///   different root).
-    pub fn add_path(
-        &mut self,
-        index: u64,
-        value: Word,
-        path: Vec<Word>,
-    ) -> Result<(), MerkleError> {
-        let depth = (path.len() + 1) as u32;
-        if depth != self.total_depth {
-            return Err(MerkleError::InvalidDepth(self.total_depth, depth));
-        }
-
-        // Actual number of node in tree
-        let pos = 2u64.pow(self.total_depth) + index;
-
-        // Index of the leaf path in map. Paths of neighboring leaves are stored in one key-value pair
-        let half_pos = pos / 2;
-
-        let mut extended_path = path;
-        if is_even(pos) {
-            extended_path.insert(0, value);
-        } else {
-            extended_path.insert(1, value);
-        }
-
-        let root_of_current_path = compute_path_root(&extended_path, depth, index);
-        if self.root == [ZERO; 4] {
-            self.root = root_of_current_path;
-        } else if self.root != root_of_current_path {
-            return Err(MerkleError::InvalidPath(extended_path));
-        }
-        self.paths.insert(half_pos, extended_path);
-
-        Ok(())
-    }
-
     /// Returns the root to which all paths in this set resolve.
-    pub fn root(&self) -> Word {
+    pub const fn root(&self) -> Word {
         self.root
     }
 
     /// Returns the depth of the Merkle tree implied by the paths stored in this set.
     ///
     /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
-    pub fn depth(&self) -> u32 {
+    pub const fn depth(&self) -> u32 {
         self.total_depth
     }
 
@@ -106,7 +62,7 @@ impl MerklePathSet {
         match self.paths.get(&index) {
             None => Err(MerkleError::NodeNotInSet(index)),
             Some(path) => {
-                if is_even(pos) {
+                if Self::is_even(pos) {
                     Ok(path[0])
                 } else {
                     Ok(path[1])
@@ -122,7 +78,7 @@ impl MerklePathSet {
     /// Returns an error if:
     /// * The specified index not valid for the depth of structure.
     /// * Node of the requested path does not exist in the set.
-    pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
+    pub fn get_path(&self, depth: u32, index: u64) -> Result<MerklePath, MerkleError> {
         if index >= 2u64.pow(self.total_depth) {
             return Err(MerkleError::InvalidIndex(self.total_depth, index));
         }
@@ -137,7 +93,7 @@ impl MerklePathSet {
             None => Err(MerkleError::NodeNotInSet(index)),
             Some(path) => {
                 let mut local_path = path.clone();
-                if is_even(pos) {
+                if Self::is_even(pos) {
                     local_path.remove(0);
                     Ok(local_path)
                 } else {
@@ -148,6 +104,53 @@ impl MerklePathSet {
         }
     }
 
+    // STATE MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters
+    /// specify the leaf node at which the path starts.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The specified index is not valid in the context of this Merkle path set (i.e., the index
+    ///   implies a greater depth than is specified for this set).
+    /// - The specified path is not consistent with other paths in the set (i.e., resolves to a
+    ///   different root).
+    pub fn add_path(
+        &mut self,
+        index: u64,
+        value: Word,
+        path: MerklePath,
+    ) -> Result<(), MerkleError> {
+        let depth = (path.len() + 1) as u32;
+        if depth != self.total_depth {
+            return Err(MerkleError::InvalidDepth(self.total_depth, depth));
+        }
+
+        // Actual number of node in tree
+        let pos = 2u64.pow(self.total_depth) + index;
+
+        // Index of the leaf path in map. Paths of neighboring leaves are stored in one key-value pair
+        let half_pos = pos / 2;
+
+        let mut extended_path = path;
+        if Self::is_even(pos) {
+            extended_path.insert(0, value);
+        } else {
+            extended_path.insert(1, value);
+        }
+
+        let root_of_current_path = Self::compute_path_root(&extended_path, depth, index);
+        if self.root == [ZERO; 4] {
+            self.root = root_of_current_path;
+        } else if self.root != root_of_current_path {
+            return Err(MerkleError::InvalidPath(extended_path));
+        }
+        self.paths.insert(half_pos, extended_path);
+
+        Ok(())
+    }
+
     /// Replaces the leaf at the specified index with the provided value.
     ///
     /// # Errors
@@ -166,16 +169,16 @@ impl MerklePathSet {
         };
 
         // Fill old_hashes vector -----------------------------------------------------------------
-        let (old_hashes, _) = compute_path_trace(path, depth, index);
+        let (old_hashes, _) = Self::compute_path_trace(path, depth, index);
 
         // Fill new_hashes vector -----------------------------------------------------------------
-        if is_even(pos) {
+        if Self::is_even(pos) {
             path[0] = value;
         } else {
             path[1] = value;
         }
 
-        let (new_hashes, new_root) = compute_path_trace(path, depth, index);
+        let (new_hashes, new_root) = Self::compute_path_trace(path, depth, index);
         self.root = new_root;
 
         // update paths ---------------------------------------------------------------------------
@@ -190,59 +193,59 @@ impl MerklePathSet {
 
         Ok(())
     }
-}
 
-// HELPER FUNCTIONS
-// --------------------------------------------------------------------------------------------
-
-fn is_even(pos: u64) -> bool {
-    pos & 1 == 0
-}
+    // HELPER FUNCTIONS
+    // --------------------------------------------------------------------------------------------
 
-/// Calculates the hash of the parent node by two sibling ones
-/// - node — current node
-/// - node_pos — position of the current node
-/// - sibling — neighboring vertex in the tree
-fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
-    if is_even(node_pos) {
-        Rpo256::merge(&[node.into(), sibling.into()]).into()
-    } else {
-        Rpo256::merge(&[sibling.into(), node.into()]).into()
+    const fn is_even(pos: u64) -> bool {
+        pos & 1 == 0
     }
-}
-
-/// Returns vector of hashes from current to the root
-fn compute_path_trace(path: &[Word], depth: u32, index: u64) -> (Vec<Word>, Word) {
-    let mut pos = 2u64.pow(depth) + index;
 
-    let mut computed_hashes = Vec::<Word>::new();
+    /// Returns hash of the root
+    fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word {
+        let mut pos = 2u64.pow(depth) + index;
 
-    let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
+        // hash that is obtained after calculating the current hash and path hash
+        let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
 
-    if path.len() != 2 {
         for path_hash in path.iter().skip(2) {
-            computed_hashes.push(comp_hash);
             pos /= 2;
-            comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash);
+            comp_hash = Self::calculate_parent_hash(comp_hash, pos, *path_hash);
         }
+
+        comp_hash
     }
 
-    (computed_hashes, comp_hash)
-}
+    /// Calculates the hash of the parent node by two sibling ones
+    /// - node — current node
+    /// - node_pos — position of the current node
+    /// - sibling — neighboring vertex in the tree
+    fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
+        if Self::is_even(node_pos) {
+            Rpo256::merge(&[node.into(), sibling.into()]).into()
+        } else {
+            Rpo256::merge(&[sibling.into(), node.into()]).into()
+        }
+    }
 
-/// Returns hash of the root
-fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word {
-    let mut pos = 2u64.pow(depth) + index;
+    /// Returns vector of hashes from current to the root
+    fn compute_path_trace(path: &[Word], depth: u32, index: u64) -> (MerklePath, Word) {
+        let mut pos = 2u64.pow(depth) + index;
 
-    // hash that is obtained after calculating the current hash and path hash
-    let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
+        let mut computed_hashes = Vec::<Word>::new();
 
-    for path_hash in path.iter().skip(2) {
-        pos /= 2;
-        comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash);
-    }
+        let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into();
+
+        if path.len() != 2 {
+            for path_hash in path.iter().skip(2) {
+                computed_hashes.push(comp_hash);
+                pos /= 2;
+                comp_hash = Self::calculate_parent_hash(comp_hash, pos, *path_hash);
+            }
+        }
 
-    comp_hash
+        (computed_hashes.into(), comp_hash)
+    }
 }
 
 // TESTS
@@ -250,7 +253,7 @@ fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word {
 
 #[cfg(test)]
 mod tests {
-    use super::calculate_parent_hash;
+    use super::*;
     use crate::merkle::int_to_node;
 
     #[test]
@@ -260,14 +263,14 @@ mod tests {
         let leaf2 = int_to_node(2);
         let leaf3 = int_to_node(3);
 
-        let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
-        let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
+        let parent0 = MerklePathSet::calculate_parent_hash(leaf0, 0, leaf1);
+        let parent1 = MerklePathSet::calculate_parent_hash(leaf2, 2, leaf3);
 
-        let root_exp = calculate_parent_hash(parent0, 0, parent1);
+        let root_exp = MerklePathSet::calculate_parent_hash(parent0, 0, parent1);
 
         let mut set = super::MerklePathSet::new(3).unwrap();
 
-        set.add_path(0, leaf0, vec![leaf1, parent1]).unwrap();
+        set.add_path(0, leaf0, vec![leaf1, parent1].into()).unwrap();
 
         assert_eq!(set.root(), root_exp);
     }
@@ -280,10 +283,10 @@ mod tests {
         let depth = 4u32;
         let mut set = super::MerklePathSet::new(depth).unwrap();
 
-        set.add_path(index, hash_6, path_6.clone()).unwrap();
+        set.add_path(index, hash_6, path_6.clone().into()).unwrap();
         let stored_path_6 = set.get_path(depth, index).unwrap();
 
-        assert_eq!(path_6, stored_path_6);
+        assert_eq!(path_6, *stored_path_6);
         assert!(set.get_path(depth, 15u64).is_err())
     }
 
@@ -295,7 +298,7 @@ mod tests {
         let depth = 4u32;
         let mut set = super::MerklePathSet::new(depth).unwrap();
 
-        set.add_path(index, hash_6, path_6).unwrap();
+        set.add_path(index, hash_6, path_6.into()).unwrap();
 
         assert_eq!(int_to_node(6u64), set.get_node(depth, index).unwrap());
         assert!(set.get_node(depth, 15u64).is_err());
@@ -307,8 +310,8 @@ mod tests {
         let hash_5 = int_to_node(5);
         let hash_6 = int_to_node(6);
         let hash_7 = int_to_node(7);
-        let hash_45 = calculate_parent_hash(hash_4, 12u64, hash_5);
-        let hash_67 = calculate_parent_hash(hash_6, 14u64, hash_7);
+        let hash_45 = MerklePathSet::calculate_parent_hash(hash_4, 12u64, hash_5);
+        let hash_67 = MerklePathSet::calculate_parent_hash(hash_6, 14u64, hash_7);
 
         let hash_0123 = int_to_node(123);
 
@@ -322,22 +325,22 @@ mod tests {
         let depth = 4u32;
         let mut set = super::MerklePathSet::new(depth).unwrap();
 
-        set.add_path(index_6, hash_6, path_6).unwrap();
-        set.add_path(index_5, hash_5, path_5).unwrap();
-        set.add_path(index_4, hash_4, path_4).unwrap();
+        set.add_path(index_6, hash_6, path_6.into()).unwrap();
+        set.add_path(index_5, hash_5, path_5.into()).unwrap();
+        set.add_path(index_4, hash_4, path_4.into()).unwrap();
 
         let new_hash_6 = int_to_node(100);
         let new_hash_5 = int_to_node(55);
 
         set.update_leaf(index_6, new_hash_6).unwrap();
         let new_path_4 = set.get_path(depth, index_4).unwrap();
-        let new_hash_67 = calculate_parent_hash(new_hash_6, 14u64, hash_7);
+        let new_hash_67 = MerklePathSet::calculate_parent_hash(new_hash_6, 14u64, hash_7);
         assert_eq!(new_hash_67, new_path_4[1]);
 
         set.update_leaf(index_5, new_hash_5).unwrap();
         let new_path_4 = set.get_path(depth, index_4).unwrap();
         let new_path_6 = set.get_path(depth, index_6).unwrap();
-        let new_hash_45 = calculate_parent_hash(new_hash_5, 13u64, hash_4);
+        let new_hash_45 = MerklePathSet::calculate_parent_hash(new_hash_5, 13u64, hash_4);
         assert_eq!(new_hash_45, new_path_6[1]);
         assert_eq!(new_hash_5, new_path_4[0]);
     }

src/merkle/simple_smt/mod.rs

@@ -1,4 +1,4 @@
-use super::{BTreeMap, MerkleError, Rpo256, RpoDigest, Vec, Word};
+use super::{BTreeMap, MerkleError, MerklePath, Rpo256, RpoDigest, Vec, Word};
 
 #[cfg(test)]
 mod tests;
@@ -102,7 +102,7 @@ impl SimpleSmt {
     /// Returns an error if:
     /// * The specified key does not exist as a branch or leaf node
     /// * The specified depth is greater than the depth of the tree.
-    pub fn get_path(&self, depth: u32, key: u64) -> Result<Vec<Word>, MerkleError> {
+    pub fn get_path(&self, depth: u32, key: u64) -> Result<MerklePath, MerkleError> {
         if depth == 0 {
             return Err(MerkleError::DepthTooSmall(depth));
         } else if depth > self.depth() {
@@ -124,7 +124,7 @@ impl SimpleSmt {
             path.push(sibling_node.into());
             curr_key >>= 1;
         }
-        Ok(path)
+        Ok(path.into())
     }
 
     /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
@@ -133,7 +133,7 @@ impl SimpleSmt {
     /// # Errors
     /// Returns an error if:
     /// * The specified key does not exist as a leaf node.
-    pub fn get_leaf_path(&self, key: u64) -> Result<Vec<Word>, MerkleError> {
+    pub fn get_leaf_path(&self, key: u64) -> Result<MerklePath, MerkleError> {
         self.get_path(self.depth(), key)
     }
 

src/merkle/simple_smt/tests.rs

@@ -103,14 +103,14 @@ fn get_path() {
     let (_, node2, node3) = compute_internal_nodes();
 
     // check depth 2
-    assert_eq!(vec![VALUES4[1], node3], tree.get_path(2, 0).unwrap());
-    assert_eq!(vec![VALUES4[0], node3], tree.get_path(2, 1).unwrap());
-    assert_eq!(vec![VALUES4[3], node2], tree.get_path(2, 2).unwrap());
-    assert_eq!(vec![VALUES4[2], node2], tree.get_path(2, 3).unwrap());
+    assert_eq!(vec![VALUES4[1], node3], *tree.get_path(2, 0).unwrap());
+    assert_eq!(vec![VALUES4[0], node3], *tree.get_path(2, 1).unwrap());
+    assert_eq!(vec![VALUES4[3], node2], *tree.get_path(2, 2).unwrap());
+    assert_eq!(vec![VALUES4[2], node2], *tree.get_path(2, 3).unwrap());
 
     // check depth 1
-    assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
-    assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
+    assert_eq!(vec![node3], *tree.get_path(1, 0).unwrap());
+    assert_eq!(vec![node2], *tree.get_path(1, 1).unwrap());
 }
 
 #[test]
@@ -191,7 +191,7 @@ fn small_tree_opening_is_consistent() {
     for (depth, key, path) in cases {
         let opening = tree.get_path(depth, key).unwrap();
 
-        assert_eq!(path, opening);
+        assert_eq!(path, *opening);
     }
 }