SMT2 ZK gadget

src/sdk/src/crypto/smt2/mod.rs

@@ -66,9 +66,19 @@ pub use util::Poseidon;
 
 use num_bigint::BigUint;
 use std::collections::HashMap;
+// Only used for the type aliases below
+use pasta_curves::pallas;
 
 use util::{FieldElement, FieldHasher};
 
+// Bit size for Fp (and Fq)
+pub const SMT_FP_DEPTH: usize = 255;
+pub type PoseidonFp = Poseidon<pallas::Base, 2>;
+pub type MemoryStorageFp = MemoryStorage<pallas::Base>;
+pub type SmtMemoryFp =
+    SparseMerkleTree<SMT_FP_DEPTH, { SMT_FP_DEPTH + 1 }, pallas::Base, PoseidonFp, MemoryStorageFp>;
+pub type PathFp = Path<SMT_FP_DEPTH, pallas::Base, PoseidonFp>;
+
 /// Pluggable storage backend for the SMT.
 /// Has a minimal interface to simply put and get objects from the store.
 pub trait StorageAdapter {
@@ -242,10 +252,11 @@ impl<const N: usize, F: FieldElement, H: FieldHasher<F, 2>> Path<N, F, H> {
             let sibling_node = self.path[i];
 
             let is_right = pos.bit((N - 1 - i) as u64);
-            //println!("is_right: {}", is_right);
             let (left, right) =
                 if is_right { (sibling_node, current_node) } else { (current_node, sibling_node) };
+            //println!("is_right: {}", is_right);
             //println!("left: {:?}, right: {:?}", left, right);
+            //println!("current_node: {:?}", current_node);
 
             current_node = self.hasher.hash([left, right]);
         }

src/zk/gadget/mod.rs

@@ -42,3 +42,4 @@ pub mod zero_cond;
 
 /// Poseidon-based sparse Merkle tree chip
 pub mod smt;
+pub mod smt2;

src/zk/gadget/smt2.rs

@@ -0,0 +1,366 @@
+use darkfi_sdk::crypto::smt2::SMT_FP_DEPTH;
+use halo2_gadgets::poseidon::{
+    primitives as poseidon, Hash as PoseidonHash, Pow5Chip as PoseidonChip,
+    Pow5Config as PoseidonConfig,
+};
+use halo2_proofs::{
+    circuit::{AssignedCell, Layouter, Value},
+    pasta::{
+        group::ff::{Field, PrimeFieldBits},
+        Fp,
+    },
+    plonk::{self, Advice, Column, ConstraintSystem, Constraints, Selector},
+    poly::Rotation,
+};
+
+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<Advice>; 2],
+    poseidon_config: PoseidonConfig<Fp, 3, 2>,
+    is_eq_config: IsEqualConfig<Fp>,
+    conditional_select_config: ConditionalSelectConfig<Fp>,
+    assert_equal_config: AssertEqualConfig<Fp>,
+}
+
+impl PathConfig {
+    fn poseidon_chip(&self) -> PoseidonChip<Fp, 3, 2> {
+        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())
+    }
+
+    fn assert_eq_chip(&self) -> AssertEqualChip<Fp> {
+        AssertEqualChip::construct(self.assert_equal_config.clone())
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct PathChip {
+    config: PathConfig,
+}
+
+impl PathChip {
+    pub fn configure(
+        meta: &mut ConstraintSystem<Fp>,
+        advices: [Column<Advice>; 2],
+        utility_advices: [Column<Advice>; NUM_OF_UTILITY_ADVICE_COLUMNS],
+        poseidon_config: PoseidonConfig<Fp, 3, 2>,
+    ) -> PathConfig {
+        let s_path = meta.selector();
+
+        for advice in &advices {
+            meta.enable_equality(*advice);
+        }
+
+        for advice in &utility_advices {
+            meta.enable_equality(*advice);
+        }
+
+        meta.create_gate("Path builder", |meta| {
+            let s_path = meta.query_selector(s_path);
+            let current_path = meta.query_advice(advices[0], Rotation::cur());
+            let bit = meta.query_advice(advices[1], Rotation::cur());
+            let next_path = meta.query_advice(advices[0], Rotation::next());
+
+            Constraints::with_selector(s_path, Some(next_path - (current_path * Fp::from(2) + bit)))
+        });
+
+        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 construct(config: PathConfig) -> Self {
+        Self { config }
+    }
+
+    fn decompose_value(value: &Fp) -> Vec<Fp> {
+        // Returns 256 bits, but the last bit is uneeded
+        let bits: Vec<bool> = value.to_le_bits().into_iter().collect();
+
+        let mut bits: Vec<Fp> = bits[..SMT_FP_DEPTH].iter().map(|x| Fp::from(*x)).collect();
+        bits.resize(SMT_FP_DEPTH, Fp::from(0));
+        bits
+    }
+
+    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]>,
+    ) -> Result<AssignedCell<Fp, Fp>, plonk::Error> {
+        let path = path.transpose_array();
+        // Witness values
+        let (bits, path, zero) = layouter.assign_region(
+            || "witness",
+            |mut region| {
+                let bits = pos.value().map(Self::decompose_value).transpose_vec(SMT_FP_DEPTH);
+                assert_eq!(bits.len(), SMT_FP_DEPTH);
+
+                let mut witness_bits = vec![];
+                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",
+                        self.config.advices[0],
+                        i,
+                        || bit,
+                    )?;
+                    witness_bits.push(bit);
+
+                    let sibling = region.assign_advice(
+                        || "witness root",
+                        self.config.advices[1],
+                        i,
+                        || sibling,
+                    )?;
+                    witness_path.push(sibling);
+                }
+
+                let zero = region.assign_advice(
+                    || "witness one",
+                    self.config.advices[0],
+                    SMT_FP_DEPTH,
+                    || Value::known(Fp::ZERO),
+                )?;
+                region.constrain_constant(zero.cell(), Fp::ZERO)?;
+
+                Ok((witness_bits, witness_path, zero))
+            },
+        )?;
+        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();
+
+        // Check path construction
+        let mut current_path = zero;
+        for bit in bits.iter().rev() {
+            current_path = layouter.assign_region(
+                || "pᵢ₊₁ = 2pᵢ + bᵢ",
+                |mut region| {
+                    self.config.s_path.enable(&mut region, 0)?;
+
+                    current_path.copy_advice(
+                        || "current path",
+                        &mut region,
+                        self.config.advices[0],
+                        0,
+                    )?;
+                    bit.copy_advice(|| "path bit", &mut region, self.config.advices[1], 0)?;
+
+                    let next_path =
+                        current_path.value().zip(bit.value()).map(|(p, b)| p * Fp::from(2) + b);
+                    region.assign_advice(|| "next path", self.config.advices[0], 1, || next_path)
+                },
+            )?;
+        }
+
+        // Check tree construction
+        let mut current_node = leaf.clone();
+        for (bit, sibling) in bits.into_iter().zip(path.into_iter().rev()) {
+            // Conditional select also constraints the bit ∈ {0, 1}
+            let left = condselect_chip.conditional_select(
+                layouter,
+                sibling.clone(),
+                current_node.clone(),
+                bit.clone(),
+            )?;
+            let right = condselect_chip.conditional_select(
+                layouter,
+                current_node.clone(),
+                sibling,
+                bit.clone(),
+            )?;
+            //println!("bit: {:?}", bit);
+            //println!("left: {:?}, right: {:?}", left, right);
+
+            let hasher = PoseidonHash::<
+                _,
+                _,
+                poseidon::P128Pow5T3,
+                poseidon::ConstantLength<2>,
+                3,
+                2,
+            >::init(
+                self.config.poseidon_chip(),
+                layouter.namespace(|| "SmtPoseidonHash init"),
+            )?;
+
+            current_node =
+                hasher.hash(layouter.namespace(|| "SmtPoseidonHash hash"), [left, right])?;
+        }
+
+        asserteq_chip.assert_equal(layouter, current_path, pos)?;
+
+        iseq_chip.is_eq_with_output(layouter, current_node, root)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use darkfi_sdk::crypto::smt2::{MemoryStorageFp, PoseidonFp, SmtMemoryFp};
+    use halo2_proofs::{circuit::floor_planner, dev::MockProver, plonk::Circuit};
+    use rand::rngs::OsRng;
+
+    struct TestCircuit {
+        root: Value<Fp>,
+        path: Value<[Fp; SMT_FP_DEPTH]>,
+        leaf: Value<Fp>,
+    }
+
+    impl Circuit<Fp> for TestCircuit {
+        type Config = PathConfig;
+        type FloorPlanner = floor_planner::V1;
+        type Params = ();
+
+        fn without_witnesses(&self) -> Self {
+            Self { root: Value::unknown(), path: Value::unknown(), leaf: Value::unknown() }
+        }
+
+        fn configure(meta: &mut ConstraintSystem<Fp>) -> Self::Config {
+            // Advice wires required by PathChip
+            let advices = [(); 2].map(|_| meta.advice_column());
+            let utility_advices = [(); NUM_OF_UTILITY_ADVICE_COLUMNS].map(|_| meta.advice_column());
+
+            // Setup poseidon config
+            let poseidon_advices = [(); 4].map(|_| meta.advice_column());
+            for advice in &poseidon_advices {
+                meta.enable_equality(*advice);
+            }
+
+            // Needed for poseidon hash
+            let col_const = meta.fixed_column();
+            meta.enable_constant(col_const);
+
+            let rc_a = [(); 3].map(|_| meta.fixed_column());
+            let rc_b = [(); 3].map(|_| meta.fixed_column());
+
+            let poseidon_config = PoseidonChip::configure::<poseidon::P128Pow5T3>(
+                meta,
+                poseidon_advices[1..4].try_into().unwrap(),
+                poseidon_advices[0],
+                rc_a,
+                rc_b,
+            );
+
+            PathChip::configure(meta, advices, utility_advices, poseidon_config)
+        }
+
+        fn synthesize(
+            &self,
+            config: Self::Config,
+            mut layouter: impl Layouter<Fp>,
+        ) -> Result<(), plonk::Error> {
+            // Initialize the Path chip
+            let path_chip = PathChip::construct(config.clone());
+
+            // Initialize the AssertEqual chip
+            let assert_eq_chip = config.assert_eq_chip();
+
+            // Witness values
+            let (root, leaf, one) = 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],
+                        0,
+                        || self.leaf,
+                    )?;
+
+                    let one = region.assign_advice(
+                        || "witness one",
+                        config.advices[1],
+                        1,
+                        || Value::known(Fp::ONE),
+                    )?;
+                    region.constrain_constant(one.cell(), Fp::ONE)?;
+
+                    Ok((root, leaf, one))
+                },
+            )?;
+
+            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)?;
+
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn test_smt_circuit() {
+        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 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 (pos, leaf) = leaves[2];
+        assert_eq!(pos, leaf);
+        assert_eq!(smt.get_leaf(&pos), leaf);
+
+        let root = smt.root();
+        let path = smt.prove_membership(&pos);
+        assert!(path.verify(&root, &leaf, &pos));
+
+        let circuit = TestCircuit {
+            root: Value::known(root),
+            path: Value::known(path.path),
+            leaf: Value::known(leaf),
+        };
+
+        const K: u32 = 14;
+        let prover = MockProver::run(K, &circuit, vec![]).unwrap();
+        prover.assert_satisfied();
+
+        //use halo2_proofs::dev::CircuitLayout;
+        //use plotters::prelude::*;
+        //let root = BitMapBackend::new("target/smt.png", (3840, 2160)).into_drawing_area();
+        //CircuitLayout::default().render(K, &circuit, &root).unwrap();
+    }
+}