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