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