darkfi/zk/gadget/
arithmetic.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 std::marker::PhantomData;
20
21use 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
30pub 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)]
58pub 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
74pub struct ArithChip<F> {
75    config: ArithConfig,
76    _marker: PhantomData<F>,
77}
78
79impl<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
92impl<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
139impl<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)]
202mod 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}