zkas/zk: add sparse_tree_is_member() opcode

proof/smt.zk

@@ -0,0 +1,21 @@
+k = 13;
+field = "pallas";
+
+constant "SMT" {
+}
+
+witness "SMT" {
+    Base root,
+    SparseMerklePath path,
+    Base leaf,
+}
+
+circuit "SMT" {
+    is_member = sparse_tree_is_member(root, path, leaf);
+
+    ONE = witness_base(1);
+    constrain_equal_base(is_member, ONE);
+
+    constrain_instance(root);
+}
+

src/sdk/src/crypto/constants.rs

@@ -32,8 +32,11 @@ pub const BLOCK_HASH_DOMAIN: &str = "DarkFi:Block";
 
 pub const MERKLE_DEPTH_ORCHARD: usize = 32;
 
+// TODO: move to merkle_node.rs
 pub const MERKLE_DEPTH: u8 = MERKLE_DEPTH_ORCHARD as u8;
 
+pub const SPARSE_MERKLE_DEPTH: usize = 3;
+
 #[allow(dead_code)]
 /// $\ell^\mathsf{Orchard}_\mathsf{base}$
 pub(crate) const L_ORCHARD_BASE: usize = 255;

src/validator/fees.rs

@@ -16,7 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use darkfi_sdk::crypto::constants::MERKLE_DEPTH_ORCHARD;
+use darkfi_sdk::crypto::constants::{MERKLE_DEPTH_ORCHARD, SPARSE_MERKLE_DEPTH};
 
 use crate::zkas::{Opcode, VarType, ZkBinary};
 
@@ -48,6 +48,7 @@ pub fn circuit_gas_use(zkbin: &ZkBinary) -> u64 {
             VarType::Scalar => 20,
             VarType::ScalarArray => unreachable!(),
             VarType::MerklePath => 10 * MERKLE_DEPTH_ORCHARD as u64,
+            VarType::SparseMerklePath => 10 * SPARSE_MERKLE_DEPTH as u64,
             VarType::Uint32 => 10,
             VarType::Uint64 => 10,
             VarType::Any => 10,
@@ -69,6 +70,7 @@ pub fn circuit_gas_use(zkbin: &ZkBinary) -> u64 {
             Opcode::EcGetY => 5,
             Opcode::PoseidonHash => 20 + 10 * opcode.1.len() as u64,
             Opcode::MerkleRoot => 10 * MERKLE_DEPTH_ORCHARD as u64,
+            Opcode::SparseTreeIsMember => 10 * SPARSE_MERKLE_DEPTH as u64,
             Opcode::BaseAdd => 15,
             Opcode::BaseMul => 15,
             Opcode::BaseSub => 15,

src/zk/debug.rs

@@ -127,6 +127,7 @@ pub fn zkas_type_checks(
             Witness::Base(_) => *binary_witness == zkas::VarType::Base,
             Witness::Scalar(_) => *binary_witness == zkas::VarType::Scalar,
             Witness::MerklePath(_) => *binary_witness == zkas::VarType::MerklePath,
+            Witness::SparseMerklePath(_) => *binary_witness == zkas::VarType::SparseMerklePath,
             Witness::Uint32(_) => *binary_witness == zkas::VarType::Uint32,
             Witness::Uint64(_) => *binary_witness == zkas::VarType::Uint64,
         };

src/zk/gadget/smt.rs

@@ -20,7 +20,7 @@
 
 use std::marker::PhantomData;
 
-use darkfi_sdk::crypto::smt::{FieldHasher, Path};
+use darkfi_sdk::crypto::smt::FieldHasher;
 use halo2_gadgets::poseidon::{
     primitives as poseidon, Hash as PoseidonHash, Pow5Chip as PoseidonChip,
     Pow5Config as PoseidonConfig,
@@ -64,6 +64,7 @@ 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>,
@@ -103,7 +104,7 @@ impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
     pub fn from_native(
         config: PathConfig<N>,
         layouter: &mut impl Layouter<Fp>,
-        native: Path<Fp, H, N>,
+        path: [(Value<Fp>, Value<Fp>); N],
     ) -> Result<Self, plonk::Error> {
         let path = layouter.assign_region(
             || "path",
@@ -115,7 +116,7 @@ impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
                             || format!("path[{}][{}]", i, 0),
                             config.advices[i],
                             0,
-                            || Value::known(native.path[i].0),
+                            || path[i].0,
                         )
                     })
                     .collect::<Result<Vec<AssignedCell<Fp, Fp>>, plonk::Error>>();
@@ -126,7 +127,7 @@ impl<H: FieldHasher<Fp, 2>, const N: usize> PathChip<H, N> {
                             || format!("path[{}][{}]", i, 1),
                             config.advices[i],
                             1,
-                            || Value::known(native.path[i].1),
+                            || path[i].1,
                         )
                     })
                     .collect::<Result<Vec<AssignedCell<Fp, Fp>>, plonk::Error>>();
@@ -217,7 +218,7 @@ mod tests {
 
     struct TestCircuit {
         root: Value<Fp>,
-        path: Path<Fp, Poseidon<Fp, 2>, HEIGHT>,
+        path: [(Value<Fp>, Value<Fp>); HEIGHT],
         leaf: Value<Fp>,
     }
 
@@ -275,7 +276,7 @@ mod tests {
             mut layouter: impl Layouter<Fp>,
         ) -> Result<(), plonk::Error> {
             // Initialize the Path chip
-            let path_chip =
+            let path_chip: PathChip<Poseidon<Fp, 2>, HEIGHT> =
                 PathChip::from_native(config.clone(), &mut layouter, self.path.clone())?;
 
             // Initialize the AssertEqual chip
@@ -335,6 +336,12 @@ mod tests {
         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();

src/zk/vm.rs

@@ -18,11 +18,14 @@
 
 use std::collections::HashSet;
 
-use darkfi_sdk::crypto::constants::{
-    sinsemilla::{OrchardCommitDomains, OrchardHashDomains},
-    util::gen_const_array,
-    ConstBaseFieldElement, OrchardFixedBases, OrchardFixedBasesFull, ValueCommitV,
-    MERKLE_DEPTH_ORCHARD,
+use darkfi_sdk::crypto::{
+    constants::{
+        sinsemilla::{OrchardCommitDomains, OrchardHashDomains},
+        util::gen_const_array,
+        ConstBaseFieldElement, OrchardFixedBases, OrchardFixedBasesFull, ValueCommitV,
+        MERKLE_DEPTH_ORCHARD, SPARSE_MERKLE_DEPTH,
+    },
+    smt,
 };
 use halo2_gadgets::{
     ecc::{
@@ -61,6 +64,7 @@ use super::{
         less_than::{LessThanChip, LessThanConfig},
         native_range_check::{NativeRangeCheckChip, NativeRangeCheckConfig},
         small_range_check::{SmallRangeCheckChip, SmallRangeCheckConfig},
+        smt as smt_gadget,
         zero_cond::{ZeroCondChip, ZeroCondConfig},
     },
     tracer::ZkTracer,
@@ -70,6 +74,10 @@ 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)]
@@ -85,6 +93,9 @@ enum VmChip {
         ),
     ),
 
+    /// Sparse merkle tree (using Poseidon)
+    SparseTree(SmtPathConfig),
+
     /// Sinsemilla chip
     Sinsemilla(
         (
@@ -164,6 +175,16 @@ impl VmConfig {
         Some(MerkleChip::construct(merkle_cfg2.clone()))
     }
 
+    fn sparse_tree_cfg(&self) -> Option<SmtPathConfig> {
+        let Some(VmChip::SparseTree(smt_config)) =
+            self.chips.iter().find(|&c| matches!(c, VmChip::SparseTree(_)))
+        else {
+            return None
+        };
+
+        Some(smt_config.clone())
+    }
+
     fn poseidon_chip(&self) -> Option<PoseidonChip<pallas::Base, 3, 2>> {
         let Some(VmChip::Poseidon(poseidon_config)) =
             self.chips.iter().find(|&c| matches!(c, VmChip::Poseidon(_)))
@@ -472,6 +493,13 @@ impl Circuit<pallas::Base> for ZkCircuit {
             (sinsemilla_cfg2, merkle_cfg2)
         };
 
+        let smt_config = SmtPathChip::configure(
+            meta,
+            advices[..SPARSE_MERKLE_DEPTH].try_into().unwrap(),
+            advices[1..5].try_into().unwrap(),
+            poseidon_config.clone(),
+        );
+
         // K-table for 64 bit range check lookups
         let k_values_table_64 = meta.lookup_table_column();
         let native_64_range_check_config =
@@ -511,6 +539,7 @@ impl Circuit<pallas::Base> for ZkCircuit {
         let chips = vec![
             VmChip::Ecc(ecc_config),
             VmChip::Merkle((merkle_cfg1, merkle_cfg2)),
+            VmChip::SparseTree(smt_config),
             VmChip::Sinsemilla((sinsemilla_cfg1, sinsemilla_cfg2)),
             VmChip::Poseidon(poseidon_config),
             VmChip::Arithmetic(arith_config),
@@ -742,6 +771,14 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     heap.push(HeapVar::MerklePath(path));
                 }
 
+                Witness::SparseMerklePath(w) => {
+                    let path_cfg = config.sparse_tree_cfg().unwrap();
+                    let path_chip = SmtPathChip::from_native(path_cfg, &mut layouter, *w)?;
+
+                    trace!(target: "zk::vm", "Pushing SparseMerklePath to heap address {}", heap.len());
+                    heap.push(HeapVar::SparseMerklePath(path_chip));
+                }
+
                 Witness::Uint32(w) => {
                     trace!(target: "zk::vm", "Pushing Uint32 to heap address {}", heap.len());
                     heap.push(HeapVar::Uint32(*w));
@@ -938,6 +975,7 @@ impl Circuit<pallas::Base> for ZkCircuit {
                 }
 
                 Opcode::MerkleRoot => {
+                    // TODO: all these trace statements could have trace!(..., args) instead
                     trace!(target: "zk::vm", "Executing `MerkleRoot{:?}` opcode", opcode.1);
                     let args = &opcode.1;
 
@@ -960,6 +998,20 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     heap.push(HeapVar::Base(root));
                 }
 
+                Opcode::SparseTreeIsMember => {
+                    trace!(target: "zk::vm", "Executing `SparseTreeIsMember{:?}` opcode", opcode.1);
+                    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 leaf = heap[args[2].1].clone().try_into()?;
+
+                    let is_member = path_chip.check_membership(&mut layouter, root, leaf)?;
+
+                    self.tracer.push_base(&is_member);
+                    heap.push(HeapVar::Base(is_member));
+                }
+
                 Opcode::BaseAdd => {
                     trace!(target: "zk::vm", "Executing `BaseAdd{:?}` opcode", opcode.1);
                     let args = &opcode.1;

src/zk/vm_heap.rs

@@ -17,7 +17,10 @@
  */
 
 //! VM heap type abstractions
-use darkfi_sdk::crypto::{constants::OrchardFixedBases, MerkleNode};
+use darkfi_sdk::crypto::{
+    constants::{OrchardFixedBases, SPARSE_MERKLE_DEPTH},
+    MerkleNode,
+};
 use halo2_gadgets::ecc::{
     chip::EccChip, FixedPoint, FixedPointBaseField, FixedPointShort, NonIdentityPoint, Point,
     ScalarFixed,
@@ -29,12 +32,15 @@ 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)]
@@ -45,6 +51,7 @@ pub enum Witness {
     Base(Value<pallas::Base>),
     Scalar(Value<pallas::Scalar>),
     MerklePath(Value<[MerkleNode; 32]>),
+    SparseMerklePath(SmtPath),
     Uint32(Value<u32>),
     Uint64(Value<u64>),
 }
@@ -58,6 +65,7 @@ impl Witness {
             Self::Base(_) => "Base",
             Self::Scalar(_) => "Scalar",
             Self::MerklePath(_) => "MerklePath",
+            Self::SparseMerklePath(_) => "SparseMerklePath",
             Self::Uint32(_) => "Uint32",
             Self::Uint64(_) => "Uint64",
         }
@@ -77,6 +85,9 @@ 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::Uint32 => ret.push(Witness::Uint32(Value::unknown())),
             VarType::Uint64 => ret.push(Witness::Uint64(Value::unknown())),
             x => return Err(ZkasDecoderError(format!("Unsupported witness type: {:?}", x))),
@@ -98,6 +109,7 @@ pub enum HeapVar {
     Base(AssignedCell<pallas::Base, pallas::Base>),
     Scalar(ScalarFixed<pallas::Affine, EccChip<OrchardFixedBases>>),
     MerklePath(Value<[pallas::Base; 32]>),
+    SparseMerklePath(SmtPathChip),
     Uint32(Value<u32>),
     Uint64(Value<u64>),
 }
@@ -129,3 +141,4 @@ 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);

src/zkas/opcode.rs

@@ -52,6 +52,9 @@ pub enum Opcode {
     /// Calculate Merkle root, given a position, Merkle path, and an element
     MerkleRoot = 0x20,
 
+    /// Check for leaf membership in sparse merkle tree
+    SparseTreeIsMember = 0x21,
+
     /// Base field element addition
     BaseAdd = 0x30,
 
@@ -109,6 +112,7 @@ impl Opcode {
             "ec_get_y" => Some(Self::EcGetY),
             "poseidon_hash" => Some(Self::PoseidonHash),
             "merkle_root" => Some(Self::MerkleRoot),
+            "sparse_tree_is_member" => Some(Self::SparseTreeIsMember),
             "base_add" => Some(Self::BaseAdd),
             "base_mul" => Some(Self::BaseMul),
             "base_sub" => Some(Self::BaseSub),
@@ -138,6 +142,7 @@ impl Opcode {
             0x09 => Some(Self::EcGetY),
             0x10 => Some(Self::PoseidonHash),
             0x20 => Some(Self::MerkleRoot),
+            0x21 => Some(Self::SparseTreeIsMember),
             0x30 => Some(Self::BaseAdd),
             0x31 => Some(Self::BaseMul),
             0x32 => Some(Self::BaseSub),
@@ -168,6 +173,7 @@ impl Opcode {
             Self::EcGetY => "ec_get_y",
             Self::PoseidonHash => "poseidon_hash",
             Self::MerkleRoot => "merkle_root",
+            Self::SparseTreeIsMember => "sparse_tree_is_member",
             Self::BaseAdd => "base_add",
             Self::BaseMul => "base_mul",
             Self::BaseSub => "base_sub",
@@ -217,6 +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::BaseAdd => (vec![VarType::Base], vec![VarType::Base, VarType::Base]),
 
             Opcode::BaseMul => (vec![VarType::Base], vec![VarType::Base, VarType::Base]),

src/zkas/parser.rs

@@ -683,6 +683,7 @@ impl Parser {
             }
 
             // Valid witness types
+            // TODO: change to TryFrom impl for VarType
             match v.1.token.as_str() {
                 "EcPoint" => {
                     ret.push(Witness {
@@ -729,6 +730,15 @@ impl Parser {
                     });
                 }
 
+                "SparseMerklePath" => {
+                    ret.push(Witness {
+                        name: k.to_string(),
+                        typ: VarType::SparseMerklePath,
+                        line: v.0.line,
+                        column: v.0.column,
+                    });
+                }
+
                 "Uint32" => {
                     ret.push(Witness {
                         name: k.to_string(),

src/zkas/types.rs

@@ -68,9 +68,12 @@ pub enum VarType {
     /// Scalar field element array
     ScalarArray = 0x13,
 
-    /// A Merkle tree path
+    /// Merkle tree path
     MerklePath = 0x20,
 
+    /// Sparse merkle tree path
+    SparseMerklePath = 0x21,
+
     /// Unsigned 32-bit integer
     Uint32 = 0x30,
 
@@ -94,6 +97,7 @@ impl VarType {
             0x12 => Some(Self::Scalar),
             0x13 => Some(Self::ScalarArray),
             0x20 => Some(Self::MerklePath),
+            0x21 => Some(Self::SparseMerklePath),
             0x30 => Some(Self::Uint32),
             0x31 => Some(Self::Uint64),
             0xff => Some(Self::Any),
@@ -114,6 +118,7 @@ impl VarType {
             Self::Scalar => "Scalar",
             Self::ScalarArray => "ScalarArray",
             Self::MerklePath => "MerklePath",
+            Self::SparseMerklePath => "SparseMerklePath",
             Self::Uint32 => "Uint32",
             Self::Uint64 => "Uint64",
             Self::Any => "Any",

tests/smt.rs

@@ -0,0 +1,85 @@
+/* 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 darkfi_sdk::crypto::{
+    constants::SPARSE_MERKLE_DEPTH,
+    smt::{Poseidon, SparseMerkleTree},
+};
+use halo2_proofs::{arithmetic::Field, circuit::Value, dev::MockProver, pasta::Fp};
+use rand::rngs::OsRng;
+
+use darkfi::{
+    zk::{
+        proof::{ProvingKey, VerifyingKey},
+        vm::ZkCircuit,
+        vm_heap::{empty_witnesses, Witness},
+        Proof,
+    },
+    zkas::ZkBinary,
+    Result,
+};
+
+#[test]
+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; 64];
+    let leaves = [Fp::random(&mut OsRng), Fp::random(&mut OsRng), Fp::random(&mut OsRng)];
+
+    let smt = SparseMerkleTree::<Fp, Poseidon<Fp, 2>, SPARSE_MERKLE_DEPTH>::new_sequential(
+        &leaves,
+        &poseidon.clone(),
+        &empty_leaf,
+    )
+    .unwrap();
+
+    let path = smt.generate_membership_proof(0);
+    let root = path.calculate_root(&leaves[0], &poseidon).unwrap();
+
+    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;
+
+    // Values for the proof
+    let prover_witnesses = vec![
+        Witness::Base(Value::known(root)),
+        Witness::SparseMerklePath(path),
+        Witness::Base(Value::known(leaves[0])),
+    ];
+
+    let public_inputs = vec![root];
+
+    let circuit = ZkCircuit::new(prover_witnesses, &zkbin);
+
+    let mockprover = MockProver::run(zkbin.k, &circuit, vec![public_inputs.clone()])?;
+    mockprover.assert_satisfied();
+
+    let proving_key = ProvingKey::build(zkbin.k, &circuit);
+    let proof = Proof::create(&proving_key, &[circuit], &public_inputs, &mut OsRng)?;
+
+    let verifier_witnesses = empty_witnesses(&zkbin)?;
+    let circuit = ZkCircuit::new(verifier_witnesses, &zkbin);
+    let verifying_key = VerifyingKey::build(zkbin.k, &circuit);
+    proof.verify(&verifying_key, &public_inputs)?;
+
+    Ok(())
+}