darkfi/zk/gadget/
zero_cond.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2026 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use halo2_proofs::{
20    circuit::{AssignedCell, Layouter},
21    pasta::group::ff::WithSmallOrderMulGroup,
22    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector},
23    poly::Rotation,
24};
25
26use super::is_zero::{IsZeroChip, IsZeroConfig};
27
28#[derive(Clone, Debug)]
29pub struct ZeroCondConfig<F> {
30    selector: Selector,
31    a: Column<Advice>,
32    b: Column<Advice>,
33    is_zero: IsZeroConfig<F>,
34    output: Column<Advice>,
35}
36
37#[derive(Clone, Debug)]
38pub struct ZeroCondChip<F: WithSmallOrderMulGroup<3> + Ord> {
39    config: ZeroCondConfig<F>,
40}
41
42impl<F: WithSmallOrderMulGroup<3> + Ord> ZeroCondChip<F> {
43    pub fn construct(config: ZeroCondConfig<F>) -> Self {
44        Self { config }
45    }
46
47    /// Configure the chip.
48    ///
49    /// Advice columns:
50    /// * `[0]` - a
51    /// * `[1]` - b
52    /// * `[2]` - is_zero output
53    /// * `[3]` - zero_cond output
54    pub fn configure(
55        meta: &mut ConstraintSystem<F>,
56        advices: [Column<Advice>; 4],
57    ) -> ZeroCondConfig<F> {
58        for i in advices {
59            meta.enable_equality(i);
60        }
61
62        let selector = meta.selector();
63
64        let is_zero = IsZeroChip::configure(
65            meta,
66            |meta| meta.query_selector(selector),
67            |meta| meta.query_advice(advices[0], Rotation::cur()),
68            advices[2],
69        );
70
71        // NOTE: a is not used here because it already went into IsZero
72        meta.create_gate("f(a, b) = if a == 0 {a} else {b}", |meta| {
73            let s = meta.query_selector(selector);
74            let b = meta.query_advice(advices[1], Rotation::cur());
75            let output = meta.query_advice(advices[3], Rotation::cur());
76
77            let one = Expression::Constant(F::ONE);
78
79            vec![s * (is_zero.expr() * output.clone() + (one - is_zero.expr()) * (output - b))]
80        });
81
82        ZeroCondConfig { selector, a: advices[0], b: advices[1], is_zero, output: advices[3] }
83    }
84
85    pub fn assign(
86        &self,
87        mut layouter: impl Layouter<F>,
88        a: AssignedCell<F, F>,
89        b: AssignedCell<F, F>,
90    ) -> Result<AssignedCell<F, F>, Error> {
91        let is_zero_chip = IsZeroChip::construct(self.config.is_zero.clone());
92
93        let out = layouter.assign_region(
94            || "f(a, b) = if a == 0 {a} else {b}",
95            |mut region| {
96                self.config.selector.enable(&mut region, 0)?;
97                let a = a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
98                let b = b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
99                is_zero_chip.assign(&mut region, 0, a.value().copied())?;
100
101                let output = a.value().copied().to_field().zip(b.value().copied()).map(|(a, b)| {
102                    if a == F::ZERO.into() {
103                        F::ZERO
104                    } else {
105                        b
106                    }
107                });
108
109                let cell = region.assign_advice(|| "output", self.config.output, 0, || output)?;
110                Ok(cell)
111            },
112        )?;
113
114        Ok(out)
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::zk::assign_free_advice;
122    use halo2_proofs::{
123        circuit::{SimpleFloorPlanner, Value},
124        dev::MockProver,
125        pasta::Fp,
126        plonk::{Circuit, Instance},
127    };
128
129    #[derive(Default)]
130    struct MyCircuit {
131        a: Value<Fp>,
132        b: Value<Fp>,
133    }
134
135    impl Circuit<Fp> for MyCircuit {
136        type Config = (ZeroCondConfig<Fp>, [Column<Advice>; 5], Column<Instance>);
137        type FloorPlanner = SimpleFloorPlanner;
138        type Params = ();
139
140        fn without_witnesses(&self) -> Self {
141            Self::default()
142        }
143
144        fn configure(meta: &mut ConstraintSystem<Fp>) -> Self::Config {
145            let advices = [
146                meta.advice_column(),
147                meta.advice_column(),
148                meta.advice_column(),
149                meta.advice_column(),
150                meta.advice_column(),
151            ];
152            for i in advices {
153                meta.enable_equality(i);
154            }
155
156            let instance = meta.instance_column();
157            meta.enable_equality(instance);
158
159            let zcc = ZeroCondChip::configure(meta, advices[1..5].try_into().unwrap());
160
161            (zcc, advices, instance)
162        }
163
164        fn synthesize(
165            &self,
166            config: Self::Config,
167            mut layouter: impl Layouter<Fp>,
168        ) -> Result<(), Error> {
169            let a = assign_free_advice(layouter.namespace(|| "load a"), config.1[0], self.a)?;
170            let b = assign_free_advice(layouter.namespace(|| "load b"), config.1[0], self.b)?;
171
172            let zcc = ZeroCondChip::construct(config.0);
173            let output = zcc.assign(layouter.namespace(|| "zero_cond"), a, b)?;
174            layouter.constrain_instance(output.cell(), config.2, 0)?;
175
176            Ok(())
177        }
178    }
179
180    #[test]
181    fn zero_cond() {
182        let a = Fp::from(0);
183        let b = Fp::from(69);
184        let p_circuit = MyCircuit { a: Value::known(a), b: Value::known(b) };
185        let public_inputs = vec![a];
186        let prover = MockProver::run(3, &p_circuit, vec![public_inputs]).unwrap();
187        prover.assert_satisfied();
188
189        let a = Fp::from(12);
190        let b = Fp::from(42);
191        let p_circuit = MyCircuit { a: Value::known(a), b: Value::known(b) };
192        let public_inputs = vec![b];
193        let prover = MockProver::run(3, &p_circuit, vec![public_inputs]).unwrap();
194        prover.assert_satisfied();
195    }
196}
darkfi/zk/gadget/zero_cond.rs

darkfi/zk/gadget/
zero_cond.rs
