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