diff --git a/crypto/Cargo.toml b/crypto/Cargo.toml new file mode 100644 index 0000000..87560ab --- /dev/null +++ b/crypto/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "crypto" +version = "0.1.0" +description="Miden Cryptographic primitives" +authors = ["miden contributors"] +readme="README.md" +license = "MIT" +repository = "https://github.com/0xPolygonMiden/crypto" +categories = ["cryptography", "no-std"] +keywords = ["miden", "crypto", "hash", "merkle"] +edition = "2021" + +[dependencies] +winterfell = { git = "https://github.com/novifinancial/winterfell"} +winter_utils = { version = "0.4", package = "winter-utils" } + +[dev-dependencies] +proptest = "1.0.0" +rand_utils = { version = "0.4", package = "winter-rand-utils" } \ No newline at end of file diff --git a/crypto/src/hash/mod.rs b/crypto/src/hash/mod.rs new file mode 100644 index 0000000..367f415 --- /dev/null +++ b/crypto/src/hash/mod.rs @@ -0,0 +1,26 @@ +use crate::StarkField; +pub use winterfell::crypto::hashers::Rp64_256 as Hasher; +pub use winterfell::crypto::{ElementHasher, Hasher as HashFn}; + +// TYPE ALIASES +// ================================================================================================ + +pub type Digest = ::Digest; + +// HELPER FUNCTIONS +// ================================================================================================ + +#[inline(always)] +fn exp_acc(base: [B; N], tail: [B; N]) -> [B; N] { + let mut result = base; + for _ in 0..M { + result.iter_mut().for_each(|r| *r = r.square()); + } + result.iter_mut().zip(tail).for_each(|(r, t)| *r *= t); + result +} + +#[inline(always)] +pub fn merge(values: &[Digest; 2]) -> Digest { + Hasher::merge(values) +} diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs new file mode 100644 index 0000000..ec50e8b --- /dev/null +++ b/crypto/src/lib.rs @@ -0,0 +1,12 @@ +pub use winterfell::math::{ + fields::{f64::BaseElement as Felt, QuadExtension}, + ExtensionOf, FieldElement, StarkField, +}; + +pub(crate) mod hash; +pub(crate) mod merkle; + +// TYPE ALIASES +// ================================================================================================ + +pub type Word = [Felt; 4]; diff --git a/crypto/src/merkle/merkle_path_set.rs b/crypto/src/merkle/merkle_path_set.rs new file mode 100644 index 0000000..8010676 --- /dev/null +++ b/crypto/src/merkle/merkle_path_set.rs @@ -0,0 +1,354 @@ +use super::{Felt, FieldElement, MerkleError, Word}; +use crate::hash::merge; +use winter_utils::collections::{BTreeMap, Vec}; + +// MERKLE PATH SET +// ================================================================================================ + +/// A set of Merkle paths. +/// +/// This struct is intended to be used as one of the variants of the MerkleSet enum. +#[derive(Clone, Debug)] +pub struct MerklePathSet { + root: Word, + total_depth: u32, + paths: BTreeMap>, +} + +impl MerklePathSet { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + + /// Returns an empty MerklePathSet. + pub fn new(depth: u32) -> Result { + let root = [Felt::ZERO; 4]; + let paths = BTreeMap::>::new(); + + Ok(Self { + root, + total_depth: depth, + paths, + }) + } + + // 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) as u64; + + 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 == [Felt::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 { + 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 { + self.total_depth + } + + /// Returns a node at the specified index. + /// + /// # Errors + /// Returns an error if: + /// * The specified index not valid for the depth of structure. + /// * Requested node does not exist in the set. + pub fn get_node(&self, depth: u32, index: u64) -> Result { + if index >= 2u64.pow(self.total_depth) { + return Err(MerkleError::InvalidIndex(self.total_depth, index)); + } + if depth != self.total_depth { + return Err(MerkleError::InvalidDepth(self.total_depth, depth)); + } + + let pos = 2u64.pow(depth) + index; + let index = (pos / 2) as u64; + + match self.paths.get(&index) { + None => Err(MerkleError::NodeNotInSet(index)), + Some(path) => { + if is_even(pos) { + Ok(path[0]) + } else { + Ok(path[1]) + } + } + } + } + + /// Returns a Merkle path to the node at the specified index. The node itself is + /// not included in the path. + /// + /// # Errors + /// 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> { + if index >= 2u64.pow(self.total_depth) { + return Err(MerkleError::InvalidIndex(self.total_depth, index)); + } + if depth != self.total_depth { + return Err(MerkleError::InvalidDepth(self.total_depth, depth)); + } + + let pos = 2u64.pow(depth) + index; + let index = pos / 2; + + match self.paths.get(&index) { + None => Err(MerkleError::NodeNotInSet(index)), + Some(path) => { + let mut local_path = path.clone(); + if is_even(pos) { + local_path.remove(0); + Ok(local_path) + } else { + local_path.remove(1); + Ok(local_path) + } + } + } + } + + /// Replaces the leaf at the specified index with the provided value. + /// + /// # Errors + /// Returns an error if: + /// * Requested node does not exist in the set. + pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<(), MerkleError> { + let depth = self.depth(); + if index >= 2u64.pow(depth) { + return Err(MerkleError::InvalidIndex(depth, index)); + } + let pos = 2u64.pow(depth) + index; + + let path = match self.paths.get_mut(&(pos / 2)) { + None => return Err(MerkleError::NodeNotInSet(index)), + Some(path) => path, + }; + + // Fill old_hashes vector ----------------------------------------------------------------- + let (old_hashes, _) = compute_path_trace(path, depth, index); + + // Fill new_hashes vector ----------------------------------------------------------------- + if is_even(pos) { + path[0] = value; + } else { + path[1] = value; + } + + let (new_hashes, new_root) = compute_path_trace(path, depth, index); + self.root = new_root; + + // update paths --------------------------------------------------------------------------- + for path in self.paths.values_mut() { + for i in (0..old_hashes.len()).rev() { + if path[i + 2] == old_hashes[i] { + path[i + 2] = new_hashes[i]; + break; + } + } + } + + Ok(()) + } +} + +// HELPER FUNCTIONS +// -------------------------------------------------------------------------------------------- + +fn is_even(pos: u64) -> bool { + pos & 1 == 0 +} + +/// 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) { + merge(&[node.into(), sibling.into()]).into() + } else { + merge(&[sibling.into(), node.into()]).into() + } +} + +/// 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(); + + let mut comp_hash = 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); + } + } + + (computed_hashes, comp_hash) +} + +/// Returns hash of the root +fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word { + let mut pos = 2u64.pow(depth) + index; + + // hash that is obtained after calculating the current hash and path hash + let mut comp_hash = merge(&[path[0].into(), path[1].into()]).into(); + + for path_hash in path.iter().skip(2) { + pos /= 2; + comp_hash = calculate_parent_hash(comp_hash, pos, *path_hash); + } + + comp_hash +} + +// TESTS +// ================================================================================================ + +#[cfg(test)] +mod tests { + use super::{calculate_parent_hash, Felt, FieldElement, Word}; + + #[test] + fn get_root() { + let leaf0 = int_to_node(0); + let leaf1 = int_to_node(1); + 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 root_exp = calculate_parent_hash(parent0, 0, parent1); + + let mut set = super::MerklePathSet::new(3).unwrap(); + + set.add_path(0, leaf0, vec![leaf1, parent1]).unwrap(); + + assert_eq!(set.root(), root_exp); + } + + #[test] + fn add_and_get_path() { + let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)]; + let hash_6 = int_to_node(6); + let index = 6u64; + let depth = 4u32; + let mut set = super::MerklePathSet::new(depth).unwrap(); + + set.add_path(index, hash_6, path_6.clone()).unwrap(); + let stored_path_6 = set.get_path(depth, index).unwrap(); + + assert_eq!(path_6, stored_path_6); + assert!(set.get_path(depth, 15u64).is_err()) + } + + #[test] + fn get_node() { + let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)]; + let hash_6 = int_to_node(6); + let index = 6u64; + let depth = 4u32; + let mut set = super::MerklePathSet::new(depth).unwrap(); + + set.add_path(index, hash_6, path_6).unwrap(); + + assert_eq!(int_to_node(6u64), set.get_node(depth, index).unwrap()); + assert!(set.get_node(depth, 15u64).is_err()); + } + + #[test] + fn update_leaf() { + let hash_4 = int_to_node(4); + 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_0123 = int_to_node(123); + + let path_6 = vec![hash_7, hash_45, hash_0123]; + let path_5 = vec![hash_4, hash_67, hash_0123]; + let path_4 = vec![hash_5, hash_67, hash_0123]; + + let index_6 = 6u64; + let index_5 = 5u64; + let index_4 = 4u64; + 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(); + + 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); + 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); + assert_eq!(new_hash_45, new_path_6[1]); + assert_eq!(new_hash_5, new_path_4[0]); + } + + // HELPER FUNCTIONS + // -------------------------------------------------------------------------------------------- + + const fn int_to_node(value: u64) -> Word { + [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] + } +} diff --git a/crypto/src/merkle/merkle_tree.rs b/crypto/src/merkle/merkle_tree.rs new file mode 100644 index 0000000..d784d30 --- /dev/null +++ b/crypto/src/merkle/merkle_tree.rs @@ -0,0 +1,260 @@ +use super::MerkleError; +use crate::{ + hash::{merge, Digest}, + Felt, FieldElement, Word, +}; +use core::slice; +use winter_utils::uninit_vector; +use winterfell::math::log2; + +// MERKLE TREE +// ================================================================================================ + +/// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two). +/// +/// This struct is intended to be used as one of the variants of the MerkleSet enum. +#[derive(Clone, Debug)] +pub struct MerkleTree { + nodes: Vec, +} + +impl MerkleTree { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + /// Returns a Merkle tree instantiated from the provided leaves. + /// + /// # Errors + /// Returns an error if the number of leaves is smaller than two or is not a power of two. + pub fn new(leaves: Vec) -> Result { + let n = leaves.len(); + if n <= 1 { + return Err(MerkleError::DepthTooSmall); + } else if !n.is_power_of_two() { + return Err(MerkleError::NumLeavesNotPowerOfTwo(n)); + } + + // create un-initialized vector to hold all tree nodes + let mut nodes = unsafe { uninit_vector(2 * n) }; + nodes[0] = [Felt::ZERO; 4]; + + // copy leaves into the second part of the nodes vector + nodes[n..].copy_from_slice(&leaves); + + // re-interpret nodes as an array of two nodes fused together + let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [Digest; 2], n) }; + + // calculate all internal tree nodes + for i in (1..n).rev() { + nodes[i] = merge(&two_nodes[i]).into(); + } + + Ok(Self { nodes }) + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the root of this Merkle tree. + pub fn root(&self) -> Word { + self.nodes[1] + } + + /// Returns the depth of this Merkle tree. + /// + /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc. + pub fn depth(&self) -> u32 { + log2(self.nodes.len() / 2) + } + + /// Returns a node at the specified depth and index. + /// + /// # Errors + /// 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_node(&self, depth: u32, index: u64) -> Result { + if depth == 0 { + return Err(MerkleError::DepthTooSmall); + } else if depth > self.depth() { + return Err(MerkleError::DepthTooBig(depth)); + } + if index >= 2u64.pow(depth) { + return Err(MerkleError::InvalidIndex(depth, index)); + } + + let pos = 2usize.pow(depth as u32) + (index as usize); + Ok(self.nodes[pos]) + } + + /// Returns a Merkle path to the node at the specified depth and index. The node itself is + /// not included in the path. + /// + /// # Errors + /// 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> { + if depth == 0 { + return Err(MerkleError::DepthTooSmall); + } else if depth > self.depth() { + return Err(MerkleError::DepthTooBig(depth)); + } + if index >= 2u64.pow(depth) { + return Err(MerkleError::InvalidIndex(depth, index)); + } + + let mut path = Vec::with_capacity(depth as usize); + let mut pos = 2usize.pow(depth as u32) + (index as usize); + + while pos > 1 { + path.push(self.nodes[pos ^ 1]); + pos >>= 1; + } + + Ok(path) + } + + /// Replaces the leaf at the specified index with the provided value. + /// + /// # Errors + /// Returns an error if the specified index is not a valid leaf index for this tree. + pub fn update_leaf(&mut self, index: u64, value: Word) -> Result<(), MerkleError> { + let depth = self.depth(); + if index >= 2u64.pow(depth) { + return Err(MerkleError::InvalidIndex(depth, index)); + } + + let mut index = 2usize.pow(depth) + index as usize; + self.nodes[index] = value; + + let n = self.nodes.len() / 2; + let two_nodes = + unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [Digest; 2], n) }; + + for _ in 0..depth { + index /= 2; + self.nodes[index] = merge(&two_nodes[index]).into(); + } + + Ok(()) + } +} + +// TESTS +// ================================================================================================ + +#[cfg(test)] +mod tests { + use super::{Felt, FieldElement, Word}; + use crate::hash::{ElementHasher, HashFn, Hasher}; + + const LEAVES4: [Word; 4] = [ + int_to_node(1), + int_to_node(2), + int_to_node(3), + int_to_node(4), + ]; + + const LEAVES8: [Word; 8] = [ + int_to_node(1), + int_to_node(2), + int_to_node(3), + int_to_node(4), + int_to_node(5), + int_to_node(6), + int_to_node(7), + int_to_node(8), + ]; + + #[test] + fn build_merkle_tree() { + let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + assert_eq!(8, tree.nodes.len()); + + // leaves were copied correctly + for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) { + assert_eq!(a, b); + } + + let (root, node2, node3) = compute_internal_nodes(); + + assert_eq!(root, tree.nodes[1]); + assert_eq!(node2, tree.nodes[2]); + assert_eq!(node3, tree.nodes[3]); + + assert_eq!(root, tree.root()); + } + + #[test] + fn get_leaf() { + let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + + // check depth 2 + assert_eq!(LEAVES4[0], tree.get_node(2, 0).unwrap()); + assert_eq!(LEAVES4[1], tree.get_node(2, 1).unwrap()); + assert_eq!(LEAVES4[2], tree.get_node(2, 2).unwrap()); + assert_eq!(LEAVES4[3], tree.get_node(2, 3).unwrap()); + + // check depth 1 + let (_, node2, node3) = compute_internal_nodes(); + + assert_eq!(node2, tree.get_node(1, 0).unwrap()); + assert_eq!(node3, tree.get_node(1, 1).unwrap()); + } + + #[test] + fn get_path() { + let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + + 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()); + + // check depth 1 + assert_eq!(vec![node3], tree.get_path(1, 0).unwrap()); + assert_eq!(vec![node2], tree.get_path(1, 1).unwrap()); + } + + #[test] + fn update_leaf() { + let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap(); + + // update one leaf + let index = 3; + let new_node = int_to_node(9); + let mut expected_leaves = LEAVES8.to_vec(); + expected_leaves[index as usize] = new_node; + let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap(); + + tree.update_leaf(index, new_node).unwrap(); + assert_eq!(expected_tree.nodes, tree.nodes); + + // update another leaf + let index = 6; + let new_node = int_to_node(10); + expected_leaves[index as usize] = new_node; + let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap(); + + tree.update_leaf(index, new_node).unwrap(); + assert_eq!(expected_tree.nodes, tree.nodes); + } + + // HELPER FUNCTIONS + // -------------------------------------------------------------------------------------------- + + fn compute_internal_nodes() -> (Word, Word, Word) { + let node2 = Hasher::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat()); + let node3 = Hasher::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat()); + let root = Hasher::merge(&[node2, node3]); + + (root.into(), node2.into(), node3.into()) + } + + const fn int_to_node(value: u64) -> Word { + [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] + } +} diff --git a/crypto/src/merkle/mod.rs b/crypto/src/merkle/mod.rs new file mode 100644 index 0000000..290064f --- /dev/null +++ b/crypto/src/merkle/mod.rs @@ -0,0 +1,18 @@ +use crate::{Felt, FieldElement, Word}; + +pub mod merkle_path_set; +pub mod merkle_tree; + +// ERRORS +// ================================================================================================ + +#[derive(Clone, Debug)] +pub enum MerkleError { + DepthTooSmall, + DepthTooBig(u32), + NumLeavesNotPowerOfTwo(usize), + InvalidIndex(u32, u64), + InvalidDepth(u32, u32), + InvalidPath(Vec), + NodeNotInSet(u64), +}