diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 5eff1c3..7adbccb 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 87cd80f..014d4c7 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 0000000..2b81bbe --- /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 8210285..216dff6 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 821b1a6..07453d7 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 7042d1b..595d021 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); } }