diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 106ec27..42f289a 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -150,9 +150,11 @@ impl MerkleTree { Ok(()) } - /// An iterator over every inner node in the tree. The iterator order is unspecified. - pub fn inner_nodes(&self) -> MerkleTreeNodes<'_> { - MerkleTreeNodes { + /// Returns n iterator over every inner node of this [MerkleTree]. + /// + /// The iterator order is unspecified. + pub fn inner_nodes(&self) -> InnerNodeIterator<'_> { + InnerNodeIterator { nodes: &self.nodes, index: 1, // index 0 is just padding, start at 1 } @@ -165,12 +167,12 @@ impl MerkleTree { /// An iterator over every inner node of the [MerkleTree]. /// /// Use this to extract the data of the tree, there is no guarantee on the order of the elements. -pub struct MerkleTreeNodes<'a> { +pub struct InnerNodeIterator<'a> { nodes: &'a Vec, index: usize, } -impl<'a> Iterator for MerkleTreeNodes<'a> { +impl<'a> Iterator for InnerNodeIterator<'a> { type Item = InnerNodeInfo; fn next(&mut self) -> Option { diff --git a/src/merkle/path.rs b/src/merkle/path.rs index 964f6d0..e873d84 100644 --- a/src/merkle/path.rs +++ b/src/merkle/path.rs @@ -1,4 +1,4 @@ -use super::{vec, MerkleError, NodeIndex, Rpo256, Vec, Word}; +use super::{vec, InnerNodeInfo, MerkleError, NodeIndex, Rpo256, Vec, Word}; use core::ops::{Deref, DerefMut}; // MERKLE PATH @@ -22,6 +22,11 @@ impl MerklePath { // PROVIDERS // -------------------------------------------------------------------------------------------- + /// Returns the depth in which this Merkle path proof is valid. + pub fn depth(&self) -> u8 { + self.nodes.len() as u8 + } + /// Computes the merkle root for this opening. pub fn compute_root(&self, index: u64, node: Word) -> Result { let mut index = NodeIndex::new(self.depth(), index)?; @@ -34,11 +39,6 @@ impl MerklePath { Ok(root) } - /// 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`. @@ -48,6 +48,20 @@ impl MerklePath { Err(_) => false, } } + + /// Returns an iterator over every inner node of this [MerklePath]. + /// + /// The iteration order is unspecified. + /// + /// # Errors + /// Returns an error if the specified index is not valid for this path. + pub fn inner_nodes(&self, index: u64, node: Word) -> Result { + Ok(InnerNodeIterator { + nodes: &self.nodes, + index: NodeIndex::new(self.depth(), index)?, + value: node, + }) + } } impl From> for MerklePath { @@ -72,6 +86,9 @@ impl DerefMut for MerklePath { } } +// ITERATORS +// ================================================================================================ + impl FromIterator for MerklePath { fn from_iter>(iter: T) -> Self { Self::new(iter.into_iter().collect()) @@ -87,6 +104,39 @@ impl IntoIterator for MerklePath { } } +/// An iterator over internal nodes of a [MerklePath]. +pub struct InnerNodeIterator<'a> { + nodes: &'a Vec, + index: NodeIndex, + value: Word, +} + +impl<'a> Iterator for InnerNodeIterator<'a> { + type Item = InnerNodeInfo; + + fn next(&mut self) -> Option { + if !self.index.is_root() { + let sibling_pos = self.nodes.len() - self.index.depth() as usize; + let (left, right) = if self.index.is_value_odd() { + (self.nodes[sibling_pos], self.value) + } else { + (self.value, self.nodes[sibling_pos]) + }; + + self.value = Rpo256::merge(&[left.into(), right.into()]).into(); + self.index.move_up(); + + Some(InnerNodeInfo { + value: self.value, + left, + right, + }) + } else { + None + } + } +} + // MERKLE PATH CONTAINERS // ================================================================================================ @@ -110,3 +160,25 @@ pub struct RootPath { /// The path from `value` to `root` (exclusive). pub path: MerklePath, } + +// TESTS +// ================================================================================================ + +#[cfg(test)] +mod tests { + use crate::merkle::{int_to_node, MerklePath}; + + #[test] + fn test_inner_nodes() { + let nodes = vec![int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)]; + let merkle_path = MerklePath::new(nodes); + + let index = 6; + let node = int_to_node(5); + let root = merkle_path.compute_root(index, node).unwrap(); + + let inner_root = merkle_path.inner_nodes(index, node).unwrap().last().unwrap().value; + + assert_eq!(root, inner_root); + } +} diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 442ff00..96ca54a 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -115,26 +115,6 @@ impl MerkleStore { MerkleStore { nodes } } - /// Appends the provided merkle path set. - pub fn with_merkle_path( - mut self, - index_value: u64, - node: Word, - path: MerklePath, - ) -> Result { - self.add_merkle_path(index_value, node, path)?; - Ok(self) - } - - /// Appends the provided merkle path set. - pub fn with_merkle_paths(mut self, paths: I) -> Result - where - I: IntoIterator, - { - self.add_merkle_paths(paths)?; - Ok(self) - } - // PUBLIC ACCESSORS // -------------------------------------------------------------------------------------------- @@ -280,8 +260,11 @@ impl MerkleStore { // STATE MUTATORS // -------------------------------------------------------------------------------------------- - /// Adds a new [InnerNodeInfo] into the store. - pub fn extend(&mut self, iter: impl Iterator) -> &mut MerkleStore { + /// Adds a sequence of nodes yielded by the provided iterator into the store. + pub fn extend(&mut self, iter: I) -> &mut MerkleStore + where + I: Iterator, + { for node in iter { let value: RpoDigest = node.value.into(); let left: RpoDigest = node.left.into(); @@ -301,31 +284,21 @@ impl MerkleStore { /// include all the nodes into the store. pub fn add_merkle_path( &mut self, - index_value: u64, - mut node: Word, + index: u64, + node: Word, path: MerklePath, ) -> Result { - let mut index = NodeIndex::new(path.len() as u8, index_value)?; + let root = path.inner_nodes(index, node)?.fold(Word::default(), |_, node| { + let value: RpoDigest = node.value.into(); + let left: RpoDigest = node.left.into(); + let right: RpoDigest = node.right.into(); - for sibling in path { - let (left, right) = match index.is_value_odd() { - true => (sibling, node), - false => (node, sibling), - }; - let parent = Rpo256::merge(&[left.into(), right.into()]); - self.nodes.insert( - parent, - Node { - left: left.into(), - right: right.into(), - }, - ); - - index.move_up(); - node = parent.into(); - } + debug_assert_eq!(Rpo256::merge(&[left, right]), value); + self.nodes.insert(value, Node { left, right }); - Ok(node) + node.value + }); + Ok(root) } /// Adds all the nodes of multiple Merkle paths into the store. diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 28ac8d7..02e12b2 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -547,12 +547,14 @@ fn test_constructors() -> Result<(), MerkleError> { (3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()), ]; - let store1 = MerkleStore::default().with_merkle_paths(paths.clone())?; - let store2 = MerkleStore::default() - .with_merkle_path(0, LEAVES4[0], mtree.get_path(NodeIndex::make(d, 0))?)? - .with_merkle_path(1, LEAVES4[1], mtree.get_path(NodeIndex::make(d, 1))?)? - .with_merkle_path(2, LEAVES4[2], mtree.get_path(NodeIndex::make(d, 2))?)? - .with_merkle_path(3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?; + let mut store1 = MerkleStore::default(); + store1.add_merkle_paths(paths.clone())?; + + let mut store2 = MerkleStore::default(); + store2.add_merkle_path(0, LEAVES4[0], mtree.get_path(NodeIndex::make(d, 0))?)?; + store2.add_merkle_path(1, LEAVES4[1], mtree.get_path(NodeIndex::make(d, 1))?)?; + store2.add_merkle_path(2, LEAVES4[2], mtree.get_path(NodeIndex::make(d, 2))?)?; + store2.add_merkle_path(3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?; let set = MerklePathSet::new(d).with_paths(paths).unwrap(); for key in [0, 1, 2, 3] {