1  /* This file is part of DarkFi (https://dark.fi)
  2   *
  3   * Copyright (C) 2020-2025 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  
 19  use halo2_proofs::{
 20      circuit::{AssignedCell, Layouter},
 21      pasta::group::ff::WithSmallOrderMulGroup,
 22      plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector},
 23      poly::Rotation,
 24  };
 25  
 26  use super::is_zero::{IsZeroChip, IsZeroConfig};
 27  
 28  #[derive(Clone, Debug)]
 29  pub 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)]
 38  pub struct ZeroCondChip<F: WithSmallOrderMulGroup<3> + Ord> {
 39      config: ZeroCondConfig<F>,
 40  }
 41  
 42  impl<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)]
119  mod 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  }