feat: add inner node iterator to MerklePath

This commit is contained in:
Bobbin Threadbare 2023-04-21 14:27:58 -07:00
parent 22c9f382c4
commit 130ae3d12a
4 changed files with 109 additions and 60 deletions

View file

@ -150,9 +150,11 @@ impl MerkleTree {
Ok(()) Ok(())
} }
/// An iterator over every inner node in the tree. The iterator order is unspecified. /// Returns n iterator over every inner node of this [MerkleTree].
pub fn inner_nodes(&self) -> MerkleTreeNodes<'_> { ///
MerkleTreeNodes { /// The iterator order is unspecified.
pub fn inner_nodes(&self) -> InnerNodeIterator<'_> {
InnerNodeIterator {
nodes: &self.nodes, nodes: &self.nodes,
index: 1, // index 0 is just padding, start at 1 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]. /// 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. /// 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>, nodes: &'a Vec<Word>,
index: usize, index: usize,
} }
impl<'a> Iterator for MerkleTreeNodes<'a> { impl<'a> Iterator for InnerNodeIterator<'a> {
type Item = InnerNodeInfo; type Item = InnerNodeInfo;
fn next(&mut self) -> Option<Self::Item> { fn next(&mut self) -> Option<Self::Item> {

View file

@ -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}; use core::ops::{Deref, DerefMut};
// MERKLE PATH // MERKLE PATH
@ -22,6 +22,11 @@ impl MerklePath {
// PROVIDERS // 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. /// Computes the merkle root for this opening.
pub fn compute_root(&self, index: u64, node: Word) -> Result<Word, MerkleError> { pub fn compute_root(&self, index: u64, node: Word) -> Result<Word, MerkleError> {
let mut index = NodeIndex::new(self.depth(), index)?; let mut index = NodeIndex::new(self.depth(), index)?;
@ -34,11 +39,6 @@ impl MerklePath {
Ok(root) 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. /// Verifies the Merkle opening proof towards the provided root.
/// ///
/// Returns `true` if `node` exists at `index` in a Merkle tree with `root`. /// Returns `true` if `node` exists at `index` in a Merkle tree with `root`.
@ -48,6 +48,20 @@ impl MerklePath {
Err(_) => false, 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 { impl From<Vec<Word>> for MerklePath {
@ -72,6 +86,9 @@ impl DerefMut for MerklePath {
} }
} }
// ITERATORS
// ================================================================================================
impl FromIterator<Word> for MerklePath { impl FromIterator<Word> for MerklePath {
fn from_iter<T: IntoIterator<Item = Word>>(iter: T) -> Self { fn from_iter<T: IntoIterator<Item = Word>>(iter: T) -> Self {
Self::new(iter.into_iter().collect()) 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<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 // MERKLE PATH CONTAINERS
// ================================================================================================ // ================================================================================================
@ -110,3 +160,25 @@ pub struct RootPath {
/// The path from `value` to `root` (exclusive). /// The path from `value` to `root` (exclusive).
pub path: MerklePath, 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);
}
}

View file

@ -115,26 +115,6 @@ impl MerkleStore {
MerkleStore { nodes } 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 // PUBLIC ACCESSORS
// -------------------------------------------------------------------------------------------- // --------------------------------------------------------------------------------------------
@ -280,8 +260,11 @@ impl MerkleStore {
// STATE MUTATORS // STATE MUTATORS
// -------------------------------------------------------------------------------------------- // --------------------------------------------------------------------------------------------
/// Adds a new [InnerNodeInfo] into the store. /// Adds a sequence of nodes yielded by the provided iterator into the store.
pub fn extend(&mut self, iter: impl Iterator<Item = InnerNodeInfo>) -> &mut MerkleStore { pub fn extend<I>(&mut self, iter: I) -> &mut MerkleStore
where
I: Iterator<Item = InnerNodeInfo>,
{
for node in iter { for node in iter {
let value: RpoDigest = node.value.into(); let value: RpoDigest = node.value.into();
let left: RpoDigest = node.left.into(); let left: RpoDigest = node.left.into();
@ -301,31 +284,21 @@ impl MerkleStore {
/// include all the nodes into the store. /// include all the nodes into the store.
pub fn add_merkle_path( pub fn add_merkle_path(
&mut self, &mut self,
index_value: u64, index: u64,
mut node: Word, node: Word,
path: MerklePath, path: MerklePath,
) -> Result<Word, MerkleError> { ) -> 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 { debug_assert_eq!(Rpo256::merge(&[left, right]), value);
let (left, right) = match index.is_value_odd() { self.nodes.insert(value, Node { left, right });
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.value
node = parent.into(); });
} Ok(root)
Ok(node)
} }
/// Adds all the nodes of multiple Merkle paths into the store. /// Adds all the nodes of multiple Merkle paths into the store.

View file

@ -547,12 +547,14 @@ fn test_constructors() -> Result<(), MerkleError> {
(3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()), (3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()),
]; ];
let store1 = MerkleStore::default().with_merkle_paths(paths.clone())?; let mut store1 = MerkleStore::default();
let store2 = MerkleStore::default() store1.add_merkle_paths(paths.clone())?;
.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))?)? let mut store2 = MerkleStore::default();
.with_merkle_path(2, LEAVES4[2], mtree.get_path(NodeIndex::make(d, 2))?)? store2.add_merkle_path(0, LEAVES4[0], mtree.get_path(NodeIndex::make(d, 0))?)?;
.with_merkle_path(3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?; 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(); let set = MerklePathSet::new(d).with_paths(paths).unwrap();
for key in [0, 1, 2, 3] { for key in [0, 1, 2, 3] {