Merge pull request #118 from 0xPolygonMiden/hacka-support-mmr-in-the-merkle-store

feat: add support for MMR to the MerkleStore

src/merkle/mmr/full.rs

@@ -11,8 +11,10 @@
 //! merged, creating a new tree with depth d+1, this process is continued until the property is
 //! restabilished.
 use super::bit::TrueBitPositionIterator;
-use super::{super::Vec, MmrPeaks, MmrProof, Rpo256, Word};
-use crate::merkle::MerklePath;
+use super::{
+    super::{InnerNodeInfo, MerklePath, Vec},
+    MmrPeaks, MmrProof, Rpo256, Word,
+};
 use core::fmt::{Display, Formatter};
 
 #[cfg(feature = "std")]
@@ -190,6 +192,16 @@ impl Mmr {
         }
     }
 
+    /// An iterator over inner nodes in the [Mmm]. The order of iteration is unspecified.
+    pub fn inner_nodes(&self) -> MmrNodes {
+        MmrNodes {
+            mmr: self,
+            forest: 0,
+            last_right: 0,
+            index: 0,
+        }
+    }
+
     // UTILITIES
     // ============================================================================================
 
@@ -246,6 +258,87 @@ where
     }
 }
 
+// ITERATOR
+// ===============================================================================================
+
+/// Yields inner nodes of the [Mmr].
+pub struct MmrNodes<'a> {
+    /// [Mmr] being yielded, when its `forest` value is matched, the iterations is finished.
+    mmr: &'a Mmr,
+    /// Keeps track of the left nodes yielded so far waiting for a right pair, this matches the
+    /// semantics of the [Mmr]'s forest attribute, since that too works as a buffer of left nodes
+    /// waiting for a pair to be hashed together.
+    forest: usize,
+    /// Keeps track of the last right node yielded, after this value is set, the next iteration
+    /// will be its parent with its corresponding left node that has been yield already.
+    last_right: usize,
+    /// The current index in the `nodes` vector.
+    index: usize,
+}
+
+impl<'a> Iterator for MmrNodes<'a> {
+    type Item = InnerNodeInfo;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        debug_assert!(
+            self.last_right.count_ones() <= 1,
+            "last_right tracks zero or one element"
+        );
+
+        // only parent nodes are emitted, remove the single node tree from the forest
+        let target = self.mmr.forest & (usize::MAX << 1);
+
+        if self.forest < target {
+            if self.last_right == 0 {
+                // yield the left leaf
+                debug_assert!(self.last_right == 0, "left must be before right");
+                self.forest |= 1;
+                self.index += 1;
+
+                // yield the right leaf
+                debug_assert!((self.forest & 1) == 1, "right must be after left");
+                self.last_right |= 1;
+                self.index += 1;
+            };
+
+            debug_assert!(
+                self.forest & self.last_right != 0,
+                "parent requires both a left and right",
+            );
+
+            // compute the number of nodes in the right tree, this is the offset to the
+            // previous left parent
+            let right_nodes = nodes_in_forest(self.last_right);
+            // the next parent position is one above the position of the pair
+            let parent = self.last_right << 1;
+
+            // the left node has been paired and the current parent yielded, removed it from the forest
+            self.forest ^= self.last_right;
+            if self.forest & parent == 0 {
+                // this iteration yielded the left parent node
+                debug_assert!(self.forest & 1 == 0, "next iteration yields a left leaf");
+                self.last_right = 0;
+                self.forest ^= parent;
+            } else {
+                // the left node of the parent level has been yielded already, this iteration
+                // was the right parent. Next iteration yields their parent.
+                self.last_right = parent;
+            }
+
+            // yields a parent
+            let value = self.mmr.nodes[self.index];
+            let right = self.mmr.nodes[self.index - 1];
+            let left = self.mmr.nodes[self.index - 1 - right_nodes];
+            self.index += 1;
+            let node = InnerNodeInfo { value, left, right };
+
+            Some(node)
+        } else {
+            None
+        }
+    }
+}
+
 // UTILITIES
 // ===============================================================================================
 

src/merkle/mmr/tests.rs

@@ -1,6 +1,9 @@
 use super::bit::TrueBitPositionIterator;
 use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest};
-use super::{super::Vec, Mmr, Rpo256, Word};
+use super::{
+    super::{InnerNodeInfo, Vec},
+    Mmr, Rpo256, Word,
+};
 use crate::merkle::{int_to_node, MerklePath};
 
 #[test]
@@ -410,6 +413,41 @@ fn test_bit_position_iterator() {
     );
 }
 
+#[test]
+fn test_mmr_inner_nodes() {
+    let mmr: Mmr = LEAVES.into();
+    let nodes: Vec<InnerNodeInfo> = mmr.inner_nodes().collect();
+
+    let h01 = *Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat());
+    let h23 = *Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat());
+    let h0123 = *Rpo256::hash_elements(&[h01, h23].concat());
+    let h45 = *Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat());
+    let postorder = vec![
+        InnerNodeInfo {
+            value: h01,
+            left: LEAVES[0],
+            right: LEAVES[1],
+        },
+        InnerNodeInfo {
+            value: h23,
+            left: LEAVES[2],
+            right: LEAVES[3],
+        },
+        InnerNodeInfo {
+            value: h0123,
+            left: h01,
+            right: h23,
+        },
+        InnerNodeInfo {
+            value: h45,
+            left: LEAVES[4],
+            right: LEAVES[5],
+        },
+    ];
+
+    assert_eq!(postorder, nodes);
+}
+
 mod property_tests {
     use super::leaf_to_corresponding_tree;
     use proptest::prelude::*;

src/merkle/store/mod.rs

@@ -1,3 +1,4 @@
+use super::mmr::{Mmr, MmrPeaks};
 use super::{
     BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree,
     NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
@@ -151,6 +152,15 @@ impl MerkleStore {
         Ok(self)
     }
 
+    /// Appends the provided [Mmr] represented by its `leaves` to the set.
+    pub fn with_mmr<I>(mut self, leaves: I) -> Result<Self, MerkleError>
+    where
+        I: IntoIterator<Item = Word>,
+    {
+        self.add_mmr(leaves)?;
+        Ok(self)
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -375,6 +385,25 @@ impl MerkleStore {
         Ok(root)
     }
 
+    /// Appends the provided [Mmr] into the store.
+    pub fn add_mmr<I>(&mut self, leaves: I) -> Result<MmrPeaks, MerkleError>
+    where
+        I: IntoIterator<Item = Word>,
+    {
+        let mmr = Mmr::from(leaves);
+        for node in mmr.inner_nodes() {
+            self.nodes.insert(
+                node.value.into(),
+                Node {
+                    left: node.left.into(),
+                    right: node.right.into(),
+                },
+            );
+        }
+
+        Ok(mmr.accumulator())
+    }
+
     /// Sets a node to `value`.
     ///
     /// # Errors