WIP refactor of SMT

src/sdk/src/crypto/mod.rs

@@ -64,6 +64,7 @@ pub mod ecvrf;
 
 /// Sparse Merkle Tree implementation
 pub mod smt;
+pub mod smt2;
 
 /// Convenience module to import all the pasta traits.
 /// You still have to import the curves.

src/sdk/src/crypto/smt2/mod.rs

@@ -0,0 +1,274 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2024 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)]`
+//!
+//! # Terminology
+//!
+//! * **level** - the depth in the tree. Type: `u32`
+//! * **location** - a `(level, position)` tuple
+//! * **position** - the leaf index, or equivalently the binary direction through the tree
+//!   with type `F`.
+//! * **index** - the internal index used in the DB which is `BigUint`. Leaf node indexes are
+//!   calculated as `leaf_idx = final_level_start_idx + position`.
+//! * **node** - either the leaf values or parent nodes `hash(left, right)`.
+
+#[cfg(test)]
+mod test;
+
+mod util;
+pub use util::Poseidon;
+
+use num_bigint::BigUint;
+use std::collections::HashMap;
+
+use util::{FieldElement, FieldHasher};
+
+/// Pluggable storage backend for the SMT.
+/// Has a minimal interface to simply put and get objects from the store.
+pub trait StorageAdapter {
+    type Value;
+
+    fn put(&mut self, key: BigUint, value: Self::Value);
+    fn get(&self, key: &BigUint) -> Option<Self::Value>;
+}
+
+/// An in-memory storage, useful for unit tests and smaller trees.
+pub struct MemoryStorage<F: FieldElement> {
+    tree: HashMap<BigUint, F>,
+}
+
+impl<F: FieldElement> MemoryStorage<F> {
+    fn new() -> Self {
+        Self { tree: HashMap::new() }
+    }
+}
+
+impl<F: FieldElement> StorageAdapter for MemoryStorage<F> {
+    type Value = F;
+
+    fn put(&mut self, key: BigUint, value: F) {
+        self.tree.insert(key, value);
+    }
+    fn get(&self, key: &BigUint) -> Option<F> {
+        self.tree.get(key).copied()
+    }
+}
+
+/// 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.
+///
+/// The trait param `N` is the depth of the tree. A tree with a depth of `N`
+/// will have `N + 1` levels.
+#[derive(Debug)]
+pub struct SparseMerkleTree<
+    const N: usize,
+    F: FieldElement,
+    H: FieldHasher<F, 2>,
+    S: StorageAdapter<Value = F>,
+> where
+    [(); N + 1]:,
+{
+    /// A map from leaf indices to leaf data stored as field elements.
+    store: S,
+    /// The hasher used to build the Merkle tree.
+    hasher: H,
+    /// An array of empty hashes hashed with themselves `N` times.
+    empty_nodes: [F; N + 1],
+}
+
+impl<const N: usize, F: FieldElement, H: FieldHasher<F, 2>, S: StorageAdapter<Value = F>>
+    SparseMerkleTree<N, F, H, S>
+where
+    [(); N + 1]:,
+{
+    /// Creates a new SMT
+    pub fn new(store: S, hasher: H, empty_leaf: F) -> Self {
+        let empty_nodes = gen_empty_nodes(&hasher, empty_leaf);
+
+        Self { store, hasher, empty_nodes }
+    }
+
+    /// Takes a batch of field elements, inserts these hashes into the tree,
+    /// and updates the Merkle root.
+    pub fn insert_batch(&mut self, leaves: Vec<(F, F)>) {
+        // Nodes that need recalculating
+        let mut dirty_idxs = Vec::new();
+        for (pos, leaf) in leaves {
+            let idx = util::leaf_pos_to_index::<N, _>(&pos);
+            self.store.put(idx.clone(), leaf);
+
+            // Mark node parent as dirty
+            let parent_idx = util::parent(&idx).unwrap();
+            dirty_idxs.push(parent_idx);
+        }
+
+        // Depth first from the bottom of the tree
+        for _ in 0..N + 1 {
+            let mut new_dirty_idxs = Vec::new();
+
+            for idx in dirty_idxs {
+                let left_idx = util::left_child(&idx);
+                let right_idx = util::right_child(&idx);
+                let left = self.get_node(&left_idx);
+                let right = self.get_node(&right_idx);
+                // Recalculate the node
+                let node = self.hasher.hash([left, right]);
+
+                self.store.put(idx.clone(), node);
+
+                // Add this node's parent to the update list
+                let parent_idx = match util::parent(&idx) {
+                    Some(idx) => idx,
+                    // We are at the root node so no parents exist
+                    None => break,
+                };
+                new_dirty_idxs.push(parent_idx);
+            }
+
+            dirty_idxs = new_dirty_idxs;
+        }
+    }
+
+    /// Returns the Merkle tree root.
+    pub fn root(&self) -> F {
+        self.get_node(&BigUint::from(0u32))
+    }
+
+    /// 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 prove_membership(&self, pos: &F) -> Path<N, F, H> {
+        let mut path = [F::ZERO; N];
+        let leaf_idx = util::leaf_pos_to_index::<N, _>(pos);
+
+        let mut current_idx = leaf_idx;
+        // Depth first from the bottom of the tree
+        for lvl in (0..N).rev() {
+            let sibling_idx = util::sibling(&current_idx).unwrap();
+            let sibling_node = self.get_node(&sibling_idx);
+            path[lvl] = sibling_node;
+
+            // Now move to the parent
+            current_idx = util::parent(&current_idx).unwrap();
+        }
+
+        Path { path, hasher: self.hasher.clone() }
+    }
+
+    /// Fast lookup for leaf. The SMT can be used as a generic container for
+    /// objects with very little overhead using this method.
+    pub fn get_leaf(&self, pos: &F) -> F {
+        let leaf_idx = util::leaf_pos_to_index::<N, _>(pos);
+        self.get_node(&leaf_idx)
+    }
+
+    fn get_node(&self, idx: &BigUint) -> F {
+        let lvl = util::log2(&idx);
+        let empty_node = self.empty_nodes[lvl as usize];
+        self.store.get(&idx).unwrap_or(empty_node)
+    }
+}
+
+/// The path contains a sequence of sibling nodes that make up a Merkle proof.
+/// Each sibling node is used to identify whether the merkle root construction
+/// is valid at the root.
+pub struct Path<const N: usize, F: FieldElement, H: FieldHasher<F, 2>> {
+    /// Path from leaf to root. It is a list of sibling nodes.
+    /// It does not contain the root node.
+    /// Similar to other conventions here, the list starts higher in the tree
+    /// and goes down. So when iterating we start from the end.
+    pub path: [F; N],
+    hasher: H,
+}
+
+impl<const N: usize, F: FieldElement, H: FieldHasher<F, 2>> Path<N, F, H> {
+    pub fn verify(&self, root: &F, leaf: &F, pos: &F) -> bool {
+        let pos = pos.as_biguint();
+        assert!(pos.bits() as usize <= N);
+
+        let mut current_node = *leaf;
+        for i in (0..N).rev() {
+            let sibling_node = self.path[i];
+
+            let is_right = pos.bit((N - 1 - i) as u64);
+            //println!("is_right: {}", is_right);
+            let (left, right) =
+                if is_right { (sibling_node, current_node) } else { (current_node, sibling_node) };
+            //println!("left: {:?}, right: {:?}", left, right);
+
+            current_node = self.hasher.hash([left, right]);
+        }
+
+        current_node == *root
+    }
+}
+
+/// 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.
+///
+/// Ordering is depth-wise starting from root going down.
+pub fn gen_empty_nodes<const N: usize, F: FieldElement, H: FieldHasher<F, 2>>(
+    hasher: &H,
+    empty_leaf: F,
+) -> [F; N + 1] {
+    let mut empty_nodes = [F::ZERO; N + 1];
+
+    let mut empty_node = empty_leaf;
+    for item in empty_nodes.iter_mut().rev() {
+        *item = empty_node;
+        empty_node = hasher.hash([empty_node, empty_node]);
+    }
+
+    empty_nodes
+}

src/sdk/src/crypto/smt2/test.rs

@@ -0,0 +1,132 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2024 Dyne.org foundation
+ *
+ * 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/>.
+ */
+
+use super::*;
+use halo2_proofs::arithmetic::Field;
+use pasta_curves::Fp;
+use rand::rngs::OsRng;
+
+#[test]
+fn empties() {
+    let hasher = Poseidon::<Fp, 2>::new();
+    let empty_leaf = Fp::from(0);
+    let empty_nodes = gen_empty_nodes::<3, _, _>(&hasher, empty_leaf);
+
+    let empty_node1 = hasher.hash([empty_leaf, empty_leaf]);
+    let empty_node2 = hasher.hash([empty_node1, empty_node1]);
+    let empty_root = hasher.hash([empty_node2, empty_node2]);
+
+    assert_eq!(empty_nodes[3], empty_leaf);
+    assert_eq!(empty_nodes[2], empty_node1);
+    assert_eq!(empty_nodes[1], empty_node2);
+    assert_eq!(empty_nodes[0], empty_root);
+}
+
+#[test]
+fn poseidon_smt() {
+    const HEIGHT: usize = 3;
+    let hasher = Poseidon::<Fp, 2>::new();
+    let empty_leaf = Fp::from(0);
+
+    let store = MemoryStorage::<Fp>::new();
+    let mut smt =
+        SparseMerkleTree::<HEIGHT, _, _, _>::new(store, hasher.clone(), empty_leaf.clone());
+
+    // Both reprs should match
+    assert_eq!(Fp::from(1).as_biguint(), BigUint::from(1u32));
+    assert_eq!(Fp::from(300).as_biguint(), BigUint::from(300u32));
+
+    let leaves = vec![
+        (Fp::from(1), Fp::random(&mut OsRng)),
+        (Fp::from(2), Fp::random(&mut OsRng)),
+        (Fp::from(3), Fp::random(&mut OsRng)),
+    ];
+    smt.insert_batch(leaves.clone());
+
+    let empty_nodes = gen_empty_nodes::<HEIGHT, _, _>(&hasher, empty_leaf);
+
+    let hash1 = leaves[0].1;
+    let hash2 = leaves[1].1;
+    let hash3 = leaves[2].1;
+
+    let hash = |l, r| hasher.hash([l, r]);
+
+    let hash01 = hash(empty_nodes[3], hash1);
+    let hash23 = hash(hash2, hash3);
+
+    let hash0123 = hash(hash01, hash23);
+    let root = hash(hash0123, empty_nodes[1]);
+    assert_eq!(root, smt.root());
+
+    //println!("hash1: {:?}", hash1);
+    //println!("hash2: {:?}", hash2);
+    //println!("hash3: {:?}", hash3);
+    //println!("hash4-7: {:?}", empty_nodes[3]);
+    //println!();
+    //println!("hash01: {:?}", hash01);
+    //println!("hash23: {:?}", hash23);
+    //println!("hash45: {:?}", empty_nodes[2]);
+    //println!("hash67: {:?}", empty_nodes[2]);
+    //println!();
+    //println!("hash0123: {:?}", hash0123);
+    //println!("hash4567: {:?}", empty_nodes[1]);
+    //println!();
+    //println!("root: {:?}", root);
+    //println!();
+
+    // Now try to construct a membership proof for leaf 3
+    let pos = leaves[2].0;
+    let path = smt.prove_membership(&pos);
+    assert_eq!(path.path[0], empty_nodes[1]);
+    assert_eq!(path.path[1], hash01);
+    assert_eq!(path.path[2], hash2);
+
+    assert_eq!(hash23, hash(path.path[2], hash3));
+    assert_eq!(hash0123, hash(path.path[1], hash(path.path[2], hash3)));
+    assert_eq!(root, hash(hash(path.path[1], hash(path.path[2], hash3)), path.path[0]));
+
+    //println!("path0: {:?}", path.path[0]);
+    //println!("path1: {:?}", path.path[1]);
+    //println!("path2: {:?}", path.path[2]);
+
+    assert!(path.verify(&root, &hash3, &pos));
+}
+
+#[test]
+fn poseidon_smt_incl_proof() {
+    const HEIGHT: usize = 3;
+    let hasher = Poseidon::<Fp, 2>::new();
+    let empty_leaf = Fp::from(0);
+
+    let store = MemoryStorage::<Fp>::new();
+    let mut smt =
+        SparseMerkleTree::<HEIGHT, _, _, _>::new(store, hasher.clone(), empty_leaf.clone());
+
+    let leaves = vec![
+        (Fp::from(1), Fp::random(&mut OsRng)),
+        (Fp::from(2), Fp::random(&mut OsRng)),
+        (Fp::from(3), Fp::random(&mut OsRng)),
+    ];
+    smt.insert_batch(leaves.clone());
+
+    let (pos, leaf) = leaves[2];
+    assert_eq!(smt.get_leaf(&pos), leaf);
+
+    let path = smt.prove_membership(&pos);
+    assert!(path.verify(&smt.root(), &leaf, &pos));
+}

src/sdk/src/crypto/smt2/util.rs

@@ -0,0 +1,132 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2024 Dyne.org foundation
+ *
+ * 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/>.
+ */
+
+use halo2_gadgets::poseidon::{
+    primitives as poseidon,
+    primitives::{ConstantLength, P128Pow5T3, Spec},
+};
+use num_bigint::BigUint;
+use pasta_curves::group::ff::{PrimeField, WithSmallOrderMulGroup};
+use std::marker::PhantomData;
+
+pub trait FieldElement: WithSmallOrderMulGroup<3> + Ord + PrimeField {
+    fn as_biguint(&self) -> BigUint;
+}
+impl FieldElement for pasta_curves::Fp {
+    fn as_biguint(&self) -> BigUint {
+        let repr = self.to_repr();
+        BigUint::from_bytes_le(&repr)
+    }
+}
+impl FieldElement for pasta_curves::Fq {
+    fn as_biguint(&self) -> BigUint {
+        let repr = self.to_repr();
+        BigUint::from_bytes_le(&repr)
+    }
+}
+
+pub trait FieldHasher<F: WithSmallOrderMulGroup<3> + Ord, const L: usize>: Clone {
+    fn hash(&self, inputs: [F; L]) -> F;
+    fn hasher() -> Self;
+}
+
+#[derive(Debug, Clone)]
+pub struct Poseidon<F: WithSmallOrderMulGroup<3> + Ord, const L: usize>(PhantomData<F>);
+
+impl<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> Poseidon<F, L> {
+    pub fn new() -> Self {
+        Poseidon(PhantomData)
+    }
+}
+
+impl<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> Default for Poseidon<F, L> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> FieldHasher<F, L> for Poseidon<F, L>
+where
+    P128Pow5T3: Spec<F, 3, 2>,
+{
+    fn hash(&self, inputs: [F; L]) -> F {
+        poseidon::Hash::<_, P128Pow5T3, ConstantLength<L>, 3, 2>::init().hash(inputs)
+    }
+
+    fn hasher() -> Self {
+        Poseidon(PhantomData)
+    }
+}
+
+#[inline]
+/// Converts a leaf position to the internal BigUint index for storage.
+pub(super) fn leaf_pos_to_index<const N: usize, F: FieldElement>(pos: &F) -> BigUint {
+    // Starting index for the last level
+    // 2^N - 1
+    let final_level_index = (BigUint::from(1u32) << (N as u64)) - 1u32;
+
+    final_level_index.clone() + pos.as_biguint()
+}
+
+/// Returns the log2 value of the given number. Used for converting the index to the level.
+#[inline]
+pub(super) fn log2(x: &BigUint) -> u64 {
+    (x + 1u32).bits() - 1
+}
+
+/// Returns the index of the left child, given an index.
+#[inline]
+pub(super) fn left_child(index: &BigUint) -> BigUint {
+    2u32 * index + 1u32
+}
+
+/// Returns the index of the right child, given an index.
+#[inline]
+pub(super) fn right_child(index: &BigUint) -> BigUint {
+    2u32 * index + 2u32
+}
+
+/// Returns true iff the given index represents a left child.
+#[inline]
+pub(super) fn is_left_child(index: &BigUint) -> bool {
+    // Any simple way to convert the (index % 2) into a u32 rather
+    // than converting 1 into a BigUint?
+    index % 2u32 == 1u32.into()
+}
+
+/// Returns the index of the parent, given an index.
+#[inline]
+pub(super) fn parent(index: &BigUint) -> Option<BigUint> {
+    if *index > 0u32.into() {
+        Some((index - 1u32) >> 1)
+    } else {
+        None
+    }
+}
+
+/// Returns the index of the sibling, given an index.
+#[inline]
+pub(super) fn sibling(index: &BigUint) -> Option<BigUint> {
+    if *index == 0u32.into() {
+        None
+    } else if is_left_child(index) {
+        Some(index + 1u32)
+    } else {
+        Some(index - 1u32)
+    }
+}

src/sdk/src/lib.rs

@@ -15,6 +15,7 @@
  * 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/>.
  */
+#![feature(generic_const_exprs)]
 
 pub use bridgetree;
 pub use num_bigint;