diff --git a/src/zk/gadget/cond_select.rs b/src/zk/gadget/cond_select.rs index 7ff7f2fea..ba4042820 100644 --- a/src/zk/gadget/cond_select.rs +++ b/src/zk/gadget/cond_select.rs @@ -39,7 +39,6 @@ pub struct ConditionalSelectConfig + Ord> { pub struct ConditionalSelectChip + Ord> { config: ConditionalSelectConfig, - _marker: PhantomData, } impl + Ord> Chip for ConditionalSelectChip { @@ -56,11 +55,8 @@ impl + Ord> Chip for ConditionalSelectChip { } impl + Ord> ConditionalSelectChip { - pub fn construct( - config: >::Config, - _loaded: >::Loaded, - ) -> Self { - Self { config, _marker: PhantomData } + pub fn construct(config: >::Config) -> Self { + Self { config } } pub fn configure( diff --git a/src/zk/gadget/is_equal.rs b/src/zk/gadget/is_equal.rs index ca5d67221..b8f1b6dce 100644 --- a/src/zk/gadget/is_equal.rs +++ b/src/zk/gadget/is_equal.rs @@ -53,10 +53,7 @@ impl + Ord> Chip for IsEqualChip { } impl + Ord> IsEqualChip { - pub fn construct( - config: >::Config, - _loaded: >::Loaded, - ) -> Self { + pub fn construct(config: >::Config) -> Self { Self { config } } @@ -158,10 +155,7 @@ impl + Ord> Chip for AssertEqualChip { } impl + Ord> AssertEqualChip { - pub fn construct( - config: >::Config, - _loaded: >::Loaded, - ) -> Self { + pub fn construct(config: >::Config) -> Self { Self { config } } diff --git a/src/zk/gadget/mod.rs b/src/zk/gadget/mod.rs index a3dfeaa97..a7d557bd4 100644 --- a/src/zk/gadget/mod.rs +++ b/src/zk/gadget/mod.rs @@ -39,3 +39,6 @@ pub mod cond_select; /// Conditional selection based on lhs (will output lhs if lhs==0, otherwise rhs) pub mod zero_cond; + +/// Poseidon-based sparse Merkle tree chip +pub mod smt; diff --git a/src/zk/gadget/smt.rs b/src/zk/gadget/smt.rs new file mode 100644 index 000000000..a5ee715aa --- /dev/null +++ b/src/zk/gadget/smt.rs @@ -0,0 +1,328 @@ +/* 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 . + */ + +use std::marker::PhantomData; + +use darkfi_sdk::crypto::smt::{FieldHasher, Path}; +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}, +}; + +use super::{ + cond_select::{ConditionalSelectChip, ConditionalSelectConfig, NUM_OF_UTILITY_ADVICE_COLUMNS}, + is_equal::{AssertEqualChip, AssertEqualConfig, IsEqualChip, IsEqualConfig}, +}; + +#[derive(Clone, Debug)] +pub struct PathConfig { + s_path: Selector, + advices: [Column; N], + poseidon_config: PoseidonConfig, + is_eq_config: IsEqualConfig, + conditional_select_config: ConditionalSelectConfig, + assert_equal_config: AssertEqualConfig, +} + +impl PathConfig { + fn poseidon_chip(&self) -> PoseidonChip { + PoseidonChip::construct(self.poseidon_config.clone()) + } + + fn is_eq_chip(&self) -> IsEqualChip { + IsEqualChip::construct(self.is_eq_config.clone()) + } + + fn conditional_select_chip(&self) -> ConditionalSelectChip { + ConditionalSelectChip::construct(self.conditional_select_config.clone()) + } + + fn assert_eq_chip(&self) -> AssertEqualChip { + AssertEqualChip::construct(self.assert_equal_config.clone()) + } +} + +pub struct PathChip, const N: usize> { + path: [(AssignedCell, AssignedCell); N], + config: PathConfig, + _hasher: PhantomData, +} + +impl, const N: usize> PathChip { + pub fn configure( + meta: &mut ConstraintSystem, + advices: [Column; N], + utility_advices: [Column; NUM_OF_UTILITY_ADVICE_COLUMNS], + poseidon_config: PoseidonConfig, + ) -> PathConfig { + let s_path = meta.selector(); + + for advice in &advices { + meta.enable_equality(*advice); + } + + for advice in &utility_advices { + meta.enable_equality(*advice); + } + + PathConfig { + 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, + [utility_advices[0], utility_advices[1]], + ), + } + } + + pub fn from_native( + config: PathConfig, + layouter: &mut impl Layouter, + native: Path, + ) -> Result { + 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, + || Value::known(native.path[i].0), + ) + }) + .collect::>, plonk::Error>>(); + + let right = (0..N) + .map(|i| { + region.assign_advice( + || format!("path[{}][{}]", i, 1), + config.advices[i], + 1, + || Value::known(native.path[i].1), + ) + }) + .collect::>, plonk::Error>>(); + + let result = left? + .into_iter() + .zip(right?.into_iter()) + .collect::, AssignedCell)>>(); + + Ok(result.try_into().unwrap()) + }, + )?; + + Ok(PathChip { path, config, _hasher: PhantomData }) + } + + pub fn calculate_root( + &self, + layouter: &mut impl Layouter, + leaf: AssignedCell, + ) -> Result, plonk::Error> { + // Check levels between leaf level and root + let mut previous_hash = leaf; + + 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())?; + + let left_or_right = condselect_chip.conditional_select( + layouter, + left_hash.clone(), + right_hash.clone(), + previous_is_left, + )?; + + asserteq_chip.assert_equal(layouter, previous_hash, left_or_right)?; + + // Update previous_hash + let hasher = PoseidonHash::< + _, + _, + poseidon::P128Pow5T3, + poseidon::ConstantLength<2>, + 3, + 2, + >::init( + self.config.poseidon_chip(), + layouter.namespace(|| "SmtPoseidonHash init"), + )?; + + previous_hash = hasher.hash( + layouter.namespace(|| "SmtPoseidonHash hash"), + [left_hash.clone(), right_hash.clone()], + )?; + } + + Ok(previous_hash) + } + + pub fn check_membership( + &self, + layouter: &mut impl Layouter, + root_hash: AssignedCell, + leaf: AssignedCell, + ) -> Result, plonk::Error> { + let computed_root = self.calculate_root(layouter, leaf)?; + + self.config.is_eq_chip().is_eq_with_output(layouter, computed_root, root_hash) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use darkfi_sdk::crypto::smt::{Poseidon, SparseMerkleTree}; + use halo2_proofs::{arithmetic::Field, circuit::floor_planner, plonk::Circuit}; + + const HEIGHT: usize = 3; + + struct TestCircuit { + root: Fp, + path: Path, HEIGHT>, + leaf: Fp, + } + + impl Circuit for TestCircuit { + type Config = PathConfig; + type FloorPlanner = floor_planner::V1; + type Params = (); + + fn without_witnesses(&self) -> Self { + todo!() + } + + fn configure(meta: &mut ConstraintSystem) -> Self::Config { + let advices = [(); HEIGHT].map(|_| meta.advice_column()); + let utility_advices = [(); NUM_OF_UTILITY_ADVICE_COLUMNS].map(|_| meta.advice_column()); + let poseidon_advices = [(); 5].map(|_| meta.advice_column()); + + for advice in &advices { + meta.enable_equality(*advice); + } + + for advice in &utility_advices { + meta.enable_equality(*advice); + } + + for advice in &poseidon_advices { + meta.enable_equality(*advice); + } + + let rc_a = [(); 3].map(|_| meta.fixed_column()); + let rc_b = [(); 3].map(|_| meta.fixed_column()); + + let poseidon_config = PoseidonChip::configure::( + meta, + poseidon_advices[1..5].try_into().unwrap(), + poseidon_advices[0], + rc_a, + rc_b, + ); + + PathChip::, HEIGHT>::configure( + meta, + advices, + utility_advices, + poseidon_config, + ) + } + + fn synthesize( + &self, + config: Self::Config, + mut layouter: impl Layouter, + ) -> Result<(), plonk::Error> { + let (root_cell, leaf_cell, one) = layouter.assign_region( + || "test circuit", + |mut region| { + let root_cell = region.assign_advice( + || "root", + config.advices[0], + 0, + || Value::known(self.root), + )?; + + let leaf_cell = region.assign_advice( + || "leaf", + config.advices[1], + 0, + || Value::known(self.leaf), + )?; + + let one = region.assign_advice( + || "one", + config.advices[2], + 0, + || Value::known(Fp::ONE), + )?; + Ok((root_cell, leaf_cell, one)) + }, + )?; + + let path_chip = + PathChip::from_native(config.clone(), &mut layouter, self.path.clone())?; + + let res = path_chip.check_membership(&mut layouter, root_cell, leaf_cell)?; + + let assert_eq_chip = config.assert_eq_chip(); + assert_eq_chip.assert_equal(&mut layouter, res, one)?; + + Ok(()) + } + } + + #[test] + fn test_smt_circuit() { + let hasher = Poseidon::::hasher(); + let leaves: [Fp; HEIGHT] = [Fp::ZERO, Fp::ZERO, Fp::ZERO]; + let empty_leaf = [0u8; 64]; + + let smt = SparseMerkleTree::, HEIGHT>::new_sequential( + &leaves, + &hasher.clone(), + &empty_leaf, + ) + .unwrap(); + + let path = smt.generate_membership_proof(0); + let root = path.calculate_root(&leaves[0], &hasher.clone()).unwrap(); + + let _circuit = TestCircuit { root, path, leaf: leaves[0] }; + } +} diff --git a/src/zk/vm.rs b/src/zk/vm.rs index fff6cddd7..1cd33334b 100644 --- a/src/zk/vm.rs +++ b/src/zk/vm.rs @@ -191,7 +191,7 @@ impl VmConfig { return None }; - Some(ConditionalSelectChip::construct(condselect_config.clone(), ())) + Some(ConditionalSelectChip::construct(condselect_config.clone())) } fn zerocond_chip(&self) -> Option> {