Merge pull request #1 from 0xPolygonMiden/merkle

Implement basic Merkle tree-related primitives

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" }

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 = <Hasher as HashFn>::Digest;
+
+// HELPER FUNCTIONS
+// ================================================================================================
+
+#[inline(always)]
+fn exp_acc<B: StarkField, const N: usize, const M: usize>(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)
+}

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];

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<u64, Vec<Word>>,
+}
+
+impl MerklePathSet {
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns an empty MerklePathSet.
+    pub fn new(depth: u32) -> Result<Self, MerkleError> {
+        let root = [Felt::ZERO; 4];
+        let paths = BTreeMap::<u64, Vec<Word>>::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<Word>,
+    ) -> 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<Word, 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) 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<Vec<Word>, 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>, Word) {
+    let mut pos = 2u64.pow(depth) + index;
+
+    let mut computed_hashes = Vec::<Word>::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]
+    }
+}

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<Word>,
+}
+
+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<Word>) -> Result<Self, MerkleError> {
+        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<Word, 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 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<Vec<Word>, 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]
+    }
+}

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<Word>),
+    NodeNotInSet(u64),
+}