switch zkVM to new SMT gadget

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/zk/gadget/smt2.rs

@@ -103,7 +103,7 @@ impl PathChip {
         bits
     }
 
-    fn check_membership(
+    pub fn check_membership(
         &self,
         layouter: &mut impl Layouter<Fp>,
         root: AssignedCell<Fp, Fp>,

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,
+    smt2::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,
+        smt2 as 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},
+    smt2::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::smt2::{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];