feat: add inner node iterator to MerklePath
Bobbin Threadbare bobbinth@protonmail.com
Fri, 21 Apr 2023 14:27:58 -0700
4 files changed,
109 insertions(+),
60 deletions(-)
M
src/merkle/merkle_tree.rs
→
src/merkle/merkle_tree.rs
@@ -150,9 +150,11 @@
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 @@
/// 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<Word>, index: usize, } -impl<'a> Iterator for MerkleTreeNodes<'a> { +impl<'a> Iterator for InnerNodeIterator<'a> { type Item = InnerNodeInfo; fn next(&mut self) -> Option<Self::Item> {
M
src/merkle/path.rs
→
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 @@
// 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<Word, MerkleError> { let mut index = NodeIndex::new(self.depth(), index)?;@@ -34,11 +39,6 @@ });
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 @@ Ok(computed_root) => root == &computed_root,
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<InnerNodeIterator, MerkleError> { + Ok(InnerNodeIterator { + nodes: &self.nodes, + index: NodeIndex::new(self.depth(), index)?, + value: node, + }) + } } impl From<Vec<Word>> for MerklePath {@@ -72,6 +86,9 @@ &mut self.nodes
} } +// ITERATORS +// ================================================================================================ + impl FromIterator<Word> for MerklePath { fn from_iter<T: IntoIterator<Item = Word>>(iter: T) -> Self { Self::new(iter.into_iter().collect())@@ -87,6 +104,39 @@ self.nodes.into_iter()
} } +/// An iterator over internal nodes of a [MerklePath]. +pub struct InnerNodeIterator<'a> { + nodes: &'a Vec<Word>, + index: NodeIndex, + value: Word, +} + +impl<'a> Iterator for InnerNodeIterator<'a> { + type Item = InnerNodeInfo; + + fn next(&mut self) -> Option<Self::Item> { + 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 root: Word,
/// 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); + } +}
M
src/merkle/store/mod.rs
→
src/merkle/store/mod.rs
@@ -115,26 +115,6 @@
MerkleStore { nodes } } - /// Appends the provided merkle path set. - pub fn with_merkle_path( - mut self, - index_value: u64, - node: Word, - path: MerklePath, - ) -> Result<Self, MerkleError> { - self.add_merkle_path(index_value, node, path)?; - Ok(self) - } - - /// Appends the provided merkle path set. - pub fn with_merkle_paths<I>(mut self, paths: I) -> Result<Self, MerkleError> - where - I: IntoIterator<Item = (u64, Word, MerklePath)>, - { - self.add_merkle_paths(paths)?; - Ok(self) - } - // PUBLIC ACCESSORS // --------------------------------------------------------------------------------------------@@ -280,8 +260,11 @@
// STATE MUTATORS // -------------------------------------------------------------------------------------------- - /// Adds a new [InnerNodeInfo] into the store. - pub fn extend(&mut self, iter: impl Iterator<Item = InnerNodeInfo>) -> &mut MerkleStore { + /// Adds a sequence of nodes yielded by the provided iterator into the store. + pub fn extend<I>(&mut self, iter: I) -> &mut MerkleStore + where + I: Iterator<Item = InnerNodeInfo>, + { for node in iter { let value: RpoDigest = node.value.into(); let left: RpoDigest = node.left.into();@@ -301,31 +284,21 @@ /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
/// 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<Word, MerkleError> { - 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(), - }, - ); + debug_assert_eq!(Rpo256::merge(&[left, right]), value); + self.nodes.insert(value, Node { left, right }); - index.move_up(); - node = parent.into(); - } - - Ok(node) + node.value + }); + Ok(root) } /// Adds all the nodes of multiple Merkle paths into the store.
M
src/merkle/store/tests.rs
→
src/merkle/store/tests.rs
@@ -547,12 +547,14 @@ (2, LEAVES4[2], mtree.get_path(NodeIndex::make(d, 2)).unwrap()),
(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] {