Merge pull request #110 from 0xPolygonMiden/bobbin-pathset-fixes

Fix MerklePathSet issues

src/merkle/index.rs

@@ -78,7 +78,7 @@ impl NodeIndex {
         self.depth
     }
 
-    /// Returns the value of the current depth.
+    /// Returns the value of this index.
     pub const fn value(&self) -> u64 {
         self.value
     }
@@ -104,7 +104,7 @@ impl NodeIndex {
     /// arrive at the leaf. From the right-to-left the bit represent the position the hash of the
     /// current element should go.
     ///
-    /// Additionally, the value that is not visisted are the sibling values necessary for a Merkle
+    /// Additionally, the value that is not visited are the sibling values necessary for a Merkle
     /// opening.
     pub fn bit_iterator(&self) -> BitIterator {
         let depth: u32 = self.depth.into();

src/merkle/path_set.rs

@@ -1,4 +1,4 @@
-use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, Vec, Word, ZERO};
+use super::{BTreeMap, MerkleError, MerklePath, NodeIndex, Rpo256, ValuePath, Vec, Word, ZERO};
 
 // MERKLE PATH SET
 // ================================================================================================
@@ -57,14 +57,6 @@ impl MerklePathSet {
         self.total_depth
     }
 
-    /// Returns all the leaf indexes of this path set.
-    pub fn indexes(&self) -> impl Iterator<Item = NodeIndex> + '_ {
-        self.paths
-            .keys()
-            .copied()
-            .map(|index| NodeIndex::new(self.total_depth, index))
-    }
-
     /// Returns a node at the specified index.
     ///
     /// # Errors
@@ -84,15 +76,23 @@ impl MerklePathSet {
             });
         }
 
-        let index_value = index.to_scalar_index();
-        let parity = index_value & 1;
-        let index_value = index_value / 2;
+        let parity = index.value() & 1;
+        let path_key = index.value() - parity;
         self.paths
-            .get(&index_value)
-            .ok_or(MerkleError::NodeNotInSet(index_value))
+            .get(&path_key)
+            .ok_or(MerkleError::NodeNotInSet(path_key))
             .map(|path| path[parity as usize])
     }
 
+    /// Returns a leaf at the specified index.
+    ///
+    /// # Errors
+    /// * The specified index is not valid for the depth of the structure.
+    /// * Leaf with the requested path does not exist in the set.
+    pub fn get_leaf(&self, index: u64) -> Result<Word, MerkleError> {
+        self.get_node(NodeIndex::new(self.depth(), index))
+    }
+
     /// Returns a Merkle path to the node at the specified index. The node itself is
     /// not included in the path.
     ///
@@ -111,18 +111,42 @@ impl MerklePathSet {
             });
         }
 
-        let index_value = index.to_scalar_index();
-        let index = index_value / 2;
-        let parity = index_value & 1;
+        let parity = index.value() & 1;
+        let path_key = index.value() - parity;
         let mut path = self
             .paths
-            .get(&index)
+            .get(&path_key)
             .cloned()
-            .ok_or(MerkleError::NodeNotInSet(index))?;
+            .ok_or(MerkleError::NodeNotInSet(index.value()))?;
         path.remove(parity as usize);
         Ok(path)
     }
 
+    /// Returns all paths in this path set together with their indexes.
+    pub fn to_paths(&self) -> Vec<(u64, ValuePath)> {
+        let mut result = Vec::with_capacity(self.paths.len() * 2);
+
+        for (&index, path) in self.paths.iter() {
+            // push path for the even index into the result
+            let path1 = ValuePath {
+                value: path[0],
+                path: MerklePath::new(path[1..].to_vec()),
+            };
+            result.push((index, path1));
+
+            // push path for the odd index into the result
+            let mut path2 = path.clone();
+            let leaf2 = path2.remove(1);
+            let path2 = ValuePath {
+                value: leaf2,
+                path: path2,
+            };
+            result.push((index + 1, path2));
+        }
+
+        result
+    }
+
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
@@ -141,7 +165,7 @@ impl MerklePathSet {
         value: Word,
         mut path: MerklePath,
     ) -> Result<(), MerkleError> {
-        let depth = (path.len() + 1) as u8;
+        let depth = path.len() as u8;
         let mut index = NodeIndex::new(depth, index_value);
         if index.depth() != self.total_depth {
             return Err(MerkleError::InvalidDepth {
@@ -151,8 +175,6 @@ impl MerklePathSet {
         }
 
         // update the current path
-        let index_value = index.to_scalar_index();
-        let upper_index_value = index_value / 2;
         let parity = index_value & 1;
         path.insert(parity as usize, value);
 
@@ -172,7 +194,8 @@ impl MerklePathSet {
         }
 
         // finish updating the path
-        self.paths.insert(upper_index_value, path);
+        let path_key = index_value - parity;
+        self.paths.insert(path_key, path);
         Ok(())
     }
 
@@ -188,10 +211,9 @@ impl MerklePathSet {
             return Err(MerkleError::InvalidIndex(index));
         }
 
-        let path = match self
-            .paths
-            .get_mut(&index.clone().move_up().to_scalar_index())
-        {
+        let parity = index.value() & 1;
+        let path_key = index.value() - parity;
+        let path = match self.paths.get_mut(&path_key) {
             Some(path) => path,
             None => return Err(MerkleError::NodeNotInSet(base_index_value)),
         };
@@ -255,7 +277,7 @@ mod tests {
 
         let root_exp = calculate_parent_hash(parent0, 0, parent1);
 
-        let set = super::MerklePathSet::new(3)
+        let set = super::MerklePathSet::new(2)
             .with_paths([(0, leaf0, vec![leaf1, parent1].into())])
             .unwrap();
 
@@ -267,7 +289,7 @@ mod tests {
         let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
         let hash_6 = int_to_node(6);
         let index = 6_u64;
-        let depth = 4_u8;
+        let depth = 3_u8;
         let set = super::MerklePathSet::new(depth)
             .with_paths([(index, hash_6, path_6.clone().into())])
             .unwrap();
@@ -282,7 +304,7 @@ mod tests {
         let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
         let hash_6 = int_to_node(6);
         let index = 6_u64;
-        let depth = 4_u8;
+        let depth = 3_u8;
         let set = MerklePathSet::new(depth)
             .with_paths([(index, hash_6, path_6.into())])
             .unwrap();
@@ -312,7 +334,7 @@ mod tests {
         let index_6 = 6_u64;
         let index_5 = 5_u64;
         let index_4 = 4_u64;
-        let depth = 4_u8;
+        let depth = 3_u8;
         let mut set = MerklePathSet::new(depth)
             .with_paths([
                 (index_6, hash_6, path_6.into()),
@@ -337,6 +359,58 @@ mod tests {
         assert_eq!(new_hash_5, new_path_4[0]);
     }
 
+    #[test]
+    fn depth_3_is_correct() {
+        let a = int_to_node(1);
+        let b = int_to_node(2);
+        let c = int_to_node(3);
+        let d = int_to_node(4);
+        let e = int_to_node(5);
+        let f = int_to_node(6);
+        let g = int_to_node(7);
+        let h = int_to_node(8);
+
+        let i = Rpo256::merge(&[a.into(), b.into()]);
+        let j = Rpo256::merge(&[c.into(), d.into()]);
+        let k = Rpo256::merge(&[e.into(), f.into()]);
+        let l = Rpo256::merge(&[g.into(), h.into()]);
+
+        let m = Rpo256::merge(&[i.into(), j.into()]);
+        let n = Rpo256::merge(&[k.into(), l.into()]);
+
+        let root = Rpo256::merge(&[m.into(), n.into()]);
+
+        let mut set = MerklePathSet::new(3);
+
+        let value = b;
+        let index = 1;
+        let path = MerklePath::new([a.into(), j.into(), n.into()].to_vec());
+        set.add_path(index, value, path.clone()).unwrap();
+        assert_eq!(value, set.get_leaf(index).unwrap());
+        assert_eq!(Word::from(root), set.root());
+
+        let value = e;
+        let index = 4;
+        let path = MerklePath::new([f.into(), l.into(), m.into()].to_vec());
+        set.add_path(index, value, path.clone()).unwrap();
+        assert_eq!(value, set.get_leaf(index).unwrap());
+        assert_eq!(Word::from(root), set.root());
+
+        let value = a;
+        let index = 0;
+        let path = MerklePath::new([b.into(), j.into(), n.into()].to_vec());
+        set.add_path(index, value, path.clone()).unwrap();
+        assert_eq!(value, set.get_leaf(index).unwrap());
+        assert_eq!(Word::from(root), set.root());
+
+        let value = h;
+        let index = 7;
+        let path = MerklePath::new([g.into(), k.into(), m.into()].to_vec());
+        set.add_path(index, value, path.clone()).unwrap();
+        assert_eq!(value, set.get_leaf(index).unwrap());
+        assert_eq!(Word::from(root), set.root());
+    }
+
     // HELPER FUNCTIONS
     // --------------------------------------------------------------------------------------------
 

src/merkle/store/mod.rs

@@ -317,11 +317,10 @@ impl MerkleStore {
     /// For further reference, check [MerkleStore::add_merkle_path].
     pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result<Word, MerkleError> {
         let root = path_set.root();
-        path_set.indexes().try_fold(root, |_, index| {
-            let node = path_set.get_node(index)?;
-            let path = path_set.get_path(index)?;
-            self.add_merkle_path(index.value(), node, path)
-        })
+        for (index, path) in path_set.to_paths() {
+            self.add_merkle_path(index, path.value, path.path)?;
+        }
+        Ok(root)
     }
 
     /// Sets a node to `value`.

src/merkle/store/tests.rs

@@ -320,28 +320,28 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
         .add_merkle_paths(paths.clone())
         .expect("the valid paths must work");
 
-    let depth = 3;
+    let depth = 2;
     let set = MerklePathSet::new(depth).with_paths(paths).unwrap();
 
     // STORE LEAVES ARE CORRECT ==============================================================
     // checks the leaves in the store corresponds to the expected values
     assert_eq!(
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 0)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 0)),
         Ok(LEAVES4[0]),
         "node 0 must be in the set"
     );
     assert_eq!(
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 1)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 1)),
         Ok(LEAVES4[1]),
         "node 1 must be in the set"
     );
     assert_eq!(
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 2)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 2)),
         Ok(LEAVES4[2]),
         "node 2 must be in the set"
     );
     assert_eq!(
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 3)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 3)),
         Ok(LEAVES4[3]),
         "node 3 must be in the set"
     );
@@ -350,29 +350,29 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
     // sanity check the values returned by the store and the set
     assert_eq!(
         set.get_node(NodeIndex::new(set.depth(), 0)),
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 0)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 0)),
         "node 0 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
         set.get_node(NodeIndex::new(set.depth(), 1)),
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 1)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 1)),
         "node 1 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
         set.get_node(NodeIndex::new(set.depth(), 2)),
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 2)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 2)),
         "node 2 must be the same for both SparseMerkleTree and MerkleStore"
     );
     assert_eq!(
         set.get_node(NodeIndex::new(set.depth(), 3)),
-        store.get_node(set.root(), NodeIndex::new(set.depth() - 1, 3)),
+        store.get_node(set.root(), NodeIndex::new(set.depth(), 3)),
         "node 3 must be the same for both SparseMerkleTree and MerkleStore"
     );
 
     // STORE MERKLE PATH MATCHS ==============================================================
     // assert the merkle path returned by the store is the same as the one in the set
     let result = store
-        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 0))
+        .get_path(set.root(), NodeIndex::new(set.depth(), 0))
         .unwrap();
     assert_eq!(
         LEAVES4[0], result.value,
@@ -385,7 +385,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
     );
 
     let result = store
-        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 1))
+        .get_path(set.root(), NodeIndex::new(set.depth(), 1))
         .unwrap();
     assert_eq!(
         LEAVES4[1], result.value,
@@ -398,7 +398,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
     );
 
     let result = store
-        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 2))
+        .get_path(set.root(), NodeIndex::new(set.depth(), 2))
         .unwrap();
     assert_eq!(
         LEAVES4[2], result.value,
@@ -411,7 +411,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
     );
 
     let result = store
-        .get_path(set.root(), NodeIndex::new(set.depth() - 1, 3))
+        .get_path(set.root(), NodeIndex::new(set.depth(), 3))
         .unwrap();
     assert_eq!(
         LEAVES4[3], result.value,
@@ -540,7 +540,7 @@ fn test_constructors() -> Result<(), MerkleError> {
         .with_merkle_path(1, LEAVES4[1], mtree.get_path(NodeIndex::new(d, 1))?)?
         .with_merkle_path(2, LEAVES4[2], mtree.get_path(NodeIndex::new(d, 2))?)?
         .with_merkle_path(3, LEAVES4[3], mtree.get_path(NodeIndex::new(d, 3))?)?;
-    let set = MerklePathSet::new(d + 1).with_paths(paths).unwrap();
+    let set = MerklePathSet::new(d).with_paths(paths).unwrap();
 
     for key in [0, 1, 2, 3] {
         let index = NodeIndex::new(d, key);
@@ -548,7 +548,7 @@ fn test_constructors() -> Result<(), MerkleError> {
         let value_path2 = store2.get_path(set.root(), index)?;
         assert_eq!(value_path1, value_path2);
 
-        let index = NodeIndex::new(d + 1, key);
+        let index = NodeIndex::new(d, key);
         assert_eq!(set.get_path(index)?, value_path1.path);
     }