smt: simplify ZK gadget. Use root = sparse_merkle_root(pos, path, leaf) instead of the more complicated is_member = sparse_tree_is_member(root, path, pos, leaf)

2026-01-07 22:04:03 -05:00 · 2024-04-02 09:38:32 +02:00
parent 4d0c36a508
commit 9188a62bb3
11 changed files with 45 additions and 87 deletions
--- a/proof/smt.zk
+++ b/proof/smt.zk
@@ -5,17 +5,12 @@ constant "SMT" {
 }

 witness "SMT" {
-    Base root,
    SparseMerklePath path,
    Base leaf,
 }

 circuit "SMT" {
-    is_member = sparse_tree_is_member(root, path, leaf, leaf);
-
-    ONE = witness_base(1);
-    constrain_equal_base(is_member, ONE);
-
+    root = sparse_merkle_root(leaf, path, leaf);
    constrain_instance(root);
 }

--- a/src/contract/dao/proof/propose-input.zk
+++ b/src/contract/dao/proof/propose-input.zk
@@ -21,7 +21,6 @@ witness "ProposeInput" {
    Uint32 leaf_pos,
    MerklePath coin_path,

-    Base null_tree_root,
    SparseMerklePath null_path,

    Base signature_secret,
@@ -45,15 +44,12 @@ circuit "ProposeInput" {
    # We need this to detect whether the above coin was already spent.
    # Use a SMT, and show that at this position, the leaf is ZERO
    ZERO = witness_base(0);
-    ONE = witness_base(1);
    nullifier = poseidon_hash(coin_secret, coin);
-    is_member = sparse_tree_is_member(
-        null_tree_root,         # Expected root
+    null_tree_root = sparse_merkle_root(
+        nullifier,              # Position
        null_path,              # Path to root
        ZERO,                   # Leaf value
-        nullifier               # Position
    );
-    constrain_equal_base(is_member, ONE);
    constrain_instance(null_tree_root);

    # Pedersen commitment for coin's coin_value
--- a/src/contract/dao/proof/vote-input.zk
+++ b/src/contract/dao/proof/vote-input.zk
@@ -23,7 +23,6 @@ witness "VoteInput" {
    Uint32 leaf_pos,
    MerklePath coin_path,

-    Base null_tree_root,
    SparseMerklePath null_path,

    Base signature_secret,
@@ -46,15 +45,12 @@ circuit "VoteInput" {
    # We need this to detect whether the above coin was already spent.
    # Use a SMT, and show that at this position, the leaf is ZERO
    ZERO = witness_base(0);
-    ONE = witness_base(1);
    nullifier = poseidon_hash(coin_secret, coin);
-    is_member = sparse_tree_is_member(
-        null_tree_root,         # Expected root
+    null_tree_root = sparse_merkle_root(
+        nullifier,              # Position
        null_path,              # Path to root
        ZERO,                   # Leaf value
-        nullifier               # Position
    );
-    constrain_equal_base(is_member, ONE);
    constrain_instance(null_tree_root);

    # Include some secret information in vote nullifier to defeat correlation
--- a/src/contract/dao/src/client/propose.rs
+++ b/src/contract/dao/src/client/propose.rs
@@ -109,7 +109,6 @@ impl<'a> DaoProposeCall<'a> {
                Witness::Base(Value::known(gov_token_blind.inner())),
                Witness::Uint32(Value::known(leaf_pos.try_into().unwrap())),
                Witness::MerklePath(Value::known(input.merkle_path.clone().try_into().unwrap())),
-                Witness::Base(Value::known(smt_null_root)),
                Witness::SparseMerklePath(Value::known(smt_null_path.path)),
                Witness::Base(Value::known(input.signature_secret.inner())),
            ];
--- a/src/contract/dao/src/client/vote.rs
+++ b/src/contract/dao/src/client/vote.rs
@@ -139,7 +139,6 @@ impl<'a> DaoVoteCall<'a> {
                Witness::Base(Value::known(gov_token_blind)),
                Witness::Uint32(Value::known(leaf_pos.try_into().unwrap())),
                Witness::MerklePath(Value::known(input.merkle_path.clone().try_into().unwrap())),
-                Witness::Base(Value::known(smt_null_root)),
                Witness::SparseMerklePath(Value::known(smt_null_path.path)),
                Witness::Base(Value::known(input.signature_secret.inner())),
            ];
--- a/src/contract/dao/src/entrypoint/vote.rs
+++ b/src/contract/dao/src/entrypoint/vote.rs
@@ -16,9 +16,8 @@
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

-use darkfi_money_contract::MONEY_CONTRACT_NULLIFIERS_TREE;
 use darkfi_sdk::{
-    crypto::{contract_id::MONEY_CONTRACT_ID, pasta_prelude::*, ContractId, PublicKey},
+    crypto::{pasta_prelude::*, ContractId, PublicKey},
    dark_tree::DarkLeaf,
    db::{db_contains_key, db_get, db_lookup, db_set},
    error::{ContractError, ContractResult},
@@ -65,7 +64,7 @@ pub(crate) fn dao_vote_get_metadata(
        return Err(DaoError::ProposalNonexistent.into())
    };
    // Get the current votes
-    let mut proposal_metadata: DaoProposalMetadata = deserialize(&data)?;
+    let proposal_metadata: DaoProposalMetadata = deserialize(&data)?;

    // Iterate through inputs
    for input in &params.inputs {
@@ -144,7 +143,6 @@ pub(crate) fn dao_vote_process_instruction(
    let mut proposal_metadata: DaoProposalMetadata = deserialize(&data)?;

    // Check the Merkle root and nullifiers for the input coins are valid
-    let money_nullifier_db = db_lookup(*MONEY_CONTRACT_ID, MONEY_CONTRACT_NULLIFIERS_TREE)?;
    let dao_vote_nullifier_db = db_lookup(cid, DAO_CONTRACT_DB_VOTE_NULLIFIERS)?;
    let mut vote_nullifiers = vec![];

--- a/src/validator/fees.rs
+++ b/src/validator/fees.rs
@@ -70,7 +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::SparseMerkleRoot => 10 * SPARSE_MERKLE_DEPTH as u64,
            Opcode::BaseAdd => 15,
            Opcode::BaseMul => 15,
            Opcode::BaseSub => 15,
--- a/src/zk/gadget/smt.rs
+++ b/src/zk/gadget/smt.rs
@@ -33,7 +33,7 @@ use halo2_proofs::{

 use super::{
    cond_select::{ConditionalSelectChip, ConditionalSelectConfig, NUM_OF_UTILITY_ADVICE_COLUMNS},
-    is_equal::{AssertEqualChip, AssertEqualConfig, IsEqualChip, IsEqualConfig},
+    is_equal::{AssertEqualChip, AssertEqualConfig},
 };

 #[derive(Clone, Debug)]
@@ -41,7 +41,6 @@ pub struct PathConfig {
    s_path: Selector,
    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>,
 }
@@ -51,10 +50,6 @@ impl PathConfig {
        PoseidonChip::construct(self.poseidon_config.clone())
    }

-    fn is_eq_chip(&self) -> IsEqualChip<Fp> {
-        IsEqualChip::construct(self.is_eq_config.clone())
-    }
-
    fn conditional_select_chip(&self) -> ConditionalSelectChip<Fp> {
        ConditionalSelectChip::construct(self.conditional_select_config.clone())
    }
@@ -99,7 +94,6 @@ impl PathChip {
            s_path,
            advices,
            poseidon_config,
-            is_eq_config: IsEqualChip::configure(meta, utility_advices),
            conditional_select_config: ConditionalSelectChip::configure(meta, utility_advices),
            assert_equal_config: AssertEqualChip::configure(
                meta,
@@ -124,10 +118,9 @@ impl PathChip {
    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]>,
+        leaf: AssignedCell<Fp, Fp>,
    ) -> Result<AssignedCell<Fp, Fp>, plonk::Error> {
        let path = path.transpose_array();
        // Witness values
@@ -141,7 +134,7 @@ impl PathChip {
                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",
+                        || "witness pos bit",
                        self.config.advices[0],
                        i,
                        || bit,
@@ -149,7 +142,7 @@ impl PathChip {
                    witness_bits.push(bit);

                    let sibling = region.assign_advice(
-                        || "witness root",
+                        || "witness path sibling",
                        self.config.advices[1],
                        i,
                        || sibling,
@@ -158,7 +151,7 @@ impl PathChip {
                }

                let zero = region.assign_advice(
-                    || "witness one",
+                    || "witness zero",
                    self.config.advices[0],
                    SMT_FP_DEPTH,
                    || Value::known(Fp::ZERO),
@@ -171,7 +164,6 @@ impl PathChip {
        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();

@@ -235,7 +227,8 @@ impl PathChip {

        asserteq_chip.assert_equal(layouter, current_path, pos)?;

-        iseq_chip.is_eq_with_output(layouter, current_node, root)
+        let root = current_node;
+        Ok(root)
    }
 }

@@ -247,9 +240,9 @@ mod tests {
    use rand::rngs::OsRng;

    struct TestCircuit {
-        root: Value<Fp>,
        path: Value<[Fp; SMT_FP_DEPTH]>,
        leaf: Value<Fp>,
+        root: Value<Fp>,
    }

    impl Circuit<Fp> for TestCircuit {
@@ -302,43 +295,31 @@ mod tests {
            let assert_eq_chip = config.assert_eq_chip();

            // Witness values
-            let (root, leaf, one) = layouter.assign_region(
+            let (leaf, root) = layouter.assign_region(
                || "witness",
                |mut region| {
-                    let root = region.assign_advice(
-                        || "witness root",
-                        config.advices[0],
-                        0,
-                        || self.root,
-                    )?;
-
                    let leaf = region.assign_advice(
                        || "witness leaf",
-                        config.advices[1],
+                        config.advices[0],
                        0,
                        || self.leaf,
                    )?;
-
-                    let one = region.assign_advice(
-                        || "witness one",
+                    let root = region.assign_advice(
+                        || "witness root",
                        config.advices[1],
-                        1,
-                        || Value::known(Fp::ONE),
+                        0,
+                        || self.root,
                    )?;
-                    region.constrain_constant(one.cell(), Fp::ONE)?;
-
-                    Ok((root, leaf, one))
+                    Ok((leaf, root))
                },
            )?;

-            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)?;
+            let calc_root =
+                path_chip.check_membership(&mut layouter, leaf.clone(), self.path, leaf.clone())?;
+            // Normally we just reveal it as a public input.
+            // But I'm too lazy to make a separate config for this unit test so
+            // do this instead.
+            assert_eq_chip.assert_equal(&mut layouter, calc_root, root)?;

            Ok(())
        }
@@ -365,9 +346,9 @@ mod tests {
        assert!(path.verify(&root, &leaf, &pos));

        let circuit = TestCircuit {
-            root: Value::known(root),
            path: Value::known(path.path),
            leaf: Value::known(leaf),
+            root: Value::known(root),
        };

        const K: u32 = 14;
--- a/src/zk/vm.rs
+++ b/src/zk/vm.rs
@@ -999,20 +999,18 @@ impl Circuit<pallas::Base> for ZkCircuit {
                    heap.push(HeapVar::Base(root));
                }

-                Opcode::SparseTreeIsMember => {
+                Opcode::SparseMerkleRoot => {
                    trace!(target: "zk::vm", "Executing `SparseTreeIsMember{:?}` opcode", opcode.1);
                    let args = &opcode.1;

-                    let root = heap[args[0].1].clone().try_into()?;
+                    let pos = heap[args[0].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 =
-                        smt_chip.check_membership(&mut layouter, root, leaf, pos, path)?;
+                    let root = smt_chip.check_membership(&mut layouter, pos, path, leaf)?;

-                    self.tracer.push_base(&is_member);
-                    heap.push(HeapVar::Base(is_member));
+                    self.tracer.push_base(&root);
+                    heap.push(HeapVar::Base(root));
                }

                Opcode::BaseAdd => {
--- a/src/zkas/opcode.rs
+++ b/src/zkas/opcode.rs
@@ -52,8 +52,8 @@ 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,
+    /// Calculate sparse Merkle root, given the position, path and a member
+    SparseMerkleRoot = 0x21,

    /// Base field element addition
    BaseAdd = 0x30,
@@ -112,7 +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),
+            "sparse_merkle_root" => Some(Self::SparseMerkleRoot),
            "base_add" => Some(Self::BaseAdd),
            "base_mul" => Some(Self::BaseMul),
            "base_sub" => Some(Self::BaseSub),
@@ -142,7 +142,7 @@ impl Opcode {
            0x09 => Some(Self::EcGetY),
            0x10 => Some(Self::PoseidonHash),
            0x20 => Some(Self::MerkleRoot),
-            0x21 => Some(Self::SparseTreeIsMember),
+            0x21 => Some(Self::SparseMerkleRoot),
            0x30 => Some(Self::BaseAdd),
            0x31 => Some(Self::BaseMul),
            0x32 => Some(Self::BaseSub),
@@ -173,7 +173,7 @@ impl Opcode {
            Self::EcGetY => "ec_get_y",
            Self::PoseidonHash => "poseidon_hash",
            Self::MerkleRoot => "merkle_root",
-            Self::SparseTreeIsMember => "sparse_tree_is_member",
+            Self::SparseMerkleRoot => "sparse_merkle_root",
            Self::BaseAdd => "base_add",
            Self::BaseMul => "base_mul",
            Self::BaseSub => "base_sub",
@@ -223,10 +223,9 @@ impl Opcode {
                (vec![VarType::Base], vec![VarType::Uint32, VarType::MerklePath, VarType::Base])
            }

-            Opcode::SparseTreeIsMember => (
-                vec![VarType::Base],
-                vec![VarType::Base, VarType::SparseMerklePath, VarType::Base, VarType::Base],
-            ),
+            Opcode::SparseMerkleRoot => {
+                (vec![VarType::Base], vec![VarType::Base, VarType::SparseMerklePath, VarType::Base])
+            }

            Opcode::BaseAdd => (vec![VarType::Base], vec![VarType::Base, VarType::Base]),

--- a/tests/smt.rs
+++ b/tests/smt.rs
@@ -55,11 +55,8 @@ fn zkvm_smt() -> Result<()> {
    assert!(path.verify(&root, &leaf, &pos));

    // Values for the proof
-    let prover_witnesses = vec![
-        Witness::Base(Value::known(root)),
-        Witness::SparseMerklePath(Value::known(path.path)),
-        Witness::Base(Value::known(leaf)),
-    ];
+    let prover_witnesses =
+        vec![Witness::SparseMerklePath(Value::known(path.path)), Witness::Base(Value::known(leaf))];

    let public_inputs = vec![root];