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