zkas/zk: add sparse_tree_is_member() opcode

This commit is contained in:
zero
2024-03-04 10:03:30 +01:00
parent 39223c6a98
commit 376784af2e
11 changed files with 223 additions and 14 deletions

21
proof/smt.zk Normal file
View File

@@ -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);
}

View File

@@ -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;

View File

@@ -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,

View File

@@ -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,
};

View File

@@ -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();

View File

@@ -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;

View File

@@ -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);

View File

@@ -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]),

View File

@@ -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(),

View File

@@ -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",

85
tests/smt.rs Normal file
View File

@@ -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(())
}