sdk: Sparse Merkle tree impl

src/sdk/src/crypto/mod.rs

@@ -68,6 +68,9 @@ pub mod schnorr;
 /// MiMC VDF
 pub mod mimc_vdf;
 
+/// Sparse Merkle Tree implementation
+pub mod smt;
+
 pub use incrementalmerkletree;
 pub use pasta_curves::{pallas, vesta};
 

src/sdk/src/crypto/smt.rs

@@ -0,0 +1,418 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2023 Dyne.org foundation
+ *
+ * Copyright (C) 2021 Webb Technologies Inc.
+ * Copyright (c) zkMove Authors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+//! This file provides a native implementation of the Sparse Merkle tree data
+//! structure.
+//!
+//! A Sparse Merkle tree is a type of Merkle tree, but it is much easier to
+//! prove non-membership in a sparse Merkle tree than in an arbitrary Merkle
+//! tree. For an explanation of sparse Merkle trees, see:
+//! `<https://medium.com/@kelvinfichter/whats-a-sparse-merkle-tree-acda70aeb837>`
+//!
+//! In this file we define the `Path` and `SparseMerkleTree` structs.
+//! These depend on your choice of a prime field F, a field hasher over F
+//! (any hash function that maps F^2 to F will do, e.g. the poseidon hash
+//! function of width 3 where an input of zero is used for padding), and the
+//! height N of the sparse Merkle tree.
+//!
+//! The path corresponding to a given leaf node is stored as an N-tuple of pairs
+//! of field elements. Each pair consists of a node lying on the path from the
+//! leaf node to the root, and that node's sibling.  For example, suppose
+//! ```text
+//!           a
+//!         /   \
+//!        b     c
+//!       / \   / \
+//!      d   e f   g
+//! ```
+//! is our Sparse Merkle tree, and `a` through `g` are field elements stored at
+//! the nodes. Then the merkle proof path `e-b-a` from leaf `e` to root `a` is
+//! stored as `[(d,e), (b,c)]`
+
+use core::marker::PhantomData;
+use std::collections::{BTreeMap, BTreeSet};
+
+use halo2_gadgets::poseidon::{
+    primitives as poseidon,
+    primitives::{ConstantLength, P128Pow5T3},
+};
+use pasta_curves::{arithmetic::FieldExt, Fp};
+
+use crate::error::{ContractError, GenericResult};
+
+pub trait FieldHasher<const L: usize> {
+    fn hash(&self, inputs: [Fp; L]) -> GenericResult<Fp>;
+    fn hasher() -> Self;
+}
+
+#[derive(Debug, Clone)]
+pub struct Poseidon<const L: usize>(PhantomData<Fp>);
+
+impl<const L: usize> Poseidon<L> {
+    pub fn new() -> Self {
+        Poseidon(PhantomData::default())
+    }
+}
+
+impl<const L: usize> FieldHasher<L> for Poseidon<L> {
+    fn hash(&self, inputs: [Fp; L]) -> GenericResult<Fp> {
+        Ok(poseidon::Hash::<_, P128Pow5T3, ConstantLength<L>, 3, 2>::init().hash(inputs))
+    }
+
+    fn hasher() -> Self {
+        Poseidon(PhantomData::default())
+    }
+}
+
+/// The Path struct.
+///
+/// The path contains a sequence of sibling nodes that make up a Merkle proof.
+/// Each pair is used to identify whether an incremental Merkle root construction
+/// is valid at each intermediate step.
+#[derive(Debug, Copy, Clone)]
+pub struct Path<H: FieldHasher<2>, const N: usize> {
+    /// The path represented as a sequence of sibling pairs.
+    pub path: [(Fp, Fp); N],
+    /// The phantom hasher type used to reconstruct the Merkle root.
+    pub marker: PhantomData<H>,
+}
+
+impl<H: FieldHasher<2>, const N: usize> Path<H, N> {
+    /// Assumes leaf contains leaf-level data, i.e. hashes of secrets stored on
+    /// leaf-level.
+    pub fn calculate_root(&self, leaf: &Fp, hasher: &H) -> GenericResult<Fp> {
+        if *leaf != self.path[0].0 && *leaf != self.path[0].1 {
+            return Err(ContractError::SmtInvalidLeaf)
+        }
+
+        let mut prev = *leaf;
+        // Check levels between leaf level and root
+        for &(ref left_hash, ref right_hash) in &self.path {
+            if &prev != left_hash && &prev != right_hash {
+                return Err(ContractError::SmtInvalidPathNodes)
+            }
+            prev = hasher.hash([*left_hash, *right_hash])?;
+        }
+
+        Ok(prev)
+    }
+
+    /// Takes in an expected `root_hash` and leaf-level data (i.e. hashes of secrets)
+    /// for a leaf and checks that the leaf belongs to a tree having the expected hash.
+    pub fn check_membership(&self, root_hash: &Fp, leaf: &Fp, hasher: &H) -> GenericResult<bool> {
+        let root = self.calculate_root(leaf, hasher)?;
+        Ok(root == *root_hash)
+    }
+
+    /// Given leaf data, determine what the index of this leaf must be in the
+    /// Merkle tree it belongs to. Before doing so, check that the leaf does
+    /// indeed belong to a tree with the given `root_hash`.
+    pub fn get_index(&self, root_hash: &Fp, leaf: &Fp, hasher: &H) -> GenericResult<Fp> {
+        if !self.check_membership(root_hash, leaf, hasher)? {
+            return Err(ContractError::SmtInvalidLeaf)
+        }
+
+        let mut prev = *leaf;
+        let mut index = Fp::zero();
+        let mut twopower = Fp::one();
+        // Check levels between leaf level and root
+        for &(ref left_hash, ref right_hash) in &self.path {
+            // Check if the previous hash is for a left or right ndoe
+            if &prev != left_hash {
+                index += twopower;
+            }
+
+            twopower = twopower + twopower;
+            prev = hasher.hash([*left_hash, *right_hash])?;
+        }
+
+        Ok(index)
+    }
+}
+
+/// The Sparse Merkle Tree struct.
+///
+/// SMT stores a set of leaves represented in a map and a set of empty
+/// hashes that it uses to represent the sparse areas of the tree.
+pub struct SparseMerkleTree<H: FieldHasher<2>, const N: usize> {
+    /// A map from leaf indices to leaf data stored as field elements.
+    pub tree: BTreeMap<u64, Fp>,
+    /// An array of default hashes hashed with themselves `N` times.
+    empty_hashes: [Fp; N],
+    /// The phantom hasher type used to build the Merkle tree.
+    marker: PhantomData<H>,
+}
+
+impl<H: FieldHasher<2>, const N: usize> SparseMerkleTree<H, N> {
+    /// Creates a new SMT from a map of indices to field elements.
+    pub fn new(
+        leaves: &BTreeMap<u32, Fp>,
+        hasher: &H,
+        empty_leaf: &[u8; 64],
+    ) -> GenericResult<Self> {
+        // Ensure the tree can hold this many leaves
+        let last_level_size = leaves.len().next_power_of_two();
+        let tree_size = 2 * last_level_size - 1;
+        let tree_height = tree_height(tree_size as u64);
+        assert!(tree_height <= N as u32);
+
+        // Initialize the Merkle tree
+        let tree = BTreeMap::new();
+        let empty_hashes = gen_empty_hashes(hasher, empty_leaf)?;
+
+        let mut smt = SparseMerkleTree::<H, N> { tree, empty_hashes, marker: PhantomData };
+
+        smt.insert_batch(leaves, hasher)?;
+
+        Ok(smt)
+    }
+
+    /// Creates a new SMT from an array of field elements.
+    pub fn new_sequential(leaves: &[Fp], hasher: &H, empty_leaf: &[u8; 64]) -> GenericResult<Self> {
+        let pairs: BTreeMap<u32, Fp> =
+            leaves.iter().enumerate().map(|(i, l)| (i as u32, *l)).collect();
+
+        let smt = Self::new(&pairs, hasher, empty_leaf)?;
+
+        Ok(smt)
+    }
+
+    /// Takes a batch of field elements, inserts these hashes into the tree,
+    /// and updates the Merkle root.
+    pub fn insert_batch(&mut self, leaves: &BTreeMap<u32, Fp>, hasher: &H) -> GenericResult<()> {
+        let last_level_index: u64 = (1u64 << N) - 1;
+
+        let mut level_idxs: BTreeSet<u64> = BTreeSet::new();
+        for (i, leaf) in leaves {
+            let true_index = last_level_index + (*i as u64);
+            self.tree.insert(true_index, *leaf);
+            level_idxs.insert(parent(true_index).unwrap());
+        }
+
+        for level in 0..N {
+            let mut new_idxs: BTreeSet<u64> = BTreeSet::new();
+            let empty_hash = self.empty_hashes[level];
+            for i in level_idxs {
+                let left_index = left_child(i);
+                let right_index = right_child(i);
+                let left = self.tree.get(&left_index).unwrap_or(&empty_hash);
+                let right = self.tree.get(&right_index).unwrap_or(&empty_hash);
+                self.tree.insert(i, hasher.hash([*left, *right])?);
+
+                let parent = match parent(i) {
+                    Some(i) => i,
+                    None => break,
+                };
+
+                new_idxs.insert(parent);
+            }
+
+            level_idxs = new_idxs;
+        }
+
+        Ok(())
+    }
+
+    /// Returns the Merkle tree root.
+    pub fn root(&self) -> Fp {
+        self.tree.get(&0).cloned().unwrap_or(*self.empty_hashes.last().unwrap())
+    }
+
+    /// Give the path leading from the leaf at `index` up to the root. This is
+    /// a "proof" in the sense of "valid path in a Merkle tree", not a ZK argument.
+    pub fn generate_membership_proof(&self, index: u64) -> Path<H, N> {
+        let mut path = [(Fp::zero(), Fp::zero()); N];
+
+        let tree_index = convert_index_to_last_level(index, N);
+
+        // Iterate from the leaf up to the root, storing all intermediate hash values.
+        let mut current_node = tree_index;
+        let mut level = 0;
+        while !is_root(current_node) {
+            let sibling_node = sibling(current_node).unwrap();
+
+            let empty_hash = &self.empty_hashes[level];
+
+            let current = self.tree.get(&current_node).cloned().unwrap_or(*empty_hash);
+            let sibling = self.tree.get(&sibling_node).cloned().unwrap_or(*empty_hash);
+
+            if is_left_child(current_node) {
+                path[level] = (current, sibling);
+            } else {
+                path[level] = (sibling, current);
+            }
+
+            current_node = parent(current_node).unwrap();
+            level += 1;
+        }
+
+        Path { path, marker: PhantomData }
+    }
+}
+
+/// A function to generate empty hashes with a given `default_leaf`.
+///
+/// Given a `FieldHasher`, generate a list of `N` hashes consisting of the
+/// `default_leaf` hashed with itself and repeated `N` times with the
+/// intermediate results. These are used to initialize the sparse portion
+/// of the SMT.
+pub fn gen_empty_hashes<H: FieldHasher<2>, const N: usize>(
+    hasher: &H,
+    default_leaf: &[u8; 64],
+) -> GenericResult<[Fp; N]> {
+    let mut empty_hashes = [Fp::zero(); N];
+
+    let mut empty_hash = Fp::from_bytes_wide(default_leaf);
+    for item in empty_hashes.iter_mut().take(N) {
+        *item = empty_hash;
+        empty_hash = hasher.hash([empty_hash, empty_hash])?;
+    }
+
+    Ok(empty_hashes)
+}
+
+fn convert_index_to_last_level(index: u64, height: usize) -> u64 {
+    index + (1u64 << height) - 1
+}
+
+/// Returns the log2 value of the given number.
+#[inline]
+fn log2(x: u64) -> u32 {
+    if x == 0 {
+        0
+    } else if x.is_power_of_two() {
+        1usize.leading_zeros() - x.leading_zeros()
+    } else {
+        0usize.leading_zeros() - x.leading_zeros()
+    }
+}
+
+/// Returns the index of the left child, given an index.
+#[inline]
+fn left_child(index: u64) -> u64 {
+    2 * index + 1
+}
+
+/// Returns the index of the right child, given an index.
+#[inline]
+fn right_child(index: u64) -> u64 {
+    2 * index + 2
+}
+
+/// Returns true iff the given index represents a left child.
+#[inline]
+fn is_left_child(index: u64) -> bool {
+    index % 2 == 1
+}
+
+/// Returns the index of the parent, given an index.
+#[inline]
+fn parent(index: u64) -> Option<u64> {
+    if index > 0 {
+        Some((index - 1) >> 1)
+    } else {
+        None
+    }
+}
+
+/// Returns the index of the sibling, given an index.
+#[inline]
+fn sibling(index: u64) -> Option<u64> {
+    if index == 0 {
+        None
+    } else if is_left_child(index) {
+        Some(index + 1)
+    } else {
+        Some(index - 1)
+    }
+}
+
+/// Returns the height of the tree, given the size of the tree.
+#[inline]
+fn tree_height(tree_size: u64) -> u32 {
+    log2(tree_size)
+}
+
+/// Returns true iff the index represents the Merkle root.
+#[inline]
+fn is_root(index: u64) -> bool {
+    index == 0
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use halo2_proofs::arithmetic::Field;
+    use rand::rngs::OsRng;
+
+    /// Helper to change leaves array to BTreeMap and then create SMT.
+    fn create_merkle_tree<H: FieldHasher<2>, const N: usize>(
+        hasher: H,
+        leaves: &[Fp],
+        default_leaf: &[u8; 64],
+    ) -> SparseMerkleTree<H, N> {
+        SparseMerkleTree::<H, N>::new_sequential(leaves, &hasher, default_leaf).unwrap()
+    }
+
+    #[test]
+    fn poseidon_smt() {
+        let poseidon = Poseidon::<2>::new();
+        let default_leaf = [0u8; 64];
+        let leaves = [Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
+        const HEIGHT: usize = 3;
+
+        let smt =
+            create_merkle_tree::<Poseidon<2>, HEIGHT>(poseidon.clone(), &leaves, &default_leaf);
+
+        let root = smt.root();
+
+        let empty_hashes =
+            gen_empty_hashes::<Poseidon<2>, HEIGHT>(&poseidon, &default_leaf).unwrap();
+
+        let hash1 = leaves[0];
+        let hash2 = leaves[1];
+        let hash3 = leaves[2];
+
+        let hash12 = poseidon.hash([hash1, hash2]).unwrap();
+        let hash34 = poseidon.hash([hash3, empty_hashes[0]]).unwrap();
+
+        let hash1234 = poseidon.hash([hash12, hash34]).unwrap();
+        let calc_root = poseidon.hash([hash1234, empty_hashes[2]]).unwrap();
+
+        assert_eq!(root, calc_root);
+    }
+
+    #[test]
+    fn poseidon_smt_incl_proof() {
+        let poseidon = Poseidon::<2>::new();
+        let default_leaf = [0u8; 64];
+        let leaves = [Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
+        const HEIGHT: usize = 3;
+
+        let smt =
+            create_merkle_tree::<Poseidon<2>, HEIGHT>(poseidon.clone(), &leaves, &default_leaf);
+
+        let proof = smt.generate_membership_proof(0);
+        let res = proof.check_membership(&smt.root(), &leaves[0], &poseidon).unwrap();
+        assert!(res)
+    }
+}

src/sdk/src/error.rs

@@ -74,6 +74,12 @@ pub enum ContractError {
 
     #[error("Invalid function call")]
     InvalidFunction,
+
+    #[error("SMT: Invalid Leaf")]
+    SmtInvalidLeaf,
+
+    #[error("SMT: Path nodes are not consistent")]
+    SmtInvalidPathNodes,
 }
 
 /// Builtin return values occupy the upper 32 bits
@@ -99,6 +105,8 @@ pub const DB_GET_FAILED: i64 = to_builtin!(13);
 pub const DB_CONTAINS_KEY_FAILED: i64 = to_builtin!(14);
 pub const INVALID_FUNCTION: i64 = to_builtin!(15);
 pub const DB_DEL_FAILED: i64 = to_builtin!(16);
+pub const SMT_INVALID_LEAF: i64 = to_builtin!(17);
+pub const SMT_INVALID_PATH_NODES: i64 = to_builtin!(18);
 
 impl From<ContractError> for i64 {
     fn from(err: ContractError) -> Self {
@@ -118,6 +126,8 @@ impl From<ContractError> for i64 {
             ContractError::DbContainsKeyFailed => DB_CONTAINS_KEY_FAILED,
             ContractError::InvalidFunction => INVALID_FUNCTION,
             ContractError::DbDelFailed => DB_DEL_FAILED,
+            ContractError::SmtInvalidLeaf => SMT_INVALID_LEAF,
+            ContractError::SmtInvalidPathNodes => SMT_INVALID_PATH_NODES,
             ContractError::Custom(error) => {
                 if error == 0 {
                     CUSTOM_ZERO
@@ -148,6 +158,8 @@ impl From<i64> for ContractError {
             DB_CONTAINS_KEY_FAILED => Self::DbContainsKeyFailed,
             INVALID_FUNCTION => Self::InvalidFunction,
             DB_DEL_FAILED => Self::DbDelFailed,
+            SMT_INVALID_LEAF => Self::SmtInvalidLeaf,
+            SMT_INVALID_PATH_NODES => Self::SmtInvalidPathNodes,
             _ => Self::Custom(error as u32),
         }
     }