From 21a8cbcb45536c15b663096a18a5a7ce378de296 Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Sat, 11 Feb 2023 12:23:10 +0100 Subject: [PATCH] feat: add merkle path wrapper A Merkle path is a vector of nodes, regardless of the Merkle tree implementation. This commit introduces an encapsulation for such vector, also to provide functionality that is common between different algorithms such as opening verification. related issue: #36 --- src/merkle/merkle_tree.rs | 18 +- src/merkle/mod.rs | 11 +- src/merkle/path.rs | 89 +++++++ .../{merkle_path_set.rs => path_set.rs} | 219 +++++++++--------- src/merkle/simple_smt/mod.rs | 8 +- src/merkle/simple_smt/tests.rs | 14 +- 6 files changed, 227 insertions(+), 132 deletions(-) create mode 100644 src/merkle/path.rs rename src/merkle/{merkle_path_set.rs => path_set.rs} (69%) diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 5eff1c35..7adbccb2 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -1,4 +1,4 @@ -use super::{Felt, MerkleError, Rpo256, RpoDigest, Vec, Word}; +use super::{Felt, MerkleError, MerklePath, Rpo256, RpoDigest, Vec, Word}; use crate::{utils::uninit_vector, FieldElement}; use core::slice; use winter_math::log2; @@ -88,7 +88,7 @@ impl MerkleTree { /// Returns an error if: /// * The specified depth is greater than the depth of the tree. /// * The specified index not valid for the specified depth. - pub fn get_path(&self, depth: u32, index: u64) -> Result, MerkleError> { + pub fn get_path(&self, depth: u32, index: u64) -> Result { if depth == 0 { return Err(MerkleError::DepthTooSmall(depth)); } else if depth > self.depth() { @@ -106,7 +106,7 @@ impl MerkleTree { pos >>= 1; } - Ok(path) + Ok(path.into()) } /// Replaces the leaf at the specified index with the provided value. @@ -206,14 +206,14 @@ mod tests { let (_, node2, node3) = compute_internal_nodes(); // check depth 2 - assert_eq!(vec![LEAVES4[1], node3], tree.get_path(2, 0).unwrap()); - assert_eq!(vec![LEAVES4[0], node3], tree.get_path(2, 1).unwrap()); - assert_eq!(vec![LEAVES4[3], node2], tree.get_path(2, 2).unwrap()); - assert_eq!(vec![LEAVES4[2], node2], tree.get_path(2, 3).unwrap()); + assert_eq!(vec![LEAVES4[1], node3], *tree.get_path(2, 0).unwrap()); + assert_eq!(vec![LEAVES4[0], node3], *tree.get_path(2, 1).unwrap()); + assert_eq!(vec![LEAVES4[3], node2], *tree.get_path(2, 2).unwrap()); + assert_eq!(vec![LEAVES4[2], node2], *tree.get_path(2, 3).unwrap()); // check depth 1 - assert_eq!(vec![node3], tree.get_path(1, 0).unwrap()); - assert_eq!(vec![node2], tree.get_path(1, 1).unwrap()); + assert_eq!(vec![node3], *tree.get_path(1, 0).unwrap()); + assert_eq!(vec![node2], *tree.get_path(1, 1).unwrap()); } #[test] diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 87cd80f5..014d4c76 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -1,6 +1,6 @@ use super::{ hash::rpo::{Rpo256, RpoDigest}, - utils::collections::{BTreeMap, Vec}, + utils::collections::{vec, BTreeMap, Vec}, Felt, Word, ZERO, }; use core::fmt; @@ -8,8 +8,11 @@ use core::fmt; mod merkle_tree; pub use merkle_tree::MerkleTree; -mod merkle_path_set; -pub use merkle_path_set::MerklePathSet; +mod path; +pub use path::MerklePath; + +mod path_set; +pub use path_set::MerklePathSet; mod simple_smt; pub use simple_smt::SimpleSmt; @@ -24,7 +27,7 @@ pub enum MerkleError { NumLeavesNotPowerOfTwo(usize), InvalidIndex(u32, u64), InvalidDepth(u32, u32), - InvalidPath(Vec), + InvalidPath(MerklePath), InvalidEntriesCount(usize, usize), NodeNotInSet(u64), } diff --git a/src/merkle/path.rs b/src/merkle/path.rs new file mode 100644 index 00000000..2b81bbe3 --- /dev/null +++ b/src/merkle/path.rs @@ -0,0 +1,89 @@ +use super::{vec, Rpo256, Vec, Word}; +use core::ops::{Deref, DerefMut}; + +// MERKLE PATH +// ================================================================================================ + +/// A merkle path container, composed of a sequence of nodes of a Merkle tree. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct MerklePath { + nodes: Vec, +} + +impl MerklePath { + // CONSTRUCTORS + // -------------------------------------------------------------------------------------------- + + /// Creates a new Merkle path from a list of nodes. + pub fn new(nodes: Vec) -> Self { + Self { nodes } + } + + // PROVIDERS + // -------------------------------------------------------------------------------------------- + + /// Computes the merkle root for this opening. + pub fn compute_root(&self, mut index: u64, node: Word) -> Word { + self.nodes.iter().copied().fold(node, |node, sibling| { + // build the input node, considering the parity of the current index. + let is_right_sibling = (index & 1) == 1; + let input = if is_right_sibling { + [sibling.into(), node.into()] + } else { + [node.into(), sibling.into()] + }; + // compute the node and move to the next iteration. + index >>= 1; + Rpo256::merge(&input).into() + }) + } + + /// 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`. + pub fn verify(&self, index: u64, node: Word, root: &Word) -> bool { + root == &self.compute_root(index, node) + } +} + +impl From> for MerklePath { + fn from(path: Vec) -> Self { + Self::new(path) + } +} + +impl Deref for MerklePath { + // we use `Vec` here instead of slice so we can call vector mutation methods directly from the + // merkle path (example: `Vec::remove`). + type Target = Vec; + + fn deref(&self) -> &Self::Target { + &self.nodes + } +} + +impl DerefMut for MerklePath { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.nodes + } +} + +impl FromIterator for MerklePath { + fn from_iter>(iter: T) -> Self { + Self::new(iter.into_iter().collect()) + } +} + +impl IntoIterator for MerklePath { + type Item = Word; + type IntoIter = vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.nodes.into_iter() + } +} diff --git a/src/merkle/merkle_path_set.rs b/src/merkle/path_set.rs similarity index 69% rename from src/merkle/merkle_path_set.rs rename to src/merkle/path_set.rs index 8210285b..216dff66 100644 --- a/src/merkle/merkle_path_set.rs +++ b/src/merkle/path_set.rs @@ -1,4 +1,4 @@ -use super::{BTreeMap, MerkleError, Rpo256, Vec, Word, ZERO}; +use super::{BTreeMap, MerkleError, MerklePath, Rpo256, Vec, Word, ZERO}; // MERKLE PATH SET // ================================================================================================ @@ -8,7 +8,7 @@ use super::{BTreeMap, MerkleError, Rpo256, Vec, Word, ZERO}; pub struct MerklePathSet { root: Word, total_depth: u32, - paths: BTreeMap>, + paths: BTreeMap, } impl MerklePathSet { @@ -18,7 +18,7 @@ impl MerklePathSet { /// Returns an empty MerklePathSet. pub fn new(depth: u32) -> Result { let root = [ZERO; 4]; - let paths = BTreeMap::>::new(); + let paths = BTreeMap::new(); Ok(Self { root, @@ -30,59 +30,15 @@ impl MerklePathSet { // PUBLIC ACCESSORS // -------------------------------------------------------------------------------------------- - /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters - /// specify the leaf node at which the path starts. - /// - /// # Errors - /// Returns an error if: - /// - The specified index is not valid in the context of this Merkle path set (i.e., the index - /// implies a greater depth than is specified for this set). - /// - The specified path is not consistent with other paths in the set (i.e., resolves to a - /// different root). - pub fn add_path( - &mut self, - index: u64, - value: Word, - path: Vec, - ) -> Result<(), MerkleError> { - let depth = (path.len() + 1) as u32; - if depth != self.total_depth { - return Err(MerkleError::InvalidDepth(self.total_depth, depth)); - } - - // Actual number of node in tree - let pos = 2u64.pow(self.total_depth) + index; - - // Index of the leaf path in map. Paths of neighboring leaves are stored in one key-value pair - let half_pos = pos / 2; - - let mut extended_path = path; - if is_even(pos) { - extended_path.insert(0, value); - } else { - extended_path.insert(1, value); - } - - let root_of_current_path = compute_path_root(&extended_path, depth, index); - if self.root == [ZERO; 4] { - self.root = root_of_current_path; - } else if self.root != root_of_current_path { - return Err(MerkleError::InvalidPath(extended_path)); - } - self.paths.insert(half_pos, extended_path); - - Ok(()) - } - /// Returns the root to which all paths in this set resolve. - pub fn root(&self) -> Word { + pub const fn root(&self) -> Word { self.root } /// Returns the depth of the Merkle tree implied by the paths stored in this set. /// /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc. - pub fn depth(&self) -> u32 { + pub const fn depth(&self) -> u32 { self.total_depth } @@ -106,7 +62,7 @@ impl MerklePathSet { match self.paths.get(&index) { None => Err(MerkleError::NodeNotInSet(index)), Some(path) => { - if is_even(pos) { + if Self::is_even(pos) { Ok(path[0]) } else { Ok(path[1]) @@ -122,7 +78,7 @@ impl MerklePathSet { /// Returns an error if: /// * The specified index not valid for the depth of structure. /// * Node of the requested path does not exist in the set. - pub fn get_path(&self, depth: u32, index: u64) -> Result, MerkleError> { + pub fn get_path(&self, depth: u32, index: u64) -> Result { if index >= 2u64.pow(self.total_depth) { return Err(MerkleError::InvalidIndex(self.total_depth, index)); } @@ -137,7 +93,7 @@ impl MerklePathSet { None => Err(MerkleError::NodeNotInSet(index)), Some(path) => { let mut local_path = path.clone(); - if is_even(pos) { + if Self::is_even(pos) { local_path.remove(0); Ok(local_path) } else { @@ -148,6 +104,53 @@ impl MerklePathSet { } } + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Adds the specified Merkle path to this [MerklePathSet]. The `index` and `value` parameters + /// specify the leaf node at which the path starts. + /// + /// # Errors + /// Returns an error if: + /// - The specified index is not valid in the context of this Merkle path set (i.e., the index + /// implies a greater depth than is specified for this set). + /// - The specified path is not consistent with other paths in the set (i.e., resolves to a + /// different root). + pub fn add_path( + &mut self, + index: u64, + value: Word, + path: MerklePath, + ) -> Result<(), MerkleError> { + let depth = (path.len() + 1) as u32; + if depth != self.total_depth { + return Err(MerkleError::InvalidDepth(self.total_depth, depth)); + } + + // Actual number of node in tree + let pos = 2u64.pow(self.total_depth) + index; + + // Index of the leaf path in map. Paths of neighboring leaves are stored in one key-value pair + let half_pos = pos / 2; + + let mut extended_path = path; + if Self::is_even(pos) { + extended_path.insert(0, value); + } else { + extended_path.insert(1, value); + } + + let root_of_current_path = Self::compute_path_root(&extended_path, depth, index); + if self.root == [ZERO; 4] { + self.root = root_of_current_path; + } else if self.root != root_of_current_path { + return Err(MerkleError::InvalidPath(extended_path)); + } + self.paths.insert(half_pos, extended_path); + + Ok(()) + } + /// Replaces the leaf at the specified index with the provided value. /// /// # Errors @@ -166,16 +169,16 @@ impl MerklePathSet { }; // Fill old_hashes vector ----------------------------------------------------------------- - let (old_hashes, _) = compute_path_trace(path, depth, index); + let (old_hashes, _) = Self::compute_path_trace(path, depth, index); // Fill new_hashes vector ----------------------------------------------------------------- - if is_even(pos) { + if Self::is_even(pos) { path[0] = value; } else { path[1] = value; } - let (new_hashes, new_root) = compute_path_trace(path, depth, index); + let (new_hashes, new_root) = Self::compute_path_trace(path, depth, index); self.root = new_root; // update paths --------------------------------------------------------------------------- @@ -190,59 +193,59 @@ impl MerklePathSet { Ok(()) } -} -// HELPER FUNCTIONS -// -------------------------------------------------------------------------------------------- - -fn is_even(pos: u64) -> bool { - pos & 1 == 0 -} + // HELPER FUNCTIONS + // -------------------------------------------------------------------------------------------- -/// Calculates the hash of the parent node by two sibling ones -/// - node — current node -/// - node_pos — position of the current node -/// - sibling — neighboring vertex in the tree -fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word { - if is_even(node_pos) { - Rpo256::merge(&[node.into(), sibling.into()]).into() - } else { - Rpo256::merge(&[sibling.into(), node.into()]).into() + const fn is_even(pos: u64) -> bool { + pos & 1 == 0 } -} - -/// Returns vector of hashes from current to the root -fn compute_path_trace(path: &[Word], depth: u32, index: u64) -> (Vec, Word) { - let mut pos = 2u64.pow(depth) + index; - let mut computed_hashes = Vec::::new(); + /// Returns hash of the root + fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word { + let mut pos = 2u64.pow(depth) + index; - let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); + // hash that is obtained after calculating the current hash and path hash + let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); - if path.len() != 2 { for path_hash in path.iter().skip(2) { - computed_hashes.push(comp_hash); pos /= 2; - comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash); + comp_hash = Self::calculate_parent_hash(comp_hash, pos, *path_hash); } + + comp_hash } - (computed_hashes, comp_hash) -} + /// Calculates the hash of the parent node by two sibling ones + /// - node — current node + /// - node_pos — position of the current node + /// - sibling — neighboring vertex in the tree + fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word { + if Self::is_even(node_pos) { + Rpo256::merge(&[node.into(), sibling.into()]).into() + } else { + Rpo256::merge(&[sibling.into(), node.into()]).into() + } + } -/// Returns hash of the root -fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word { - let mut pos = 2u64.pow(depth) + index; + /// Returns vector of hashes from current to the root + fn compute_path_trace(path: &[Word], depth: u32, index: u64) -> (MerklePath, Word) { + let mut pos = 2u64.pow(depth) + index; - // hash that is obtained after calculating the current hash and path hash - let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); + let mut computed_hashes = Vec::::new(); - for path_hash in path.iter().skip(2) { - pos /= 2; - comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash); - } + let mut comp_hash = Rpo256::merge(&[path[0].into(), path[1].into()]).into(); + + if path.len() != 2 { + for path_hash in path.iter().skip(2) { + computed_hashes.push(comp_hash); + pos /= 2; + comp_hash = Self::calculate_parent_hash(comp_hash, pos, *path_hash); + } + } - comp_hash + (computed_hashes.into(), comp_hash) + } } // TESTS @@ -250,7 +253,7 @@ fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word { #[cfg(test)] mod tests { - use super::calculate_parent_hash; + use super::*; use crate::merkle::int_to_node; #[test] @@ -260,14 +263,14 @@ mod tests { let leaf2 = int_to_node(2); let leaf3 = int_to_node(3); - let parent0 = calculate_parent_hash(leaf0, 0, leaf1); - let parent1 = calculate_parent_hash(leaf2, 2, leaf3); + let parent0 = MerklePathSet::calculate_parent_hash(leaf0, 0, leaf1); + let parent1 = MerklePathSet::calculate_parent_hash(leaf2, 2, leaf3); - let root_exp = calculate_parent_hash(parent0, 0, parent1); + let root_exp = MerklePathSet::calculate_parent_hash(parent0, 0, parent1); let mut set = super::MerklePathSet::new(3).unwrap(); - set.add_path(0, leaf0, vec![leaf1, parent1]).unwrap(); + set.add_path(0, leaf0, vec![leaf1, parent1].into()).unwrap(); assert_eq!(set.root(), root_exp); } @@ -280,10 +283,10 @@ mod tests { let depth = 4u32; let mut set = super::MerklePathSet::new(depth).unwrap(); - set.add_path(index, hash_6, path_6.clone()).unwrap(); + set.add_path(index, hash_6, path_6.clone().into()).unwrap(); let stored_path_6 = set.get_path(depth, index).unwrap(); - assert_eq!(path_6, stored_path_6); + assert_eq!(path_6, *stored_path_6); assert!(set.get_path(depth, 15u64).is_err()) } @@ -295,7 +298,7 @@ mod tests { let depth = 4u32; let mut set = super::MerklePathSet::new(depth).unwrap(); - set.add_path(index, hash_6, path_6).unwrap(); + set.add_path(index, hash_6, path_6.into()).unwrap(); assert_eq!(int_to_node(6u64), set.get_node(depth, index).unwrap()); assert!(set.get_node(depth, 15u64).is_err()); @@ -307,8 +310,8 @@ mod tests { let hash_5 = int_to_node(5); let hash_6 = int_to_node(6); let hash_7 = int_to_node(7); - let hash_45 = calculate_parent_hash(hash_4, 12u64, hash_5); - let hash_67 = calculate_parent_hash(hash_6, 14u64, hash_7); + let hash_45 = MerklePathSet::calculate_parent_hash(hash_4, 12u64, hash_5); + let hash_67 = MerklePathSet::calculate_parent_hash(hash_6, 14u64, hash_7); let hash_0123 = int_to_node(123); @@ -322,22 +325,22 @@ mod tests { let depth = 4u32; let mut set = super::MerklePathSet::new(depth).unwrap(); - set.add_path(index_6, hash_6, path_6).unwrap(); - set.add_path(index_5, hash_5, path_5).unwrap(); - set.add_path(index_4, hash_4, path_4).unwrap(); + set.add_path(index_6, hash_6, path_6.into()).unwrap(); + set.add_path(index_5, hash_5, path_5.into()).unwrap(); + set.add_path(index_4, hash_4, path_4.into()).unwrap(); let new_hash_6 = int_to_node(100); let new_hash_5 = int_to_node(55); set.update_leaf(index_6, new_hash_6).unwrap(); let new_path_4 = set.get_path(depth, index_4).unwrap(); - let new_hash_67 = calculate_parent_hash(new_hash_6, 14u64, hash_7); + let new_hash_67 = MerklePathSet::calculate_parent_hash(new_hash_6, 14u64, hash_7); assert_eq!(new_hash_67, new_path_4[1]); set.update_leaf(index_5, new_hash_5).unwrap(); let new_path_4 = set.get_path(depth, index_4).unwrap(); let new_path_6 = set.get_path(depth, index_6).unwrap(); - let new_hash_45 = calculate_parent_hash(new_hash_5, 13u64, hash_4); + let new_hash_45 = MerklePathSet::calculate_parent_hash(new_hash_5, 13u64, hash_4); assert_eq!(new_hash_45, new_path_6[1]); assert_eq!(new_hash_5, new_path_4[0]); } diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 821b1a6f..07453d7b 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -1,4 +1,4 @@ -use super::{BTreeMap, MerkleError, Rpo256, RpoDigest, Vec, Word}; +use super::{BTreeMap, MerkleError, MerklePath, Rpo256, RpoDigest, Vec, Word}; #[cfg(test)] mod tests; @@ -102,7 +102,7 @@ impl SimpleSmt { /// Returns an error if: /// * The specified key does not exist as a branch or leaf node /// * The specified depth is greater than the depth of the tree. - pub fn get_path(&self, depth: u32, key: u64) -> Result, MerkleError> { + pub fn get_path(&self, depth: u32, key: u64) -> Result { if depth == 0 { return Err(MerkleError::DepthTooSmall(depth)); } else if depth > self.depth() { @@ -124,7 +124,7 @@ impl SimpleSmt { path.push(sibling_node.into()); curr_key >>= 1; } - Ok(path) + Ok(path.into()) } /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not @@ -133,7 +133,7 @@ impl SimpleSmt { /// # Errors /// Returns an error if: /// * The specified key does not exist as a leaf node. - pub fn get_leaf_path(&self, key: u64) -> Result, MerkleError> { + pub fn get_leaf_path(&self, key: u64) -> Result { self.get_path(self.depth(), key) } diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 7042d1b6..595d0218 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -103,14 +103,14 @@ fn get_path() { let (_, node2, node3) = compute_internal_nodes(); // check depth 2 - assert_eq!(vec![VALUES4[1], node3], tree.get_path(2, 0).unwrap()); - assert_eq!(vec![VALUES4[0], node3], tree.get_path(2, 1).unwrap()); - assert_eq!(vec![VALUES4[3], node2], tree.get_path(2, 2).unwrap()); - assert_eq!(vec![VALUES4[2], node2], tree.get_path(2, 3).unwrap()); + assert_eq!(vec![VALUES4[1], node3], *tree.get_path(2, 0).unwrap()); + assert_eq!(vec![VALUES4[0], node3], *tree.get_path(2, 1).unwrap()); + assert_eq!(vec![VALUES4[3], node2], *tree.get_path(2, 2).unwrap()); + assert_eq!(vec![VALUES4[2], node2], *tree.get_path(2, 3).unwrap()); // check depth 1 - assert_eq!(vec![node3], tree.get_path(1, 0).unwrap()); - assert_eq!(vec![node2], tree.get_path(1, 1).unwrap()); + assert_eq!(vec![node3], *tree.get_path(1, 0).unwrap()); + assert_eq!(vec![node2], *tree.get_path(1, 1).unwrap()); } #[test] @@ -191,7 +191,7 @@ fn small_tree_opening_is_consistent() { for (depth, key, path) in cases { let opening = tree.get_path(depth, key).unwrap(); - assert_eq!(path, opening); + assert_eq!(path, *opening); } }