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