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