Merge branch 'smt2'

proof/smt.zk

@@ -1,4 +1,4 @@
-k = 13;
+k = 14;
 field = "pallas";
 
 constant "SMT" {
@@ -11,7 +11,7 @@ witness "SMT" {
 }
 
 circuit "SMT" {
-    is_member = sparse_tree_is_member(root, path, leaf);
+    is_member = sparse_tree_is_member(root, path, leaf, leaf);
 
     ONE = witness_base(1);
     constrain_equal_base(is_member, ONE);

src/sdk/src/crypto/smt.rs

@@ -1,454 +0,0 @@
-/* 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)`.
-
-use core::marker::PhantomData;
-use std::collections::{BTreeMap, BTreeSet};
-
-use halo2_gadgets::poseidon::{
-    primitives as poseidon,
-    primitives::{ConstantLength, P128Pow5T3, Spec},
-};
-use pasta_curves::group::ff::{PrimeField, WithSmallOrderMulGroup};
-
-use crate::error::{ContractError, GenericResult};
-
-pub trait FieldElement: WithSmallOrderMulGroup<3> + Ord + PrimeField {}
-impl FieldElement for pasta_curves::Fp {}
-impl FieldElement for pasta_curves::Fq {}
-
-pub trait FieldHasher<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> {
-    fn hash(&self, inputs: [F; L]) -> GenericResult<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]) -> GenericResult<F> {
-        Ok(poseidon::Hash::<_, P128Pow5T3, ConstantLength<L>, 3, 2>::init().hash(inputs))
-    }
-
-    fn hasher() -> Self {
-        Poseidon(PhantomData)
-    }
-}
-
-/// 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<F: WithSmallOrderMulGroup<3> + Ord, H: FieldHasher<F, 2>, const N: usize> {
-    /// The path represented as a sequence of sibling pairs.
-    pub path: [(F, F); N],
-    /// The phantom hasher type used to reconstruct the Merkle root.
-    pub marker: PhantomData<H>,
-}
-
-impl<F: WithSmallOrderMulGroup<3> + Ord, H: FieldHasher<F, 2>, const N: usize> Path<F, H, N> {
-    /// Assumes leaf contains leaf-level data, i.e. hashes of secrets stored on
-    /// leaf-level.
-    pub fn calculate_root(&self, leaf: &F, hasher: &H) -> GenericResult<F> {
-        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 (left_hash, 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: &F, leaf: &F, 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: &F, leaf: &F, hasher: &H) -> GenericResult<F> {
-        if !self.check_membership(root_hash, leaf, hasher)? {
-            return Err(ContractError::SmtInvalidLeaf)
-        }
-
-        let mut prev = *leaf;
-        let mut index = F::ZERO;
-        let mut twopower = F::ONE;
-        // Check levels between leaf level and root
-        for (left_hash, 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.
-#[derive(Debug)]
-pub struct SparseMerkleTree<F: FieldElement, H: FieldHasher<F, 2>, const N: usize> {
-    /// A map from leaf indices to leaf data stored as field elements.
-    pub tree: BTreeMap<u64, F>,
-    /// An array of default hashes hashed with themselves `N` times.
-    empty_hashes: [F; N],
-    /// The phantom hasher type used to build the Merkle tree.
-    marker: PhantomData<H>,
-}
-
-impl<F: FieldElement, H: FieldHasher<F, 2>, const N: usize> SparseMerkleTree<F, H, N> {
-    /// Creates a new SMT from a map of indices to field elements.
-    pub fn new(leaves: &BTreeMap<u32, F>, hasher: &H, empty_leaf: F::Repr) -> 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::<F, 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: &[F], hasher: &H, empty_leaf: F::Repr) -> GenericResult<Self> {
-        let pairs: BTreeMap<u32, F> =
-            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, F>, 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) -> F {
-        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<F, H, N> {
-        let mut path = [(F::ZERO, F::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<
-    F: WithSmallOrderMulGroup<3> + Ord + PrimeField,
-    H: FieldHasher<F, 2>,
-    const N: usize,
->(
-    hasher: &H,
-    default_leaf: F::Repr,
-) -> GenericResult<[F; N]> {
-    let mut empty_hashes = [F::ZERO; N];
-
-    let empty_hash = F::from_repr(default_leaf);
-    let mut empty_hash = if empty_hash.is_some().into() {
-        empty_hash.unwrap()
-    } else {
-        return Err(ContractError::Internal)
-    };
-    for item in empty_hashes.iter_mut() {
-        *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 pasta_curves::Fp;
-    use rand::rngs::OsRng;
-
-    /// Helper to change leaves array to BTreeMap and then create SMT.
-    fn create_merkle_tree<F: FieldElement, H: FieldHasher<F, 2>, const N: usize>(
-        hasher: H,
-        leaves: &[F],
-        default_leaf: F::Repr,
-    ) -> SparseMerkleTree<F, H, N> {
-        SparseMerkleTree::<F, H, N>::new_sequential(leaves, &hasher, default_leaf).unwrap()
-    }
-
-    #[test]
-    fn poseidon_smt() {
-        let poseidon = Poseidon::<Fp, 2>::new();
-        let default_leaf = [0u8; 32];
-        let leaves = [Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
-        const HEIGHT: usize = 3;
-
-        let smt = create_merkle_tree::<Fp, Poseidon<Fp, 2>, HEIGHT>(
-            poseidon.clone(),
-            &leaves,
-            default_leaf.clone(),
-        );
-
-        let root = smt.root();
-
-        let empty_hashes =
-            gen_empty_hashes::<Fp, Poseidon<Fp, 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::<Fp, 2>::new();
-        let default_leaf = [0u8; 32];
-        let leaves = [Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
-        const HEIGHT: usize = 3;
-
-        let smt = create_merkle_tree::<Fp, Poseidon<Fp, 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/crypto/smt/mod.rs

@@ -0,0 +1,289 @@
+/* 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;
+// Only used for the type aliases below
+use pasta_curves::pallas;
+
+use util::{FieldElement, FieldHasher};
+
+// Bit size for Fp (and Fq)
+pub const SMT_FP_DEPTH: usize = 255;
+pub type PoseidonFp = Poseidon<pallas::Base, 2>;
+pub type MemoryStorageFp = MemoryStorage<pallas::Base>;
+pub type SmtMemoryFp =
+    SparseMerkleTree<SMT_FP_DEPTH, { SMT_FP_DEPTH + 1 }, pallas::Base, PoseidonFp, MemoryStorageFp>;
+pub type PathFp = Path<SMT_FP_DEPTH, pallas::Base, PoseidonFp>;
+
+/// 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> {
+    pub 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,
+    // M = N + 1
+    const M: usize,
+    F: FieldElement,
+    H: FieldHasher<F, 2>,
+    S: StorageAdapter<Value = F>,
+> {
+    /// 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; M],
+}
+
+impl<
+        const N: usize,
+        const M: usize,
+        F: FieldElement,
+        H: FieldHasher<F, 2>,
+        S: StorageAdapter<Value = F>,
+    > SparseMerkleTree<N, M, F, H, S>
+{
+    /// Creates a new SMT
+    pub fn new(store: S, hasher: H, empty_leaf: F) -> Self {
+        assert_eq!(M, N + 1);
+        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);
+            let (left, right) =
+                if is_right { (sibling_node, current_node) } else { (current_node, sibling_node) };
+            //println!("is_right: {}", is_right);
+            //println!("left: {:?}, right: {:?}", left, right);
+            //println!("current_node: {:?}", current_node);
+
+            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 M: usize, F: FieldElement, H: FieldHasher<F, 2>>(
+    hasher: &H,
+    empty_leaf: F,
+) -> [F; M] {
+    let mut empty_nodes = [F::ZERO; M];
+
+    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/smt/test.rs

@@ -0,0 +1,138 @@
+/* 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 + 1 }, _, _>(&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, { HEIGHT + 1 }, _, _, _>::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 + 1 }, _, _>(&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, { HEIGHT + 1 }, _, _, _>::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/smt/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/zk/gadget/smt.rs

@@ -1,34 +1,16 @@
-/* This file is part of DarkFi (https://dark.fi)
- *
- * Copyright (C) 2020-2024 Dyne.org foundation
- * Copyright (C) 2022 zkMove Authors (Apache-2.0)
- * Copyright (C) 2021 Webb Technologies Inc. (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/>.
- */
-
-use std::marker::PhantomData;
-
-use darkfi_sdk::crypto::smt::FieldHasher;
+use darkfi_sdk::crypto::smt::SMT_FP_DEPTH;
 use halo2_gadgets::poseidon::{
     primitives as poseidon, Hash as PoseidonHash, Pow5Chip as PoseidonChip,
     Pow5Config as PoseidonConfig,
 };
 use halo2_proofs::{
     circuit::{AssignedCell, Layouter, Value},
-    pasta::Fp,
-    plonk::{self, Advice, Column, ConstraintSystem, Selector},
+    pasta::{
+        group::ff::{Field, PrimeFieldBits},
+        Fp,
+    },
+    plonk::{self, Advice, Column, ConstraintSystem, Constraints, Selector},
+    poly::Rotation,
 };
 
 use super::{
@@ -37,16 +19,16 @@ use super::{
 };
 
 #[derive(Clone, Debug)]
-pub struct PathConfig<const N: usize> {
+pub struct PathConfig {
     s_path: Selector,
-    advices: [Column<Advice>; N],
+    advices: [Column<Advice>; 2],
     poseidon_config: PoseidonConfig<Fp, 3, 2>,
     is_eq_config: IsEqualConfig<Fp>,
     conditional_select_config: ConditionalSelectConfig<Fp>,
     assert_equal_config: AssertEqualConfig<Fp>,
 }
 
-impl<const N: usize> PathConfig<N> {
+impl PathConfig {
     fn poseidon_chip(&self) -> PoseidonChip<Fp, 3, 2> {
         PoseidonChip::construct(self.poseidon_config.clone())
     }
@@ -65,19 +47,17 @@ impl<const N: usize> PathConfig<N> {
 }
 
 #[derive(Clone, Debug)]
-pub struct PathChip<H: FieldHasher<Fp, 2>, const N: usize> {
-    path: [(AssignedCell<Fp, Fp>, AssignedCell<Fp, Fp>); N],
-    config: PathConfig<N>,
-    _hasher: PhantomData<H>,
+pub struct PathChip {
+    config: PathConfig,
 }
 
-impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
+impl PathChip {
     pub fn configure(
         meta: &mut ConstraintSystem<Fp>,
-        advices: [Column<Advice>; N],
+        advices: [Column<Advice>; 2],
         utility_advices: [Column<Advice>; NUM_OF_UTILITY_ADVICE_COLUMNS],
         poseidon_config: PoseidonConfig<Fp, 3, 2>,
-    ) -> PathConfig<N> {
+    ) -> PathConfig {
         let s_path = meta.selector();
 
         for advice in &advices {
@@ -88,6 +68,15 @@ impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
             meta.enable_equality(*advice);
         }
 
+        meta.create_gate("Path builder", |meta| {
+            let s_path = meta.query_selector(s_path);
+            let current_path = meta.query_advice(advices[0], Rotation::cur());
+            let bit = meta.query_advice(advices[1], Rotation::cur());
+            let next_path = meta.query_advice(advices[0], Rotation::next());
+
+            Constraints::with_selector(s_path, Some(next_path - (current_path * Fp::from(2) + bit)))
+        });
+
         PathConfig {
             s_path,
             advices,
@@ -101,76 +90,115 @@ impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
         }
     }
 
-    pub fn from_native(
-        config: PathConfig<N>,
-        layouter: &mut impl Layouter<Fp>,
-        path: [(Value<Fp>, Value<Fp>); N],
-    ) -> Result<Self, plonk::Error> {
-        let path = layouter.assign_region(
-            || "path",
-            |mut region| {
-                config.s_path.enable(&mut region, 0)?;
-                let left = (0..N)
-                    .map(|i| {
-                        region.assign_advice(
-                            || format!("path[{}][{}]", i, 0),
-                            config.advices[i],
-                            0,
-                            || path[i].0,
-                        )
-                    })
-                    .collect::<Result<Vec<AssignedCell<Fp, Fp>>, plonk::Error>>();
-
-                let right = (0..N)
-                    .map(|i| {
-                        region.assign_advice(
-                            || format!("path[{}][{}]", i, 1),
-                            config.advices[i],
-                            1,
-                            || path[i].1,
-                        )
-                    })
-                    .collect::<Result<Vec<AssignedCell<Fp, Fp>>, plonk::Error>>();
-
-                let result = left?
-                    .into_iter()
-                    .zip(right?.into_iter())
-                    .collect::<Vec<(AssignedCell<Fp, Fp>, AssignedCell<Fp, Fp>)>>();
-
-                Ok(result.try_into().unwrap())
-            },
-        )?;
+    pub fn construct(config: PathConfig) -> Self {
+        Self { config }
+    }
+
+    fn decompose_value(value: &Fp) -> Vec<Fp> {
+        // Returns 256 bits, but the last bit is uneeded
+        let bits: Vec<bool> = value.to_le_bits().into_iter().collect();
 
-        Ok(PathChip { path, config, _hasher: PhantomData })
+        let mut bits: Vec<Fp> = bits[..SMT_FP_DEPTH].iter().map(|x| Fp::from(*x)).collect();
+        bits.resize(SMT_FP_DEPTH, Fp::from(0));
+        bits
     }
 
-    pub fn calculate_root(
+    pub fn check_membership(
         &self,
         layouter: &mut impl Layouter<Fp>,
+        root: AssignedCell<Fp, Fp>,
         leaf: AssignedCell<Fp, Fp>,
+        pos: AssignedCell<Fp, Fp>,
+        path: Value<[Fp; SMT_FP_DEPTH]>,
     ) -> Result<AssignedCell<Fp, Fp>, plonk::Error> {
-        // Check levels between leaf level and root
-        let mut previous_hash = leaf;
+        let path = path.transpose_array();
+        // Witness values
+        let (bits, path, zero) = layouter.assign_region(
+            || "witness",
+            |mut region| {
+                let bits = pos.value().map(Self::decompose_value).transpose_vec(SMT_FP_DEPTH);
+                assert_eq!(bits.len(), SMT_FP_DEPTH);
+
+                let mut witness_bits = vec![];
+                let mut witness_path = vec![];
+                for (i, (bit, sibling)) in bits.into_iter().zip(path.into_iter()).enumerate() {
+                    let bit = region.assign_advice(
+                        || "witness root",
+                        self.config.advices[0],
+                        i,
+                        || bit,
+                    )?;
+                    witness_bits.push(bit);
+
+                    let sibling = region.assign_advice(
+                        || "witness root",
+                        self.config.advices[1],
+                        i,
+                        || sibling,
+                    )?;
+                    witness_path.push(sibling);
+                }
+
+                let zero = region.assign_advice(
+                    || "witness one",
+                    self.config.advices[0],
+                    SMT_FP_DEPTH,
+                    || Value::known(Fp::ZERO),
+                )?;
+                region.constrain_constant(zero.cell(), Fp::ZERO)?;
+
+                Ok((witness_bits, witness_path, zero))
+            },
+        )?;
+        assert_eq!(bits.len(), path.len());
+        assert_eq!(bits.len(), SMT_FP_DEPTH);
 
         let iseq_chip = self.config.is_eq_chip();
         let condselect_chip = self.config.conditional_select_chip();
         let asserteq_chip = self.config.assert_eq_chip();
 
-        for (left_hash, right_hash) in self.path.iter() {
-            // Check if previous_hash matches the correct current hash
-            let previous_is_left =
-                iseq_chip.is_eq_with_output(layouter, previous_hash.clone(), left_hash.clone())?;
+        // Check path construction
+        let mut current_path = zero;
+        for bit in bits.iter().rev() {
+            current_path = layouter.assign_region(
+                || "pᵢ₊₁ = 2pᵢ + bᵢ",
+                |mut region| {
+                    self.config.s_path.enable(&mut region, 0)?;
 
-            let left_or_right = condselect_chip.conditional_select(
-                layouter,
-                left_hash.clone(),
-                right_hash.clone(),
-                previous_is_left,
+                    current_path.copy_advice(
+                        || "current path",
+                        &mut region,
+                        self.config.advices[0],
+                        0,
+                    )?;
+                    bit.copy_advice(|| "path bit", &mut region, self.config.advices[1], 0)?;
+
+                    let next_path =
+                        current_path.value().zip(bit.value()).map(|(p, b)| p * Fp::from(2) + b);
+                    region.assign_advice(|| "next path", self.config.advices[0], 1, || next_path)
+                },
             )?;
+        }
 
-            asserteq_chip.assert_equal(layouter, previous_hash, left_or_right)?;
+        // Check tree construction
+        let mut current_node = leaf.clone();
+        for (bit, sibling) in bits.into_iter().zip(path.into_iter().rev()) {
+            // Conditional select also constraints the bit ∈ {0, 1}
+            let left = condselect_chip.conditional_select(
+                layouter,
+                sibling.clone(),
+                current_node.clone(),
+                bit.clone(),
+            )?;
+            let right = condselect_chip.conditional_select(
+                layouter,
+                current_node.clone(),
+                sibling,
+                bit.clone(),
+            )?;
+            //println!("bit: {:?}", bit);
+            //println!("left: {:?}, right: {:?}", left, right);
 
-            // Update previous_hash
             let hasher = PoseidonHash::<
                 _,
                 _,
@@ -183,71 +211,50 @@ impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
                 layouter.namespace(|| "SmtPoseidonHash init"),
             )?;
 
-            previous_hash = hasher.hash(
-                layouter.namespace(|| "SmtPoseidonHash hash"),
-                [left_hash.clone(), right_hash.clone()],
-            )?;
+            current_node =
+                hasher.hash(layouter.namespace(|| "SmtPoseidonHash hash"), [left, right])?;
         }
 
-        Ok(previous_hash)
-    }
-
-    pub fn check_membership(
-        &self,
-        layouter: &mut impl Layouter<Fp>,
-        root_hash: AssignedCell<Fp, Fp>,
-        leaf: AssignedCell<Fp, Fp>,
-    ) -> Result<AssignedCell<Fp, Fp>, plonk::Error> {
-        let computed_root = self.calculate_root(layouter, leaf)?;
+        asserteq_chip.assert_equal(layouter, current_path, pos)?;
 
-        self.config.is_eq_chip().is_eq_with_output(layouter, computed_root, root_hash)
+        iseq_chip.is_eq_with_output(layouter, current_node, root)
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-
-    use darkfi_sdk::crypto::smt::{Poseidon, SparseMerkleTree};
-    use halo2_proofs::{
-        arithmetic::Field, circuit::floor_planner, dev::MockProver, plonk::Circuit,
-    };
+    use darkfi_sdk::crypto::smt::{MemoryStorageFp, PoseidonFp, SmtMemoryFp};
+    use halo2_proofs::{circuit::floor_planner, dev::MockProver, plonk::Circuit};
     use rand::rngs::OsRng;
 
-    const HEIGHT: usize = 3;
-
     struct TestCircuit {
         root: Value<Fp>,
-        path: [(Value<Fp>, Value<Fp>); HEIGHT],
+        path: Value<[Fp; SMT_FP_DEPTH]>,
         leaf: Value<Fp>,
     }
 
     impl Circuit<Fp> for TestCircuit {
-        type Config = PathConfig<HEIGHT>;
+        type Config = PathConfig;
         type FloorPlanner = floor_planner::V1;
         type Params = ();
 
         fn without_witnesses(&self) -> Self {
-            Self { root: Value::unknown(), path: self.path.clone(), leaf: Value::unknown() }
+            Self { root: Value::unknown(), path: Value::unknown(), leaf: Value::unknown() }
         }
 
         fn configure(meta: &mut ConstraintSystem<Fp>) -> Self::Config {
-            let advices = [(); HEIGHT].map(|_| meta.advice_column());
+            // Advice wires required by PathChip
+            let advices = [(); 2].map(|_| meta.advice_column());
             let utility_advices = [(); NUM_OF_UTILITY_ADVICE_COLUMNS].map(|_| meta.advice_column());
-            let poseidon_advices = [(); 4].map(|_| meta.advice_column());
-
-            for advice in &advices {
-                meta.enable_equality(*advice);
-            }
-
-            for advice in &utility_advices {
-                meta.enable_equality(*advice);
-            }
 
+            // Setup poseidon config
+            let poseidon_advices = [(); 4].map(|_| meta.advice_column());
             for advice in &poseidon_advices {
                 meta.enable_equality(*advice);
             }
 
+            // Needed for poseidon hash
             let col_const = meta.fixed_column();
             meta.enable_constant(col_const);
 
@@ -262,12 +269,7 @@ mod tests {
                 rc_b,
             );
 
-            PathChip::<Poseidon<Fp, 2>, HEIGHT>::configure(
-                meta,
-                advices,
-                utility_advices,
-                poseidon_config,
-            )
+            PathChip::configure(meta, advices, utility_advices, poseidon_config)
         }
 
         fn synthesize(
@@ -276,8 +278,7 @@ mod tests {
             mut layouter: impl Layouter<Fp>,
         ) -> Result<(), plonk::Error> {
             // Initialize the Path chip
-            let path_chip: PathChip<Poseidon<Fp, 2>, HEIGHT> =
-                PathChip::from_native(config.clone(), &mut layouter, self.path.clone())?;
+            let path_chip = PathChip::construct(config.clone());
 
             // Initialize the AssertEqual chip
             let assert_eq_chip = config.assert_eq_chip();
@@ -286,11 +287,11 @@ mod tests {
             let (root, leaf, one) = layouter.assign_region(
                 || "witness",
                 |mut region| {
-                    let one = region.assign_advice(
-                        || "witness one",
-                        config.advices[2],
+                    let root = region.assign_advice(
+                        || "witness root",
+                        config.advices[0],
                         0,
-                        || Value::known(Fp::ONE),
+                        || self.root,
                     )?;
 
                     let leaf = region.assign_advice(
@@ -300,18 +301,25 @@ mod tests {
                         || self.leaf,
                     )?;
 
-                    let root = region.assign_advice(
-                        || "witness root",
-                        config.advices[0],
-                        0,
-                        || self.root,
+                    let one = region.assign_advice(
+                        || "witness one",
+                        config.advices[1],
+                        1,
+                        || Value::known(Fp::ONE),
                     )?;
+                    region.constrain_constant(one.cell(), Fp::ONE)?;
 
                     Ok((root, leaf, one))
                 },
             )?;
 
-            let is_valid = path_chip.check_membership(&mut layouter, root, leaf)?;
+            let is_valid = path_chip.check_membership(
+                &mut layouter,
+                root,
+                leaf.clone(),
+                leaf.clone(),
+                self.path,
+            )?;
             assert_eq_chip.assert_equal(&mut layouter, is_valid, one)?;
 
             Ok(())
@@ -320,31 +328,39 @@ mod tests {
 
     #[test]
     fn test_smt_circuit() {
-        let hasher = Poseidon::<Fp, 2>::hasher();
-        let leaves: [Fp; HEIGHT] = [(); HEIGHT].map(|_| Fp::random(&mut OsRng));
-        let empty_leaf = [0u8; 32];
-
-        let smt = SparseMerkleTree::<Fp, Poseidon<Fp, 2>, HEIGHT>::new_sequential(
-            &leaves,
-            &hasher.clone(),
-            empty_leaf,
-        )
-        .unwrap();
-
-        //println!("{:#?}", smt);
-
-        let path = smt.generate_membership_proof(0);
-        let root = path.calculate_root(&leaves[0], &hasher.clone()).unwrap();
-
-        let mut witnessed_path = [(Value::unknown(), Value::unknown()); HEIGHT];
-        for (i, (left, right)) in path.path.into_iter().enumerate() {
-            witnessed_path[i] = (Value::known(left), Value::known(right));
-        }
-        let path = witnessed_path;
-
-        let circuit = TestCircuit { root: Value::known(root), path, leaf: Value::known(leaves[0]) };
-
-        let prover = MockProver::run(13, &circuit, vec![]).unwrap();
+        let hasher = PoseidonFp::new();
+        let empty_leaf = Fp::from(0);
+
+        let store = MemoryStorageFp::new();
+        let mut smt = SmtMemoryFp::new(store, hasher.clone(), empty_leaf.clone());
+
+        let leaves = vec![Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
+        // Use the leaf value as its position in the SMT
+        // Therefore we need an additional constraint that leaf == pos
+        let leaves: Vec<_> = leaves.into_iter().map(|l| (l, l)).collect();
+        smt.insert_batch(leaves.clone());
+
+        let (pos, leaf) = leaves[2];
+        assert_eq!(pos, leaf);
+        assert_eq!(smt.get_leaf(&pos), leaf);
+
+        let root = smt.root();
+        let path = smt.prove_membership(&pos);
+        assert!(path.verify(&root, &leaf, &pos));
+
+        let circuit = TestCircuit {
+            root: Value::known(root),
+            path: Value::known(path.path),
+            leaf: Value::known(leaf),
+        };
+
+        const K: u32 = 14;
+        let prover = MockProver::run(K, &circuit, vec![]).unwrap();
         prover.assert_satisfied();
+
+        //use halo2_proofs::dev::CircuitLayout;
+        //use plotters::prelude::*;
+        //let root = BitMapBackend::new("target/smt.png", (3840, 2160)).into_drawing_area();
+        //CircuitLayout::default().render(K, &circuit, &root).unwrap();
     }
 }

src/zk/vm.rs

@@ -23,9 +23,9 @@ use darkfi_sdk::crypto::{
         sinsemilla::{OrchardCommitDomains, OrchardHashDomains},
         util::gen_const_array,
         ConstBaseFieldElement, OrchardFixedBases, OrchardFixedBasesFull, ValueCommitV,
-        MERKLE_DEPTH_ORCHARD, SPARSE_MERKLE_DEPTH,
+        MERKLE_DEPTH_ORCHARD,
     },
-    smt,
+    smt::SMT_FP_DEPTH,
 };
 use halo2_gadgets::{
     ecc::{
@@ -64,7 +64,7 @@ use super::{
         less_than::{LessThanChip, LessThanConfig},
         native_range_check::{NativeRangeCheckChip, NativeRangeCheckConfig},
         small_range_check::{SmallRangeCheckChip, SmallRangeCheckConfig},
-        smt as smt_gadget,
+        smt,
         zero_cond::{ZeroCondChip, ZeroCondConfig},
     },
     tracer::ZkTracer,
@@ -74,10 +74,6 @@ use crate::zkas::{
     Opcode, ZkBinary,
 };
 
-type SmtPathConfig = smt_gadget::PathConfig<SPARSE_MERKLE_DEPTH>;
-pub(super) type SmtPathChip =
-    smt_gadget::PathChip<smt::Poseidon<pallas::Base, 2>, SPARSE_MERKLE_DEPTH>;
-
 /// Available chips/gadgets in the zkvm
 #[derive(Debug, Clone)]
 #[allow(clippy::large_enum_variant)]
@@ -94,7 +90,7 @@ enum VmChip {
     ),
 
     /// Sparse merkle tree (using Poseidon)
-    SparseTree(SmtPathConfig),
+    SparseTree(smt::PathConfig),
 
     /// Sinsemilla chip
     Sinsemilla(
@@ -175,14 +171,14 @@ impl VmConfig {
         Some(MerkleChip::construct(merkle_cfg2.clone()))
     }
 
-    fn sparse_tree_cfg(&self) -> Option<SmtPathConfig> {
-        let Some(VmChip::SparseTree(smt_config)) =
+    fn smt_chip(&self) -> Option<smt::PathChip> {
+        let Some(VmChip::SparseTree(config)) =
             self.chips.iter().find(|&c| matches!(c, VmChip::SparseTree(_)))
         else {
             return None
         };
 
-        Some(smt_config.clone())
+        Some(smt::PathChip::construct(config.clone()))
     }
 
     fn poseidon_chip(&self) -> Option<PoseidonChip<pallas::Base, 3, 2>> {
@@ -493,10 +489,10 @@ impl Circuit<pallas::Base> for ZkCircuit {
             (sinsemilla_cfg2, merkle_cfg2)
         };
 
-        let smt_config = SmtPathChip::configure(
+        let smt_config = smt::PathChip::configure(
             meta,
-            advices[..SPARSE_MERKLE_DEPTH].try_into().unwrap(),
-            advices[1..5].try_into().unwrap(),
+            advices[0..2].try_into().unwrap(),
+            advices[2..6].try_into().unwrap(),
             poseidon_config.clone(),
         );
 
@@ -627,6 +623,9 @@ impl Circuit<pallas::Base> for ZkCircuit {
         // Construct the zero_cond selection chip
         let zerocond_chip = config.zerocond_chip();
 
+        // Construct sparse Merkle tree chip
+        let smt_chip = config.smt_chip().unwrap();
+
         // ==========================
         // Constants setup
         // ==========================
@@ -772,11 +771,11 @@ impl Circuit<pallas::Base> for ZkCircuit {
                 }
 
                 Witness::SparseMerklePath(w) => {
-                    let path_cfg = config.sparse_tree_cfg().unwrap();
-                    let path_chip = SmtPathChip::from_native(path_cfg, &mut layouter, *w)?;
+                    let path: Value<[pallas::Base; SMT_FP_DEPTH]> =
+                        w.map(|typed_path| gen_const_array(|i| typed_path[i]));
 
                     trace!(target: "zk::vm", "Pushing SparseMerklePath to heap address {}", heap.len());
-                    heap.push(HeapVar::SparseMerklePath(path_chip));
+                    heap.push(HeapVar::SparseMerklePath(path));
                 }
 
                 Witness::Uint32(w) => {
@@ -980,7 +979,8 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     let args = &opcode.1;
 
                     let leaf_pos = heap[args[0].1].clone().try_into()?;
-                    let merkle_path = heap[args[1].1].clone().try_into()?;
+                    let merkle_path: Value<[Fp; MERKLE_DEPTH_ORCHARD]> =
+                        heap[args[1].1].clone().try_into()?;
                     let leaf = heap[args[2].1].clone().try_into()?;
 
                     let merkle_inputs = MerklePath::construct(
@@ -1003,10 +1003,12 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     let args = &opcode.1;
 
                     let root = heap[args[0].1].clone().try_into()?;
-                    let path_chip: SmtPathChip = heap[args[1].1].clone().try_into()?;
+                    let path: Value<[Fp; SMT_FP_DEPTH]> = heap[args[1].1].clone().try_into()?;
                     let leaf = heap[args[2].1].clone().try_into()?;
+                    let pos = heap[args[3].1].clone().try_into()?;
 
-                    let is_member = path_chip.check_membership(&mut layouter, root, leaf)?;
+                    let is_member =
+                        smt_chip.check_membership(&mut layouter, root, leaf, pos, path)?;
 
                     self.tracer.push_base(&is_member);
                     heap.push(HeapVar::Base(is_member));

src/zk/vm_heap.rs

@@ -18,7 +18,8 @@
 
 //! VM heap type abstractions
 use darkfi_sdk::crypto::{
-    constants::{OrchardFixedBases, SPARSE_MERKLE_DEPTH},
+    constants::{OrchardFixedBases, MERKLE_DEPTH_ORCHARD},
+    smt::SMT_FP_DEPTH,
     MerkleNode,
 };
 use halo2_gadgets::ecc::{
@@ -32,15 +33,12 @@ use halo2_proofs::{
 };
 use log::error;
 
-use super::vm::SmtPathChip;
 use crate::{
     zkas::{decoder::ZkBinary, types::VarType},
     Error::ZkasDecoderError,
     Result,
 };
 
-type SmtPath = [(Value<pallas::Base>, Value<pallas::Base>); SPARSE_MERKLE_DEPTH];
-
 /// These represent the witness types outside of the circuit
 #[allow(clippy::large_enum_variant)]
 #[derive(Clone)]
@@ -50,8 +48,8 @@ pub enum Witness {
     EcFixedPoint(Value<pallas::Point>),
     Base(Value<pallas::Base>),
     Scalar(Value<pallas::Scalar>),
-    MerklePath(Value<[MerkleNode; 32]>),
-    SparseMerklePath(SmtPath),
+    MerklePath(Value<[MerkleNode; MERKLE_DEPTH_ORCHARD]>),
+    SparseMerklePath(Value<[pallas::Base; SMT_FP_DEPTH]>),
     Uint32(Value<u32>),
     Uint64(Value<u64>),
 }
@@ -85,9 +83,7 @@ pub fn empty_witnesses(zkbin: &ZkBinary) -> Result<Vec<Witness>> {
             VarType::Base => ret.push(Witness::Base(Value::unknown())),
             VarType::Scalar => ret.push(Witness::Scalar(Value::unknown())),
             VarType::MerklePath => ret.push(Witness::MerklePath(Value::unknown())),
-            VarType::SparseMerklePath => ret.push(Witness::SparseMerklePath(
-                [(Value::unknown(), Value::unknown()); SPARSE_MERKLE_DEPTH],
-            )),
+            VarType::SparseMerklePath => ret.push(Witness::SparseMerklePath(Value::unknown())),
             VarType::Uint32 => ret.push(Witness::Uint32(Value::unknown())),
             VarType::Uint64 => ret.push(Witness::Uint64(Value::unknown())),
             x => return Err(ZkasDecoderError(format!("Unsupported witness type: {:?}", x))),
@@ -108,8 +104,8 @@ pub enum HeapVar {
     EcFixedPointBase(FixedPointBaseField<pallas::Affine, EccChip<OrchardFixedBases>>),
     Base(AssignedCell<pallas::Base, pallas::Base>),
     Scalar(ScalarFixed<pallas::Affine, EccChip<OrchardFixedBases>>),
-    MerklePath(Value<[pallas::Base; 32]>),
-    SparseMerklePath(SmtPathChip),
+    MerklePath(Value<[pallas::Base; MERKLE_DEPTH_ORCHARD]>),
+    SparseMerklePath(Value<[pallas::Base; SMT_FP_DEPTH]>),
     Uint32(Value<u32>),
     Uint64(Value<u64>),
 }
@@ -140,5 +136,5 @@ impl_try_from!(EcFixedPointBase, FixedPointBaseField<pallas::Affine, EccChip<Orc
 impl_try_from!(Scalar, ScalarFixed<pallas::Affine, EccChip<OrchardFixedBases>>);
 impl_try_from!(Base, AssignedCell<pallas::Base, pallas::Base>);
 impl_try_from!(Uint32, Value<u32>);
-impl_try_from!(MerklePath, Value<[pallas::Base; 32]>);
-impl_try_from!(SparseMerklePath, SmtPathChip);
+impl_try_from!(MerklePath, Value<[pallas::Base; MERKLE_DEPTH_ORCHARD]>);
+impl_try_from!(SparseMerklePath, Value<[pallas::Base; SMT_FP_DEPTH]>);

src/zkas/opcode.rs

@@ -223,9 +223,10 @@ impl Opcode {
                 (vec![VarType::Base], vec![VarType::Uint32, VarType::MerklePath, VarType::Base])
             }
 
-            Opcode::SparseTreeIsMember => {
-                (vec![VarType::Base], vec![VarType::Base, VarType::SparseMerklePath, VarType::Base])
-            }
+            Opcode::SparseTreeIsMember => (
+                vec![VarType::Base],
+                vec![VarType::Base, VarType::SparseMerklePath, VarType::Base, VarType::Base],
+            ),
 
             Opcode::BaseAdd => (vec![VarType::Base], vec![VarType::Base, VarType::Base]),
 

tests/smt.rs

@@ -16,10 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use darkfi_sdk::crypto::{
-    constants::SPARSE_MERKLE_DEPTH,
-    smt::{Poseidon, SparseMerkleTree},
-};
+use darkfi_sdk::crypto::smt::{MemoryStorageFp, PoseidonFp, SmtMemoryFp};
 use halo2_proofs::{arithmetic::Field, circuit::Value, dev::MockProver, pasta::Fp};
 use rand::rngs::OsRng;
 
@@ -39,31 +36,31 @@ fn zkvm_smt() -> Result<()> {
     let bincode = include_bytes!("../proof/smt.zk.bin");
     let zkbin = ZkBinary::decode(bincode)?;
 
-    let poseidon = Poseidon::<Fp, 2>::new();
-    let empty_leaf = [0u8; 32];
-    let leaves = [Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
+    let hasher = PoseidonFp::new();
+    let empty_leaf = Fp::from(0);
+
+    let store = MemoryStorageFp::new();
+    let mut smt = SmtMemoryFp::new(store, hasher.clone(), empty_leaf.clone());
 
-    let smt = SparseMerkleTree::<Fp, Poseidon<Fp, 2>, SPARSE_MERKLE_DEPTH>::new_sequential(
-        &leaves,
-        &poseidon.clone(),
-        empty_leaf,
-    )
-    .unwrap();
+    let leaves = vec![Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
+    // Use the leaf value as its position in the SMT
+    // Therefore we need an additional constraint that leaf == pos
+    let leaves: Vec<_> = leaves.into_iter().map(|l| (l, l)).collect();
+    smt.insert_batch(leaves.clone());
 
-    let path = smt.generate_membership_proof(0);
-    let root = path.calculate_root(&leaves[0], &poseidon).unwrap();
+    let (pos, leaf) = leaves[2];
+    assert_eq!(pos, leaf);
+    assert_eq!(smt.get_leaf(&pos), leaf);
 
-    let mut witnessed_path = [(Value::unknown(), Value::unknown()); SPARSE_MERKLE_DEPTH];
-    for (i, (left, right)) in path.path.into_iter().enumerate() {
-        witnessed_path[i] = (Value::known(left), Value::known(right));
-    }
-    let path = witnessed_path;
+    let root = smt.root();
+    let path = smt.prove_membership(&pos);
+    assert!(path.verify(&root, &leaf, &pos));
 
     // Values for the proof
     let prover_witnesses = vec![
         Witness::Base(Value::known(root)),
-        Witness::SparseMerklePath(path),
-        Witness::Base(Value::known(leaves[0])),
+        Witness::SparseMerklePath(Value::known(path.path)),
+        Witness::Base(Value::known(leaf)),
     ];
 
     let public_inputs = vec![root];