From 429d3bab6fb91541ad59660b5ba6a3a0064a30e9 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Fri, 31 Mar 2023 04:47:12 +0200 Subject: [PATCH] feat: add support for MMR to the MerkleStore --- src/merkle/mmr/full.rs | 97 ++++++++++++++++++++++++++++++++++++++++- src/merkle/mmr/tests.rs | 40 ++++++++++++++++- src/merkle/store/mod.rs | 29 ++++++++++++ 3 files changed, 163 insertions(+), 3 deletions(-) diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs index 92c350c..76f1e3d 100644 --- a/src/merkle/mmr/full.rs +++ b/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 { + 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 // =============================================================================================== diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index 577c6c4..1fb879e 100644 --- a/src/merkle/mmr/tests.rs +++ b/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 = 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::*; diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 557942c..d1cffb6 100644 --- a/src/merkle/store/mod.rs +++ b/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(mut self, leaves: I) -> Result + where + I: IntoIterator, + { + 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(&mut self, leaves: I) -> Result + where + I: IntoIterator, + { + 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