zk/gadget: Implement `f(a,b) = if a==0 {a} else {b}` as zero_cond chip.

src/zk/gadget/mod.rs

@@ -34,5 +34,5 @@ pub mod is_zero;
 /// Conditional selection
 pub mod cond_select;
 
-// Sparse Merkle Tree (using poseidon hash) gadget
-//pub mod smt_poseidon;
+/// Conditional selection based on lhs (will output lhs if lhs==0, otherwise rhs)
+pub mod zero_cond;

src/zk/gadget/zero_cond.rs

@@ -0,0 +1,197 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2023 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 halo2_proofs::{
+    circuit::{AssignedCell, Layouter, Value},
+    pasta::group::ff::WithSmallOrderMulGroup,
+    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector},
+    poly::Rotation,
+};
+
+use super::is_zero::{IsZeroChip, IsZeroConfig};
+
+#[derive(Clone, Debug)]
+pub struct ZeroCondConfig<F> {
+    selector: Selector,
+    a: Column<Advice>,
+    b: Column<Advice>,
+    _z: Column<Advice>,
+    is_zero: IsZeroConfig<F>,
+    output: Column<Advice>,
+}
+
+#[derive(Clone, Debug)]
+pub struct ZeroCondChip<F: WithSmallOrderMulGroup<3> + Ord> {
+    config: ZeroCondConfig<F>,
+}
+
+impl<F: WithSmallOrderMulGroup<3> + Ord> ZeroCondChip<F> {
+    pub fn construct(config: ZeroCondConfig<F>) -> Self {
+        Self { config }
+    }
+
+    pub fn configure(
+        meta: &mut ConstraintSystem<F>,
+        advices: [Column<Advice>; 4],
+    ) -> ZeroCondConfig<F> {
+        for i in advices {
+            meta.enable_equality(i);
+        }
+
+        let selector = meta.selector();
+
+        let is_zero = IsZeroChip::configure(
+            meta,
+            |meta| meta.query_selector(selector),
+            |meta| meta.query_advice(advices[0], Rotation::cur()),
+            advices[2],
+        );
+
+        meta.create_gate("f(a, b) = if a == 0 {a} else {b}", |meta| {
+            let s = meta.query_selector(selector);
+            let a = meta.query_advice(advices[0], Rotation::cur());
+            let b = meta.query_advice(advices[1], Rotation::cur());
+            let output = meta.query_advice(advices[3], Rotation::cur());
+
+            let one = Expression::Constant(F::ONE);
+            vec![
+                s * (is_zero.expr() * (output.clone() - a) + (one - is_zero.expr()) * (output - b)),
+            ]
+        });
+
+        ZeroCondConfig {
+            selector,
+            a: advices[0],
+            b: advices[1],
+            _z: advices[2],
+            is_zero,
+            output: advices[3],
+        }
+    }
+
+    pub fn assign(
+        &self,
+        mut layouter: impl Layouter<F>,
+        a: AssignedCell<F, F>,
+        b: AssignedCell<F, F>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        let is_zero_chip = IsZeroChip::construct(self.config.is_zero.clone());
+
+        let out = layouter.assign_region(
+            || "f(a, b) = if a == 0 {a} else {b}",
+            |mut region| {
+                self.config.selector.enable(&mut region, 0)?;
+                let a = a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
+                let b = b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
+                is_zero_chip.assign(&mut region, 0, a.value().copied())?;
+
+                let output = if a.value().copied().to_field() == Value::known(F::ZERO.into()) {
+                    a.value().copied()
+                } else {
+                    b.value().copied()
+                };
+
+                let cell = region.assign_advice(|| "output", self.config.output, 0, || output)?;
+                Ok(cell)
+            },
+        )?;
+
+        Ok(out)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::zk::assign_free_advice;
+    use halo2_proofs::{
+        circuit::floor_planner::V1,
+        dev::MockProver,
+        pasta::Fp,
+        plonk::{Circuit, Instance},
+    };
+
+    #[derive(Default)]
+    struct MyCircuit {
+        a: Value<Fp>,
+        b: Value<Fp>,
+    }
+
+    impl Circuit<Fp> for MyCircuit {
+        type Config = (ZeroCondConfig<Fp>, [Column<Advice>; 5], Column<Instance>);
+        type FloorPlanner = V1;
+
+        fn without_witnesses(&self) -> Self {
+            Self::default()
+        }
+
+        fn configure(meta: &mut ConstraintSystem<Fp>) -> Self::Config {
+            let advices = [
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+                meta.advice_column(),
+            ];
+            for i in advices {
+                meta.enable_equality(i);
+            }
+
+            let instance = meta.instance_column();
+            meta.enable_equality(instance);
+
+            let zcc = ZeroCondChip::configure(meta, advices[1..5].try_into().unwrap());
+
+            (zcc, advices, instance)
+        }
+
+        fn synthesize(
+            &self,
+            config: Self::Config,
+            mut layouter: impl Layouter<Fp>,
+        ) -> Result<(), Error> {
+            let a = assign_free_advice(layouter.namespace(|| "load a"), config.1[0], self.a)?;
+            let b = assign_free_advice(layouter.namespace(|| "load b"), config.1[0], self.b)?;
+
+            let zcc = ZeroCondChip::construct(config.0);
+            let output = zcc.assign(layouter.namespace(|| "zero_cond"), a, b)?;
+            layouter.constrain_instance(output.cell(), config.2, 0)?;
+
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn zero_cond() -> crate::Result<()> {
+        let a = Fp::from(0);
+        let b = Fp::from(69);
+        let p_circuit = MyCircuit { a: Value::known(a), b: Value::known(b) };
+        let public_inputs = vec![a];
+        let prover = MockProver::run(3, &p_circuit, vec![public_inputs])?;
+        prover.assert_satisfied();
+
+        let a = Fp::from(12);
+        let b = Fp::from(42);
+        let p_circuit = MyCircuit { a: Value::known(a), b: Value::known(b) };
+        let public_inputs = vec![b];
+        let prover = MockProver::run(3, &p_circuit, vec![public_inputs])?;
+        prover.assert_satisfied();
+
+        Ok(())
+    }
+}