plonky2-u32/src/gates/arithmetic_u32.rs
src/gates/arithmetic_u32.rsRaw
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 itertools::unfold;
8use plonky2::field::extension::Extendable;
9use plonky2::field::packed::PackedField;
10use plonky2::field::types::Field;
11use plonky2::gates::gate::Gate;
12use plonky2::gates::packed_util::PackedEvaluableBase;
13use plonky2::gates::util::StridedConstraintConsumer;
14use plonky2::hash::hash_types::RichField;
15use plonky2::iop::ext_target::ExtensionTarget;
16use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
17use plonky2::iop::target::Target;
18use plonky2::iop::wire::Wire;
19use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
20use plonky2::plonk::circuit_builder::CircuitBuilder;
21use plonky2::plonk::circuit_data::CircuitConfig;
22use plonky2::plonk::vars::{
23 EvaluationTargets, EvaluationVars, EvaluationVarsBase, EvaluationVarsBaseBatch,
24 EvaluationVarsBasePacked,
25};
26
27/// A gate to perform a basic mul-add on 32-bit values (we assume they are range-checked beforehand).
28#[derive(Copy, Clone, Debug)]
29pub struct U32ArithmeticGate<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> U32ArithmeticGate<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 = Self::routed_wires_per_op() + Self::num_limbs();
44 (config.num_wires / wires_per_op).min(config.num_routed_wires / Self::routed_wires_per_op())
45 }
46
47 pub fn wire_ith_multiplicand_0(&self, i: usize) -> usize {
48 debug_assert!(i < self.num_ops);
49 Self::routed_wires_per_op() * i
50 }
51 pub fn wire_ith_multiplicand_1(&self, i: usize) -> usize {
52 debug_assert!(i < self.num_ops);
53 Self::routed_wires_per_op() * i + 1
54 }
55 pub fn wire_ith_addend(&self, i: usize) -> usize {
56 debug_assert!(i < self.num_ops);
57 Self::routed_wires_per_op() * i + 2
58 }
59
60 pub fn wire_ith_output_low_half(&self, i: usize) -> usize {
61 debug_assert!(i < self.num_ops);
62 Self::routed_wires_per_op() * i + 3
63 }
64
65 pub fn wire_ith_output_high_half(&self, i: usize) -> usize {
66 debug_assert!(i < self.num_ops);
67 Self::routed_wires_per_op() * i + 4
68 }
69
70 pub fn wire_ith_inverse(&self, i: usize) -> usize {
71 debug_assert!(i < self.num_ops);
72 Self::routed_wires_per_op() * i + 5
73 }
74
75 pub fn limb_bits() -> usize {
76 2
77 }
78 pub fn num_limbs() -> usize {
79 64 / Self::limb_bits()
80 }
81 pub fn routed_wires_per_op() -> usize {
82 6
83 }
84 pub fn wire_ith_output_jth_limb(&self, i: usize, j: usize) -> usize {
85 debug_assert!(i < self.num_ops);
86 debug_assert!(j < Self::num_limbs());
87 Self::routed_wires_per_op() * self.num_ops + Self::num_limbs() * i + j
88 }
89}
90
91impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for U32ArithmeticGate<F, D> {
92 fn id(&self) -> String {
93 format!("{self:?}")
94 }
95
96 fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
97 dst.write_usize(self.num_ops)
98 }
99
100 fn deserialize(src: &mut Buffer) -> IoResult<Self> {
101 let num_ops = src.read_usize()?;
102 Ok(Self { num_ops, _phantom: PhantomData })
103 }
104
105 fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
106 let mut constraints = Vec::with_capacity(self.num_constraints());
107 for i in 0..self.num_ops {
108 let multiplicand_0 = vars.local_wires[self.wire_ith_multiplicand_0(i)];
109 let multiplicand_1 = vars.local_wires[self.wire_ith_multiplicand_1(i)];
110 let addend = vars.local_wires[self.wire_ith_addend(i)];
111
112 let computed_output = multiplicand_0 * multiplicand_1 + addend;
113
114 let output_low = vars.local_wires[self.wire_ith_output_low_half(i)];
115 let output_high = vars.local_wires[self.wire_ith_output_high_half(i)];
116 let inverse = vars.local_wires[self.wire_ith_inverse(i)];
117
118 // Check canonicity of combined_output = output_high * 2^32 + output_low
119 let combined_output = {
120 let base = F::Extension::from_canonical_u64(1 << 32u64);
121 let one = F::Extension::ONE;
122 let u32_max = F::Extension::from_canonical_u32(u32::MAX);
123
124 // This is zero if and only if the high limb is `u32::MAX`.
125 // u32::MAX - output_high
126 let diff = u32_max - output_high;
127 // If this is zero, the diff is invertible, so the high limb is not `u32::MAX`.
128 // inverse * diff - 1
129 let hi_not_max = inverse * diff - one;
130 // If this is zero, either the high limb is not `u32::MAX`, or the low limb is zero.
131 // hi_not_max * limb_0_u32
132 let hi_not_max_or_lo_zero = hi_not_max * output_low;
133
134 constraints.push(hi_not_max_or_lo_zero);
135
136 output_high * base + output_low
137 };
138
139 constraints.push(combined_output - computed_output);
140
141 let mut combined_low_limbs = F::Extension::ZERO;
142 let mut combined_high_limbs = F::Extension::ZERO;
143 let midpoint = Self::num_limbs() / 2;
144 let base = F::Extension::from_canonical_u64(1u64 << Self::limb_bits());
145 for j in (0..Self::num_limbs()).rev() {
146 let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
147 let max_limb = 1 << Self::limb_bits();
148 let product = (0..max_limb)
149 .map(|x| this_limb - F::Extension::from_canonical_usize(x))
150 .product();
151 constraints.push(product);
152
153 if j < midpoint {
154 combined_low_limbs = base * combined_low_limbs + this_limb;
155 } else {
156 combined_high_limbs = base * combined_high_limbs + this_limb;
157 }
158 }
159 constraints.push(combined_low_limbs - output_low);
160 constraints.push(combined_high_limbs - output_high);
161 }
162
163 constraints
164 }
165
166 fn eval_unfiltered_base_one(
167 &self,
168 _vars: EvaluationVarsBase<F>,
169 _yield_constr: StridedConstraintConsumer<F>,
170 ) {
171 panic!("use eval_unfiltered_base_packed instead");
172 }
173
174 fn eval_unfiltered_base_batch(&self, vars_base: EvaluationVarsBaseBatch<F>) -> Vec<F> {
175 self.eval_unfiltered_base_batch_packed(vars_base)
176 }
177
178 fn eval_unfiltered_circuit(
179 &self,
180 builder: &mut CircuitBuilder<F, D>,
181 vars: EvaluationTargets<D>,
182 ) -> Vec<ExtensionTarget<D>> {
183 let mut constraints = Vec::with_capacity(self.num_constraints());
184
185 for i in 0..self.num_ops {
186 let multiplicand_0 = vars.local_wires[self.wire_ith_multiplicand_0(i)];
187 let multiplicand_1 = vars.local_wires[self.wire_ith_multiplicand_1(i)];
188 let addend = vars.local_wires[self.wire_ith_addend(i)];
189
190 let computed_output = builder.mul_add_extension(multiplicand_0, multiplicand_1, addend);
191
192 let output_low = vars.local_wires[self.wire_ith_output_low_half(i)];
193 let output_high = vars.local_wires[self.wire_ith_output_high_half(i)];
194 let inverse = vars.local_wires[self.wire_ith_inverse(i)];
195
196 // Check canonicity of combined_output = output_high * 2^32 + output_low
197 let combined_output = {
198 let base: F::Extension = F::from_canonical_u64(1 << 32u64).into();
199 let base_target = builder.constant_extension(base);
200 let one = builder.one_extension();
201 let u32_max =
202 builder.constant_extension(F::Extension::from_canonical_u32(u32::MAX));
203
204 // This is zero if and only if the high limb is `u32::MAX`.
205 let diff = builder.sub_extension(u32_max, output_high);
206 // If this is zero, the diff is invertible, so the high limb is not `u32::MAX`.
207 let hi_not_max = builder.mul_sub_extension(inverse, diff, one);
208 // If this is zero, either the high limb is not `u32::MAX`, or the low limb is zero.
209 let hi_not_max_or_lo_zero = builder.mul_extension(hi_not_max, output_low);
210
211 constraints.push(hi_not_max_or_lo_zero);
212
213 builder.mul_add_extension(output_high, base_target, output_low)
214 };
215
216 constraints.push(builder.sub_extension(combined_output, computed_output));
217
218 let mut combined_low_limbs = builder.zero_extension();
219 let mut combined_high_limbs = builder.zero_extension();
220 let midpoint = Self::num_limbs() / 2;
221 let base = builder
222 .constant_extension(F::Extension::from_canonical_u64(1u64 << Self::limb_bits()));
223 for j in (0..Self::num_limbs()).rev() {
224 let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
225 let max_limb = 1 << Self::limb_bits();
226
227 let mut product = builder.one_extension();
228 for x in 0..max_limb {
229 let x_target =
230 builder.constant_extension(F::Extension::from_canonical_usize(x));
231 let diff = builder.sub_extension(this_limb, x_target);
232 product = builder.mul_extension(product, diff);
233 }
234 constraints.push(product);
235
236 if j < midpoint {
237 combined_low_limbs =
238 builder.mul_add_extension(base, combined_low_limbs, this_limb);
239 } else {
240 combined_high_limbs =
241 builder.mul_add_extension(base, combined_high_limbs, this_limb);
242 }
243 }
244
245 constraints.push(builder.sub_extension(combined_low_limbs, output_low));
246 constraints.push(builder.sub_extension(combined_high_limbs, output_high));
247 }
248
249 constraints
250 }
251
252 fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F>> {
253 (0..self.num_ops)
254 .map(|i| {
255 let g: WitnessGeneratorRef<F> = WitnessGeneratorRef::new(
256 U32ArithmeticGenerator {
257 gate: *self,
258 row,
259 i,
260 _phantom: PhantomData,
261 }
262 .adapter(),
263 );
264 g
265 })
266 .collect()
267 }
268
269 fn num_wires(&self) -> usize {
270 self.num_ops * (Self::routed_wires_per_op() + Self::num_limbs())
271 }
272
273 fn num_constants(&self) -> usize {
274 0
275 }
276
277 fn degree(&self) -> usize {
278 1 << Self::limb_bits()
279 }
280
281 fn num_constraints(&self) -> usize {
282 self.num_ops * (4 + Self::num_limbs())
283 }
284}
285
286impl<F: RichField + Extendable<D>, const D: usize> PackedEvaluableBase<F, D>
287 for U32ArithmeticGate<F, D>
288{
289 fn eval_unfiltered_base_packed<P: PackedField<Scalar = F>>(
290 &self,
291 vars: EvaluationVarsBasePacked<P>,
292 mut yield_constr: StridedConstraintConsumer<P>,
293 ) {
294 for i in 0..self.num_ops {
295 let multiplicand_0 = vars.local_wires[self.wire_ith_multiplicand_0(i)];
296 let multiplicand_1 = vars.local_wires[self.wire_ith_multiplicand_1(i)];
297 let addend = vars.local_wires[self.wire_ith_addend(i)];
298
299 let computed_output = multiplicand_0 * multiplicand_1 + addend;
300
301 let output_low = vars.local_wires[self.wire_ith_output_low_half(i)];
302 let output_high = vars.local_wires[self.wire_ith_output_high_half(i)];
303 let inverse = vars.local_wires[self.wire_ith_inverse(i)];
304
305 let combined_output = {
306 let base = P::from(F::from_canonical_u64(1 << 32u64));
307 let one = P::ONES;
308 let u32_max = P::from(F::from_canonical_u32(u32::MAX));
309
310 // This is zero if and only if the high limb is `u32::MAX`.
311 // u32::MAX - output_high
312 let diff = u32_max - output_high;
313 // If this is zero, the diff is invertible, so the high limb is not `u32::MAX`.
314 // inverse * diff - 1
315 let hi_not_max = inverse * diff - one;
316 // If this is zero, either the high limb is not `u32::MAX`, or the low limb is zero.
317 // hi_not_max * limb_0_u32
318 let hi_not_max_or_lo_zero = hi_not_max * output_low;
319
320 yield_constr.one(hi_not_max_or_lo_zero);
321
322 output_high * base + output_low
323 };
324
325 yield_constr.one(combined_output - computed_output);
326
327 let mut combined_low_limbs = P::ZEROS;
328 let mut combined_high_limbs = P::ZEROS;
329 let midpoint = Self::num_limbs() / 2;
330 let base = F::from_canonical_u64(1u64 << Self::limb_bits());
331 for j in (0..Self::num_limbs()).rev() {
332 let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)];
333 let max_limb = 1 << Self::limb_bits();
334 let product = (0..max_limb)
335 .map(|x| this_limb - F::from_canonical_usize(x))
336 .product();
337 yield_constr.one(product);
338
339 if j < midpoint {
340 combined_low_limbs = combined_low_limbs * base + this_limb;
341 } else {
342 combined_high_limbs = combined_high_limbs * base + this_limb;
343 }
344 }
345 yield_constr.one(combined_low_limbs - output_low);
346 yield_constr.one(combined_high_limbs - output_high);
347 }
348 }
349}
350
351#[derive(Clone, Debug)]
352struct U32ArithmeticGenerator<F: RichField + Extendable<D>, const D: usize> {
353 gate: U32ArithmeticGate<F, D>,
354 row: usize,
355 i: usize,
356 _phantom: PhantomData<F>,
357}
358
359impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
360 for U32ArithmeticGenerator<F, D>
361{
362 fn id(&self) -> String {
363 "U32ArithmeticGenerator".to_string()
364 }
365
366 fn dependencies(&self) -> Vec<Target> {
367 let local_target = |column| Target::wire(self.row, column);
368
369 vec![
370 local_target(self.gate.wire_ith_multiplicand_0(self.i)),
371 local_target(self.gate.wire_ith_multiplicand_1(self.i)),
372 local_target(self.gate.wire_ith_addend(self.i)),
373 ]
374 }
375
376 fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
377 let local_wire = |column| Wire {
378 row: self.row,
379 column,
380 };
381
382 let get_local_wire = |column| witness.get_wire(local_wire(column));
383
384 let multiplicand_0 = get_local_wire(self.gate.wire_ith_multiplicand_0(self.i));
385 let multiplicand_1 = get_local_wire(self.gate.wire_ith_multiplicand_1(self.i));
386 let addend = get_local_wire(self.gate.wire_ith_addend(self.i));
387
388 let output = multiplicand_0 * multiplicand_1 + addend;
389 let mut output_u64 = output.to_canonical_u64();
390
391 let output_high_u64 = output_u64 >> 32;
392 let output_low_u64 = output_u64 & ((1 << 32) - 1);
393
394 let output_high = F::from_canonical_u64(output_high_u64);
395 let output_low = F::from_canonical_u64(output_low_u64);
396
397 let output_high_wire = local_wire(self.gate.wire_ith_output_high_half(self.i));
398 let output_low_wire = local_wire(self.gate.wire_ith_output_low_half(self.i));
399
400 out_buffer.set_wire(output_high_wire, output_high);
401 out_buffer.set_wire(output_low_wire, output_low);
402
403 let diff = u32::MAX as u64 - output_high_u64;
404 let inverse = if diff == 0 {
405 F::ZERO
406 } else {
407 F::from_canonical_u64(diff).inverse()
408 };
409 let inverse_wire = local_wire(self.gate.wire_ith_inverse(self.i));
410 out_buffer.set_wire(inverse_wire, inverse);
411
412 let num_limbs = U32ArithmeticGate::<F, D>::num_limbs();
413 let limb_base = 1 << U32ArithmeticGate::<F, D>::limb_bits();
414 let output_limbs_u64 = unfold((), move |_| {
415 let ret = output_u64 % limb_base;
416 output_u64 /= limb_base;
417 Some(ret)
418 })
419 .take(num_limbs);
420 let output_limbs_f = output_limbs_u64.map(F::from_canonical_u64);
421
422 for (j, output_limb) in output_limbs_f.enumerate() {
423 let wire = local_wire(self.gate.wire_ith_output_jth_limb(self.i, j));
424 out_buffer.set_wire(wire, output_limb);
425 }
426 }
427
428 fn serialize(&self, dst: &mut Vec<u8>) -> IoResult<()> {
429 self.gate.serialize(dst)?;
430 dst.write_usize(self.row)?;
431 dst.write_usize(self.i)
432 }
433
434 fn deserialize(src: &mut Buffer) -> IoResult<Self> {
435 let gate = U32ArithmeticGate::deserialize(src)?;
436 let row = src.read_usize()?;
437 let i = src.read_usize()?;
438 Ok(Self { gate, row, i, _phantom: PhantomData })
439 }
440}
441
442#[cfg(test)]
443mod tests {
444 use anyhow::Result;
445 use plonky2::field::goldilocks_field::GoldilocksField;
446 use plonky2::field::types::Sample;
447 use plonky2::gates::gate_testing::{test_eval_fns, test_low_degree};
448 use plonky2::hash::hash_types::HashOut;
449 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
450 use rand::rngs::OsRng;
451 use rand::Rng;
452
453 use super::*;
454
455 #[test]
456 fn low_degree() {
457 test_low_degree::<GoldilocksField, _, 4>(U32ArithmeticGate::<GoldilocksField, 4> {
458 num_ops: 3,
459 _phantom: PhantomData,
460 })
461 }
462
463 #[test]
464 fn eval_fns() -> Result<()> {
465 const D: usize = 2;
466 type C = PoseidonGoldilocksConfig;
467 type F = <C as GenericConfig<D>>::F;
468 test_eval_fns::<F, C, _, D>(U32ArithmeticGate::<GoldilocksField, D> {
469 num_ops: 3,
470 _phantom: PhantomData,
471 })
472 }
473
474 fn get_wires<
475 F: RichField + Extendable<D>,
476 FF: From<F>,
477 const D: usize,
478 const NUM_U32_ARITHMETIC_OPS: usize,
479 >(
480 multiplicands_0: Vec<u64>,
481 multiplicands_1: Vec<u64>,
482 addends: Vec<u64>,
483 ) -> Vec<FF> {
484 let mut v0 = Vec::new();
485 let mut v1 = Vec::new();
486
487 let limb_bits = U32ArithmeticGate::<F, D>::limb_bits();
488 let num_limbs = U32ArithmeticGate::<F, D>::num_limbs();
489 let limb_base = 1 << limb_bits;
490 for c in 0..NUM_U32_ARITHMETIC_OPS {
491 let m0 = multiplicands_0[c];
492 let m1 = multiplicands_1[c];
493 let a = addends[c];
494
495 let mut output = m0 * m1 + a;
496 let output_low = output & ((1 << 32) - 1);
497 let output_high = output >> 32;
498 let diff = u32::MAX as u64 - output_high;
499 let inverse = if diff == 0 {
500 F::ZERO
501 } else {
502 F::from_canonical_u64(diff).inverse()
503 };
504
505 let mut output_limbs = Vec::with_capacity(num_limbs);
506 for _i in 0..num_limbs {
507 output_limbs.push(output % limb_base);
508 output /= limb_base;
509 }
510 let mut output_limbs_f: Vec<_> = output_limbs
511 .into_iter()
512 .map(F::from_canonical_u64)
513 .collect();
514
515 v0.push(F::from_canonical_u64(m0));
516 v0.push(F::from_canonical_u64(m1));
517 v0.push(F::from_noncanonical_u64(a));
518 v0.push(F::from_canonical_u64(output_low));
519 v0.push(F::from_canonical_u64(output_high));
520 v0.push(inverse);
521 v1.append(&mut output_limbs_f);
522 }
523
524 v0.iter().chain(v1.iter()).map(|&x| x.into()).collect()
525 }
526
527 #[test]
528 fn test_gate_constraint() {
529 const D: usize = 2;
530 type C = PoseidonGoldilocksConfig;
531 type F = <C as GenericConfig<D>>::F;
532 type FF = <C as GenericConfig<D>>::FE;
533 const NUM_U32_ARITHMETIC_OPS: usize = 3;
534
535 let mut rng = OsRng;
536 let multiplicands_0: Vec<_> = (0..NUM_U32_ARITHMETIC_OPS)
537 .map(|_| rng.gen::<u32>() as u64)
538 .collect();
539 let multiplicands_1: Vec<_> = (0..NUM_U32_ARITHMETIC_OPS)
540 .map(|_| rng.gen::<u32>() as u64)
541 .collect();
542 let addends: Vec<_> = (0..NUM_U32_ARITHMETIC_OPS)
543 .map(|_| rng.gen::<u32>() as u64)
544 .collect();
545
546 let gate = U32ArithmeticGate::<F, D> {
547 num_ops: NUM_U32_ARITHMETIC_OPS,
548 _phantom: PhantomData,
549 };
550
551 let vars = EvaluationVars {
552 local_constants: &[],
553 local_wires: &get_wires::<F, FF, D, NUM_U32_ARITHMETIC_OPS>(
554 multiplicands_0,
555 multiplicands_1,
556 addends,
557 ),
558 public_inputs_hash: &HashOut::rand(),
559 };
560
561 assert!(
562 gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()),
563 "Gate constraints are not satisfied."
564 );
565 }
566
567 #[test]
568 fn test_canonicity() {
569 const D: usize = 2;
570 type C = PoseidonGoldilocksConfig;
571 type F = <C as GenericConfig<D>>::F;
572 type FF = <C as GenericConfig<D>>::FE;
573 const NUM_U32_ARITHMETIC_OPS: usize = 3;
574
575 let multiplicands_0 = vec![0; NUM_U32_ARITHMETIC_OPS];
576 let multiplicands_1 = vec![0; NUM_U32_ARITHMETIC_OPS];
577 // A non-canonical addend will produce a non-canonical output using
578 // get_wires.
579 let addends = vec![0xFFFFFFFF00000001; NUM_U32_ARITHMETIC_OPS];
580
581 let gate = U32ArithmeticGate::<F, D> {
582 num_ops: NUM_U32_ARITHMETIC_OPS,
583 _phantom: PhantomData,
584 };
585
586 let vars = EvaluationVars {
587 local_constants: &[],
588 local_wires: &get_wires::<F, FF, D, NUM_U32_ARITHMETIC_OPS>(
589 multiplicands_0,
590 multiplicands_1,
591 addends,
592 ),
593 public_inputs_hash: &HashOut::rand(),
594 };
595
596 assert!(
597 !gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()),
598 "Non-canonical output should not pass constraints."
599 );
600 }
601}