1use alloc::string::{String, ToString};
  2use alloc::vec::Vec;
  3use alloc::{format, vec};
  4use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
  5use core::marker::PhantomData;
  6
  7use plonky2::field::extension::Extendable;
  8use plonky2::field::packed::PackedField;
  9use plonky2::field::types::Field;
 10use plonky2::gates::gate::Gate;
 11use plonky2::gates::packed_util::PackedEvaluableBase;
 12use plonky2::gates::util::StridedConstraintConsumer;
 13use plonky2::hash::hash_types::RichField;
 14use plonky2::iop::ext_target::ExtensionTarget;
 15use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 16use plonky2::iop::target::Target;
 17use plonky2::iop::wire::Wire;
 18use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
 19use plonky2::plonk::circuit_builder::CircuitBuilder;
 20use plonky2::plonk::circuit_data::CircuitConfig;
 21use plonky2::plonk::vars::{
 22    EvaluationTargets, EvaluationVars, EvaluationVarsBase, EvaluationVarsBaseBatch,
 23    EvaluationVarsBasePacked,
 24};
 25
 26/// A gate to perform a subtraction on 32-bit limbs: given `x`, `y`, and `borrow`, it returns
 27/// the result `x - y - borrow` and, if this underflows, a new `borrow`. Inputs are not range-checked.
 28#[derive(Copy, Clone, Debug)]
 29pub struct U32SubtractionGate<F: RichField + Extendable<D>, const D: usize> {
 30    pub num_ops: usize,
 31    _phantom: PhantomData<F>,
 32}
 33
 34impl<F: RichField + Extendable<D>, const D: usize> U32SubtractionGate<F, D> {
 35    pub fn new_from_config(config: &CircuitConfig) -> Self {
 36        Self {
 37            num_ops: Self::num_ops(config),
 38            _phantom: PhantomData,
 39        }
 40    }
 41
 42    pub(crate) fn num_ops(config: &CircuitConfig) -> usize {
 43        let wires_per_op = 5 + Self::num_limbs();
 44        let routed_wires_per_op = 5;
 45        (config.num_wires / wires_per_op).min(config.num_routed_wires / routed_wires_per_op)
 46    }
 47
 48    pub fn wire_ith_input_x(&self, i: usize) -> usize {
 49        debug_assert!(i < self.num_ops);
 50        5 * i
 51    }
 52    pub fn wire_ith_input_y(&self, i: usize) -> usize {
 53        debug_assert!(i < self.num_ops);
 54        5 * i + 1
 55    }
 56    pub fn wire_ith_input_borrow(&self, i: usize) -> usize {
 57        debug_assert!(i < self.num_ops);
 58        5 * i + 2
 59    }
 60
 61    pub fn wire_ith_output_result(&self, i: usize) -> usize {
 62        debug_assert!(i < self.num_ops);
 63        5 * i + 3
 64    }
 65    pub fn wire_ith_output_borrow(&self, i: usize) -> usize {
 66        debug_assert!(i < self.num_ops);
 67        5 * i + 4
 68    }
 69
 70    pub fn limb_bits() -> usize {
 71        2
 72    }
 73    // We have limbs for the 32 bits of `output_result`.
 74    pub fn num_limbs() -> usize {
 75        32 / Self::limb_bits()
 76    }
 77
 78    pub fn wire_ith_output_jth_limb(&self, i: usize, j: usize) -> usize {
 79        debug_assert!(i < self.num_ops);
 80        debug_assert!(j < Self::num_limbs());
 81        5 * self.num_ops + Self::num_limbs() * i + j
 82    }
 83}
 84
 85impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for U32SubtractionGate<F, D> {
 86    fn id(&self) -> String {
 87        format!("{self:?}")
 88    }
 89
 90    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
 91        dst.write_usize(self.num_ops)
 92    }
 93
 94    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
 95        let num_ops = src.read_usize()?;
 96        Ok(Self { num_ops, _phantom: PhantomData })
 97    }
 98
 99    fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
100        let mut constraints = Vec::with_capacity(self.num_constraints());
101        for i in 0..self.num_ops {
102            let input_x = vars.local_wires[self.wire_ith_input_x(i)];
103            let input_y = vars.local_wires[self.wire_ith_input_y(i)];
104            let input_borrow = vars.local_wires[self.wire_ith_input_borrow(i)];
105
106            let result_initial = input_x - input_y - input_borrow;
107            let base = F::Extension::from_canonical_u64(1 << 32u64);
108
109            let output_result = vars.local_wires[self.wire_ith_output_result(i)];
110            let output_borrow = vars.local_wires[self.wire_ith_output_borrow(i)];
111
112            constraints.push(output_result - (result_initial + base * output_borrow));
113
114            // Range-check output_result to be at most 32 bits.
115            let mut combined_limbs = F::Extension::ZERO;
116            let limb_base = F::Extension::from_canonical_u64(1u64 << Self::limb_bits());
117            for j in (0..Self::num_limbs()).rev() {
118                let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
119                let max_limb = 1 << Self::limb_bits();
120                let product = (0..max_limb)
121                    .map(|x| this_limb - F::Extension::from_canonical_usize(x))
122                    .product();
123                constraints.push(product);
124
125                combined_limbs = limb_base * combined_limbs + this_limb;
126            }
127            constraints.push(combined_limbs - output_result);
128
129            // Range-check output_borrow to be one bit.
130            constraints.push(output_borrow * (F::Extension::ONE - output_borrow));
131        }
132
133        constraints
134    }
135
136    fn eval_unfiltered_base_one(
137        &self,
138        _vars: EvaluationVarsBase<F>,
139        _yield_constr: StridedConstraintConsumer<F>,
140    ) {
141        panic!("use eval_unfiltered_base_packed instead");
142    }
143
144    fn eval_unfiltered_base_batch(&self, vars_base: EvaluationVarsBaseBatch<F>) -> Vec<F> {
145        self.eval_unfiltered_base_batch_packed(vars_base)
146    }
147
148    fn eval_unfiltered_circuit(
149        &self,
150        builder: &mut CircuitBuilder<F, D>,
151        vars: EvaluationTargets<D>,
152    ) -> Vec<ExtensionTarget<D>> {
153        let mut constraints = Vec::with_capacity(self.num_constraints());
154        for i in 0..self.num_ops {
155            let input_x = vars.local_wires[self.wire_ith_input_x(i)];
156            let input_y = vars.local_wires[self.wire_ith_input_y(i)];
157            let input_borrow = vars.local_wires[self.wire_ith_input_borrow(i)];
158
159            let diff = builder.sub_extension(input_x, input_y);
160            let result_initial = builder.sub_extension(diff, input_borrow);
161            let base = builder.constant_extension(F::Extension::from_canonical_u64(1 << 32u64));
162
163            let output_result = vars.local_wires[self.wire_ith_output_result(i)];
164            let output_borrow = vars.local_wires[self.wire_ith_output_borrow(i)];
165
166            let computed_output = builder.mul_add_extension(base, output_borrow, result_initial);
167            constraints.push(builder.sub_extension(output_result, computed_output));
168
169            // Range-check output_result to be at most 32 bits.
170            let mut combined_limbs = builder.zero_extension();
171            let limb_base = builder
172                .constant_extension(F::Extension::from_canonical_u64(1u64 << Self::limb_bits()));
173            for j in (0..Self::num_limbs()).rev() {
174                let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
175                let max_limb = 1 << Self::limb_bits();
176                let mut product = builder.one_extension();
177                for x in 0..max_limb {
178                    let x_target =
179                        builder.constant_extension(F::Extension::from_canonical_usize(x));
180                    let diff = builder.sub_extension(this_limb, x_target);
181                    product = builder.mul_extension(product, diff);
182                }
183                constraints.push(product);
184
185                combined_limbs = builder.mul_add_extension(limb_base, combined_limbs, this_limb);
186            }
187            constraints.push(builder.sub_extension(combined_limbs, output_result));
188
189            // Range-check output_borrow to be one bit.
190            let one = builder.one_extension();
191            let not_borrow = builder.sub_extension(one, output_borrow);
192            constraints.push(builder.mul_extension(output_borrow, not_borrow));
193        }
194
195        constraints
196    }
197
198    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
199        (0..self.num_ops)
200            .map(|i| {
201                let g = WitnessGeneratorRef::new(
202                    U32SubtractionGenerator {
203                        gate: *self,
204                        row,
205                        i,
206                        _phantom: PhantomData,
207                    }
208                    .adapter(),
209                );
210                g
211            })
212            .collect()
213    }
214
215    fn num_wires(&self) -> usize {
216        self.num_ops * (5 + Self::num_limbs())
217    }
218
219    fn num_constants(&self) -> usize {
220        0
221    }
222
223    fn degree(&self) -> usize {
224        1 << Self::limb_bits()
225    }
226
227    fn num_constraints(&self) -> usize {
228        self.num_ops * (3 + Self::num_limbs())
229    }
230}
231
232impl<F: RichField + Extendable<D>, const D: usize> PackedEvaluableBase<F, D>
233    for U32SubtractionGate<F, D>
234{
235    fn eval_unfiltered_base_packed<P: PackedField<Scalar = F>>(
236        &self,
237        vars: EvaluationVarsBasePacked<P>,
238        mut yield_constr: StridedConstraintConsumer<P>,
239    ) {
240        for i in 0..self.num_ops {
241            let input_x = vars.local_wires[self.wire_ith_input_x(i)];
242            let input_y = vars.local_wires[self.wire_ith_input_y(i)];
243            let input_borrow = vars.local_wires[self.wire_ith_input_borrow(i)];
244
245            let result_initial = input_x - input_y - input_borrow;
246            let base = F::from_canonical_u64(1 << 32u64);
247
248            let output_result = vars.local_wires[self.wire_ith_output_result(i)];
249            let output_borrow = vars.local_wires[self.wire_ith_output_borrow(i)];
250
251            yield_constr.one(output_result - (result_initial + output_borrow * base));
252
253            // Range-check output_result to be at most 32 bits.
254            let mut combined_limbs = P::ZEROS;
255            let limb_base = F::from_canonical_u64(1u64 << Self::limb_bits());
256            for j in (0..Self::num_limbs()).rev() {
257                let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
258                let max_limb = 1 << Self::limb_bits();
259                let product = (0..max_limb)
260                    .map(|x| this_limb - F::from_canonical_usize(x))
261                    .product();
262                yield_constr.one(product);
263
264                combined_limbs = combined_limbs * limb_base + this_limb;
265            }
266            yield_constr.one(combined_limbs - output_result);
267
268            // Range-check output_borrow to be one bit.
269            yield_constr.one(output_borrow * (P::ONES - output_borrow));
270        }
271    }
272}
273
274#[derive(Clone, Debug)]
275struct U32SubtractionGenerator<F: RichField + Extendable<D>, const D: usize> {
276    gate: U32SubtractionGate<F, D>,
277    row: usize,
278    i: usize,
279    _phantom: PhantomData<F>,
280}
281
282impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
283    for U32SubtractionGenerator<F, D>
284{
285    fn id(&self) -> String {
286        "U32SubtractionGenerator".to_string()
287    }
288
289    fn dependencies(&self) -> Vec<Target> {
290        let local_target = |column| Target::wire(self.row, column);
291
292        vec![
293            local_target(self.gate.wire_ith_input_x(self.i)),
294            local_target(self.gate.wire_ith_input_y(self.i)),
295            local_target(self.gate.wire_ith_input_borrow(self.i)),
296        ]
297    }
298
299    fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
300        let local_wire = |column| Wire {
301            row: self.row,
302            column,
303        };
304
305        let get_local_wire = |column| witness.get_wire(local_wire(column));
306
307        let input_x = get_local_wire(self.gate.wire_ith_input_x(self.i));
308        let input_y = get_local_wire(self.gate.wire_ith_input_y(self.i));
309        let input_borrow = get_local_wire(self.gate.wire_ith_input_borrow(self.i));
310
311        let result_initial = input_x - input_y - input_borrow;
312        let result_initial_u64 = result_initial.to_canonical_u64();
313        let output_borrow = if result_initial_u64 > 1 << 32u64 {
314            F::ONE
315        } else {
316            F::ZERO
317        };
318
319        let base = F::from_canonical_u64(1 << 32u64);
320        let output_result = result_initial + base * output_borrow;
321
322        let output_result_wire = local_wire(self.gate.wire_ith_output_result(self.i));
323        let output_borrow_wire = local_wire(self.gate.wire_ith_output_borrow(self.i));
324
325        out_buffer.set_wire(output_result_wire, output_result);
326        out_buffer.set_wire(output_borrow_wire, output_borrow);
327
328        let output_result_u64 = output_result.to_canonical_u64();
329
330        let num_limbs = U32SubtractionGate::<F, D>::num_limbs();
331        let limb_base = 1 << U32SubtractionGate::<F, D>::limb_bits();
332        let output_limbs: Vec<_> = (0..num_limbs)
333            .scan(output_result_u64, |acc, _| {
334                let tmp = *acc % limb_base;
335                *acc /= limb_base;
336                Some(F::from_canonical_u64(tmp))
337            })
338            .collect();
339
340        for j in 0..num_limbs {
341            let wire = local_wire(self.gate.wire_ith_output_jth_limb(self.i, j));
342            out_buffer.set_wire(wire, output_limbs[j]);
343        }
344    }
345
346    fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
347        self.gate.serialize(dst)?;
348        dst.write_usize(self.row)?;
349        dst.write_usize(self.i)
350    }
351
352    fn deserialize(src: &mut Buffer) -> IoResult<Self> {
353        let gate = U32SubtractionGate::deserialize(src)?;
354        let row = src.read_usize()?;
355        let i = src.read_usize()?;
356        Ok(Self { gate, row, i, _phantom: PhantomData })
357    }
358
359}
360
361#[cfg(test)]
362mod tests {
363    use anyhow::Result;
364    use plonky2::field::extension::quartic::QuarticExtension;
365    use plonky2::field::goldilocks_field::GoldilocksField;
366    use plonky2::field::types::{PrimeField64, Sample};
367    use plonky2::gates::gate_testing::{test_eval_fns, test_low_degree};
368    use plonky2::hash::hash_types::HashOut;
369    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
370    use rand::rngs::OsRng;
371    use rand::Rng;
372
373    use super::*;
374
375    #[test]
376    fn low_degree() {
377        test_low_degree::<GoldilocksField, _, 4>(U32SubtractionGate::<GoldilocksField, 4> {
378            num_ops: 3,
379            _phantom: PhantomData,
380        })
381    }
382
383    #[test]
384    fn eval_fns() -> Result<()> {
385        const D: usize = 2;
386        type C = PoseidonGoldilocksConfig;
387        type F = <C as GenericConfig<D>>::F;
388        test_eval_fns::<F, C, _, D>(U32SubtractionGate::<GoldilocksField, D> {
389            num_ops: 3,
390            _phantom: PhantomData,
391        })
392    }
393
394    #[test]
395    fn test_gate_constraint() {
396        type F = GoldilocksField;
397        type FF = QuarticExtension<GoldilocksField>;
398        const D: usize = 4;
399        const NUM_U32_SUBTRACTION_OPS: usize = 3;
400
401        fn get_wires(inputs_x: Vec<u64>, inputs_y: Vec<u64>, borrows: Vec<u64>) -> Vec<FF> {
402            let mut v0 = Vec::new();
403            let mut v1 = Vec::new();
404
405            let limb_bits = U32SubtractionGate::<F, D>::limb_bits();
406            let num_limbs = U32SubtractionGate::<F, D>::num_limbs();
407            let limb_base = 1 << limb_bits;
408            for c in 0..NUM_U32_SUBTRACTION_OPS {
409                let input_x = F::from_canonical_u64(inputs_x[c]);
410                let input_y = F::from_canonical_u64(inputs_y[c]);
411                let input_borrow = F::from_canonical_u64(borrows[c]);
412
413                let result_initial = input_x - input_y - input_borrow;
414                let result_initial_u64 = result_initial.to_canonical_u64();
415                let output_borrow = if result_initial_u64 > 1 << 32u64 {
416                    F::ONE
417                } else {
418                    F::ZERO
419                };
420
421                let base = F::from_canonical_u64(1 << 32u64);
422                let output_result = result_initial + base * output_borrow;
423
424                let output_result_u64 = output_result.to_canonical_u64();
425
426                let mut output_limbs: Vec<_> = (0..num_limbs)
427                    .scan(output_result_u64, |acc, _| {
428                        let tmp = *acc % limb_base;
429                        *acc /= limb_base;
430                        Some(F::from_canonical_u64(tmp))
431                    })
432                    .collect();
433
434                v0.push(input_x);
435                v0.push(input_y);
436                v0.push(input_borrow);
437                v0.push(output_result);
438                v0.push(output_borrow);
439                v1.append(&mut output_limbs);
440            }
441
442            v0.iter().chain(v1.iter()).map(|&x| x.into()).collect()
443        }
444
445        let mut rng = OsRng;
446        let inputs_x = (0..NUM_U32_SUBTRACTION_OPS)
447            .map(|_| rng.gen::<u32>() as u64)
448            .collect();
449        let inputs_y = (0..NUM_U32_SUBTRACTION_OPS)
450            .map(|_| rng.gen::<u32>() as u64)
451            .collect();
452        let borrows = (0..NUM_U32_SUBTRACTION_OPS)
453            .map(|_| (rng.gen::<u32>() % 2) as u64)
454            .collect();
455
456        let gate = U32SubtractionGate::<F, D> {
457            num_ops: NUM_U32_SUBTRACTION_OPS,
458            _phantom: PhantomData,
459        };
460
461        let vars = EvaluationVars {
462            local_constants: &[],
463            local_wires: &get_wires(inputs_x, inputs_y, borrows),
464            public_inputs_hash: &HashOut::rand(),
465        };
466
467        assert!(
468            gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()),
469            "Gate constraints are not satisfied."
470        );
471    }
472}