From 1df4318399eeaf1e689a5959c364e0901f6174e6 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Thu, 30 Mar 2023 23:51:44 +0200 Subject: [PATCH] feat: add node iterator to MerkleTree --- src/merkle/merkle_tree.rs | 83 +++++++++++++++++++++++++++++++++++++-- src/merkle/mod.rs | 3 ++ src/merkle/node.rs | 9 +++++ src/merkle/store/mod.rs | 33 +++++----------- 4 files changed, 101 insertions(+), 27 deletions(-) create mode 100644 src/merkle/node.rs diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 434c8e7..6b767ae 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -1,4 +1,6 @@ -use super::{Felt, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word}; +use super::{ + Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word, +}; use crate::{ utils::{string::String, uninit_vector, word_to_hex}, FieldElement, @@ -12,7 +14,7 @@ use winter_math::log2; /// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two). #[derive(Debug, Clone, PartialEq, Eq)] pub struct MerkleTree { - pub(crate) nodes: Vec, + nodes: Vec, } impl MerkleTree { @@ -158,6 +160,47 @@ impl MerkleTree { Ok(()) } + + /// An iterator over every inner node in the tree. The iterator order is unspecified. + pub fn inner_nodes(&self) -> MerkleTreeNodes<'_> { + MerkleTreeNodes { + nodes: &self.nodes, + index: 1, // index 0 is just padding, start at 1 + } + } +} + +// ITERATORS +// ================================================================================================ + +/// 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> { + nodes: &'a Vec, + index: usize, +} + +impl<'a> Iterator for MerkleTreeNodes<'a> { + type Item = InnerNodeInfo; + + fn next(&mut self) -> Option { + if self.index < self.nodes.len() / 2 { + let value = self.index; + let left = self.index * 2; + let right = left + 1; + + self.index += 1; + + Some(InnerNodeInfo { + value: self.nodes[value], + left: self.nodes[left], + right: self.nodes[right], + }) + } else { + None + } + } } /// Utility to vizualize a [MerkleTree] in text. @@ -212,7 +255,7 @@ pub fn path_to_text(path: &MerklePath) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::merkle::int_to_node; + use crate::merkle::{int_to_node, InnerNodeInfo}; use core::mem::size_of; use proptest::prelude::*; @@ -323,6 +366,40 @@ mod tests { assert_eq!(expected_tree.nodes, tree.nodes); } + #[test] + fn nodes() -> Result<(), MerkleError> { + let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + let root = tree.root(); + let l1n0 = tree.get_node(NodeIndex::new(1, 0))?; + let l1n1 = tree.get_node(NodeIndex::new(1, 1))?; + let l2n0 = tree.get_node(NodeIndex::new(2, 0))?; + let l2n1 = tree.get_node(NodeIndex::new(2, 1))?; + let l2n2 = tree.get_node(NodeIndex::new(2, 2))?; + let l2n3 = tree.get_node(NodeIndex::new(2, 3))?; + + let nodes: Vec = tree.inner_nodes().collect(); + let expected = vec![ + InnerNodeInfo { + value: root, + left: l1n0, + right: l1n1, + }, + InnerNodeInfo { + value: l1n0, + left: l2n0, + right: l2n1, + }, + InnerNodeInfo { + value: l1n1, + left: l2n2, + right: l2n3, + }, + ]; + assert_eq!(nodes, expected); + + Ok(()) + } + proptest! { #[test] fn arbitrary_word_can_be_represented_as_digest( diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index b0c7b2d..c12ac7a 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -32,6 +32,9 @@ pub use mmr::{Mmr, MmrPeaks}; mod store; pub use store::MerkleStore; +mod node; +pub use node::InnerNodeInfo; + // ERRORS // ================================================================================================ diff --git a/src/merkle/node.rs b/src/merkle/node.rs new file mode 100644 index 0000000..681e306 --- /dev/null +++ b/src/merkle/node.rs @@ -0,0 +1,9 @@ +use super::Word; + +/// Representation of a node with two children used for iterating over containers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InnerNodeInfo { + pub value: Word, + pub left: Word, + pub right: Word, +} diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index e21335f..557942c 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -249,33 +249,18 @@ impl MerkleStore { return Err(MerkleError::DepthTooSmall(leaves.len() as u8)); } - let layers = leaves.len().ilog2(); let tree = MerkleTree::new(leaves)?; - - let mut depth = 0; - let mut parent_offset = 1; - let mut child_offset = 2; - while depth < layers { - let layer_size = 1usize << depth; - for _ in 0..layer_size { - // merkle tree is using level form representation, so left and right siblings are - // next to each other - let left = tree.nodes[child_offset]; - let right = tree.nodes[child_offset + 1]; - self.nodes.insert( - tree.nodes[parent_offset].into(), - Node { - left: left.into(), - right: right.into(), - }, - ); - parent_offset += 1; - child_offset += 2; - } - depth += 1; + for node in tree.inner_nodes() { + self.nodes.insert( + node.value.into(), + Node { + left: node.left.into(), + right: node.right.into(), + }, + ); } - Ok(tree.nodes[1]) + Ok(tree.root()) } /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the