/ src / zk / gadget / arithmetic.rs
arithmetic.rs
  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 std::marker::PhantomData;
 20  
 21  use halo2_proofs::{
 22      circuit::{AssignedCell, Chip, Layouter},
 23      pasta::group::ff::WithSmallOrderMulGroup,
 24      plonk,
 25      plonk::{Advice, Column, ConstraintSystem, Constraints, Selector},
 26      poly::Rotation,
 27  };
 28  
 29  /// Arithmetic instructions implemented in the chip
 30  pub trait ArithInstruction<F: WithSmallOrderMulGroup<3> + Ord>: Chip<F> {
 31      /// Add two field elements and return their sum
 32      fn add(
 33          &self,
 34          layouter: impl Layouter<F>,
 35          a: &AssignedCell<F, F>,
 36          b: &AssignedCell<F, F>,
 37      ) -> Result<AssignedCell<F, F>, plonk::Error>;
 38  
 39      /// Subtract two field elements and return their difference
 40      fn sub(
 41          &self,
 42          layouter: impl Layouter<F>,
 43          a: &AssignedCell<F, F>,
 44          b: &AssignedCell<F, F>,
 45      ) -> Result<AssignedCell<F, F>, plonk::Error>;
 46  
 47      /// Multiply two field elements and return their product
 48      fn mul(
 49          &self,
 50          layouter: impl Layouter<F>,
 51          a: &AssignedCell<F, F>,
 52          b: &AssignedCell<F, F>,
 53      ) -> Result<AssignedCell<F, F>, plonk::Error>;
 54  }
 55  
 56  /// Configuration for the Arithmetic Chip
 57  #[derive(Clone, Debug)]
 58  pub struct ArithConfig {
 59      /// lhs
 60      a: Column<Advice>,
 61      /// rhs
 62      b: Column<Advice>,
 63      /// out
 64      c: Column<Advice>,
 65      /// Selector for the `add` operation
 66      q_add: Selector,
 67      /// Selector for the `sub` operation
 68      q_sub: Selector,
 69      /// Selector for the `mul` operation
 70      q_mul: Selector,
 71  }
 72  
 73  /// Arithmetic Chip
 74  pub struct ArithChip<F> {
 75      config: ArithConfig,
 76      _marker: PhantomData<F>,
 77  }
 78  
 79  impl<F: WithSmallOrderMulGroup<3> + Ord> Chip<F> for ArithChip<F> {
 80      type Config = ArithConfig;
 81      type Loaded = ();
 82  
 83      fn config(&self) -> &Self::Config {
 84          &self.config
 85      }
 86  
 87      fn loaded(&self) -> &Self::Loaded {
 88          &()
 89      }
 90  }
 91  
 92  impl<F: WithSmallOrderMulGroup<3> + Ord> ArithChip<F> {
 93      /// Configure the Arithmetic chip with the given columns
 94      pub fn configure(
 95          meta: &mut ConstraintSystem<F>,
 96          a: Column<Advice>,
 97          b: Column<Advice>,
 98          c: Column<Advice>,
 99      ) -> ArithConfig {
100          let q_add = meta.selector();
101          let q_sub = meta.selector();
102          let q_mul = meta.selector();
103  
104          meta.create_gate("Field element addition: c = a + b", |meta| {
105              let q_add = meta.query_selector(q_add);
106              let a = meta.query_advice(a, Rotation::cur());
107              let b = meta.query_advice(b, Rotation::cur());
108              let c = meta.query_advice(c, Rotation::cur());
109  
110              Constraints::with_selector(q_add, Some(a + b - c))
111          });
112  
113          meta.create_gate("Field element subtraction: c = a - b", |meta| {
114              let q_sub = meta.query_selector(q_sub);
115              let a = meta.query_advice(a, Rotation::cur());
116              let b = meta.query_advice(b, Rotation::cur());
117              let c = meta.query_advice(c, Rotation::cur());
118  
119              Constraints::with_selector(q_sub, Some(a - b - c))
120          });
121  
122          meta.create_gate("Field element multiplication: c = a * b", |meta| {
123              let q_mul = meta.query_selector(q_mul);
124              let a = meta.query_advice(a, Rotation::cur());
125              let b = meta.query_advice(b, Rotation::cur());
126              let c = meta.query_advice(c, Rotation::cur());
127  
128              Constraints::with_selector(q_mul, Some(a * b - c))
129          });
130  
131          ArithConfig { a, b, c, q_add, q_sub, q_mul }
132      }
133  
134      pub fn construct(config: ArithConfig) -> Self {
135          Self { config, _marker: PhantomData }
136      }
137  }
138  
139  impl<F: WithSmallOrderMulGroup<3> + Ord> ArithInstruction<F> for ArithChip<F> {
140      fn add(
141          &self,
142          mut layouter: impl Layouter<F>,
143          a: &AssignedCell<F, F>,
144          b: &AssignedCell<F, F>,
145      ) -> Result<AssignedCell<F, F>, plonk::Error> {
146          layouter.assign_region(
147              || "c = a + b",
148              |mut region| {
149                  self.config.q_add.enable(&mut region, 0)?;
150  
151                  a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
152                  b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
153  
154                  let scalar_val = a.value().zip(b.value()).map(|(a, b)| *a + b);
155                  region.assign_advice(|| "c", self.config.c, 0, || scalar_val)
156              },
157          )
158      }
159  
160      fn sub(
161          &self,
162          mut layouter: impl Layouter<F>,
163          a: &AssignedCell<F, F>,
164          b: &AssignedCell<F, F>,
165      ) -> Result<AssignedCell<F, F>, plonk::Error> {
166          layouter.assign_region(
167              || "c = a - b",
168              |mut region| {
169                  self.config.q_sub.enable(&mut region, 0)?;
170  
171                  a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
172                  b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
173  
174                  let scalar_val = a.value().zip(b.value()).map(|(a, b)| *a - b);
175                  region.assign_advice(|| "c", self.config.c, 0, || scalar_val)
176              },
177          )
178      }
179  
180      fn mul(
181          &self,
182          mut layouter: impl Layouter<F>,
183          a: &AssignedCell<F, F>,
184          b: &AssignedCell<F, F>,
185      ) -> Result<AssignedCell<F, F>, plonk::Error> {
186          layouter.assign_region(
187              || "c = a * b",
188              |mut region| {
189                  self.config.q_mul.enable(&mut region, 0)?;
190  
191                  a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
192                  b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
193  
194                  let scalar_val = a.value().zip(b.value()).map(|(a, b)| *a * b);
195                  region.assign_advice(|| "c", self.config.c, 0, || scalar_val)
196              },
197          )
198      }
199  }
200  
201  #[cfg(test)]
202  mod tests {
203      use super::*;
204      use crate::zk::assign_free_advice;
205      use darkfi_sdk::pasta::pallas;
206      use halo2_proofs::{
207          arithmetic::Field,
208          circuit::{floor_planner, Value},
209          dev::{CircuitLayout, MockProver},
210          plonk::{Circuit, Instance as InstanceColumn},
211      };
212  
213      #[derive(Clone)]
214      struct ArithCircuitConfig {
215          primary: Column<InstanceColumn>,
216          advices: [Column<Advice>; 3],
217          arith_config: ArithConfig,
218      }
219  
220      #[derive(Default)]
221      struct ArithCircuit {
222          pub one: Value<pallas::Base>,
223          pub minus_one: Value<pallas::Base>,
224          pub factor: Value<pallas::Base>,
225      }
226  
227      impl Circuit<pallas::Base> for ArithCircuit {
228          type Config = ArithCircuitConfig;
229          type FloorPlanner = floor_planner::V1;
230          type Params = ();
231  
232          fn without_witnesses(&self) -> Self {
233              Self::default()
234          }
235  
236          fn configure(meta: &mut ConstraintSystem<pallas::Base>) -> Self::Config {
237              let advices = [meta.advice_column(), meta.advice_column(), meta.advice_column()];
238  
239              let primary = meta.instance_column();
240              meta.enable_equality(primary);
241  
242              for advice in advices.iter() {
243                  meta.enable_equality(*advice);
244              }
245  
246              let arith_config = ArithChip::configure(meta, advices[0], advices[1], advices[2]);
247  
248              Self::Config { primary, advices, arith_config }
249          }
250  
251          fn synthesize(
252              &self,
253              config: Self::Config,
254              mut layouter: impl Layouter<pallas::Base>,
255          ) -> Result<(), plonk::Error> {
256              let arith_chip = ArithChip::construct(config.arith_config.clone());
257  
258              let one = assign_free_advice(
259                  layouter.namespace(|| "Load Fp(1)"),
260                  config.advices[0],
261                  self.one,
262              )?;
263  
264              let minus_one = assign_free_advice(
265                  layouter.namespace(|| "Load Fp(-1)"),
266                  config.advices[1],
267                  self.minus_one,
268              )?;
269  
270              let factor = assign_free_advice(
271                  layouter.namespace(|| "Load Fp(factor)"),
272                  config.advices[2],
273                  self.factor,
274              )?;
275  
276              let diff =
277                  arith_chip.sub(layouter.namespace(|| "one - minus_one"), &one, &minus_one)?;
278              layouter.constrain_instance(diff.cell(), config.primary, 0)?;
279  
280              let zero =
281                  arith_chip.add(layouter.namespace(|| "one + minus_one"), &one, &minus_one)?;
282              layouter.constrain_instance(zero.cell(), config.primary, 1)?;
283  
284              let min_1_min_1 = arith_chip.add(
285                  layouter.namespace(|| "minus_one + minus_one"),
286                  &minus_one,
287                  &minus_one,
288              )?;
289              layouter.constrain_instance(min_1_min_1.cell(), config.primary, 2)?;
290  
291              let product =
292                  arith_chip.mul(layouter.namespace(|| "minus_one * factor"), &minus_one, &factor)?;
293              layouter.constrain_instance(product.cell(), config.primary, 3)?;
294  
295              Ok(())
296          }
297      }
298  
299      #[test]
300      fn arithmetic_chip() -> crate::Result<()> {
301          let one = pallas::Base::ONE;
302          let minus_one = -pallas::Base::ONE;
303          let factor = pallas::Base::from(644211);
304  
305          let public_inputs =
306              vec![one - minus_one, pallas::Base::ZERO, minus_one + minus_one, minus_one * factor];
307  
308          let circuit = ArithCircuit {
309              one: Value::known(one),
310              minus_one: Value::known(minus_one),
311              factor: Value::known(factor),
312          };
313  
314          use plotters::prelude::*;
315          let root = BitMapBackend::new("target/arithmetic_circuit_layout.png", (3840, 2160))
316              .into_drawing_area();
317          root.fill(&WHITE).unwrap();
318          let root = root.titled("Arithmetic Circuit Layout", ("sans-serif", 60)).unwrap();
319          CircuitLayout::default().render(4, &circuit, &root).unwrap();
320  
321          let prover = MockProver::run(4, &circuit, vec![public_inputs.clone()])?;
322          prover.assert_satisfied();
323  
324          Ok(())
325      }
326  }