plonky2-u32/src/gadgets/arithmetic_u32.rs
src/gadgets/arithmetic_u32.rsRaw
1use alloc::vec;
2use alloc::vec::Vec;
3use core::marker::PhantomData;
4
5use plonky2::field::extension::Extendable;
6use plonky2::hash::hash_types::RichField;
7use plonky2::iop::generator::{GeneratedValues, SimpleGenerator};
8use plonky2::iop::target::Target;
9use plonky2::iop::witness::{PartitionWitness, Witness};
10use plonky2::plonk::circuit_builder::CircuitBuilder;
11
12use crate::gates::add_many_u32::U32AddManyGate;
13use crate::gates::arithmetic_u32::U32ArithmeticGate;
14use crate::gates::subtraction_u32::U32SubtractionGate;
15use crate::witness::GeneratedValuesU32;
16
17#[derive(Clone, Copy, Debug)]
18pub struct U32Target(pub Target);
19
20pub trait CircuitBuilderU32<F: RichField + Extendable<D>, const D: usize> {
21 fn add_virtual_u32_target(&mut self) -> U32Target;
22
23 fn add_virtual_u32_targets(&mut self, n: usize) -> Vec<U32Target>;
24
25 /// Returns a U32Target for the value `c`, which is assumed to be at most 32 bits.
26 fn constant_u32(&mut self, c: u32) -> U32Target;
27
28 fn zero_u32(&mut self) -> U32Target;
29
30 fn one_u32(&mut self) -> U32Target;
31
32 fn connect_u32(&mut self, x: U32Target, y: U32Target);
33
34 fn assert_zero_u32(&mut self, x: U32Target);
35
36 /// Checks for special cases where the value of
37 /// `x * y + z`
38 /// can be determined without adding a `U32ArithmeticGate`.
39 fn arithmetic_u32_special_cases(
40 &mut self,
41 x: U32Target,
42 y: U32Target,
43 z: U32Target,
44 ) -> Option<(U32Target, U32Target)>;
45
46 // Returns x * y + z.
47 fn mul_add_u32(&mut self, x: U32Target, y: U32Target, z: U32Target) -> (U32Target, U32Target);
48
49 fn add_u32(&mut self, a: U32Target, b: U32Target) -> (U32Target, U32Target);
50
51 fn add_many_u32(&mut self, to_add: &[U32Target]) -> (U32Target, U32Target);
52
53 fn add_u32s_with_carry(
54 &mut self,
55 to_add: &[U32Target],
56 carry: U32Target,
57 ) -> (U32Target, U32Target);
58
59 fn mul_u32(&mut self, a: U32Target, b: U32Target) -> (U32Target, U32Target);
60
61 // Returns x - y - borrow, as a pair (result, borrow), where borrow is 0 or 1 depending on whether borrowing from the next digit is required (iff y + borrow > x).
62 fn sub_u32(&mut self, x: U32Target, y: U32Target, borrow: U32Target) -> (U32Target, U32Target);
63}
64
65impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilderU32<F, D>
66 for CircuitBuilder<F, D>
67{
68 fn add_virtual_u32_target(&mut self) -> U32Target {
69 U32Target(self.add_virtual_target())
70 }
71
72 fn add_virtual_u32_targets(&mut self, n: usize) -> Vec<U32Target> {
73 self.add_virtual_targets(n)
74 .into_iter()
75 .map(U32Target)
76 .collect()
77 }
78
79 /// Returns a U32Target for the value `c`, which is assumed to be at most 32 bits.
80 fn constant_u32(&mut self, c: u32) -> U32Target {
81 U32Target(self.constant(F::from_canonical_u32(c)))
82 }
83
84 fn zero_u32(&mut self) -> U32Target {
85 U32Target(self.zero())
86 }
87
88 fn one_u32(&mut self) -> U32Target {
89 U32Target(self.one())
90 }
91
92 fn connect_u32(&mut self, x: U32Target, y: U32Target) {
93 self.connect(x.0, y.0)
94 }
95
96 fn assert_zero_u32(&mut self, x: U32Target) {
97 self.assert_zero(x.0)
98 }
99
100 /// Checks for special cases where the value of
101 /// `x * y + z`
102 /// can be determined without adding a `U32ArithmeticGate`.
103 fn arithmetic_u32_special_cases(
104 &mut self,
105 x: U32Target,
106 y: U32Target,
107 z: U32Target,
108 ) -> Option<(U32Target, U32Target)> {
109 let x_const = self.target_as_constant(x.0);
110 let y_const = self.target_as_constant(y.0);
111 let z_const = self.target_as_constant(z.0);
112
113 // If both terms are constant, return their (constant) sum.
114 let first_term_const = if let (Some(xx), Some(yy)) = (x_const, y_const) {
115 Some(xx * yy)
116 } else {
117 None
118 };
119
120 if let (Some(a), Some(b)) = (first_term_const, z_const) {
121 let sum = (a + b).to_canonical_u64();
122 let (low, high) = (sum as u32, (sum >> 32) as u32);
123 return Some((self.constant_u32(low), self.constant_u32(high)));
124 }
125
126 None
127 }
128
129 // Returns x * y + z.
130 fn mul_add_u32(&mut self, x: U32Target, y: U32Target, z: U32Target) -> (U32Target, U32Target) {
131 if let Some(result) = self.arithmetic_u32_special_cases(x, y, z) {
132 return result;
133 }
134
135 let gate = U32ArithmeticGate::<F, D>::new_from_config(&self.config);
136 let (row, copy) = self.find_slot(gate, &[], &[]);
137
138 self.connect(Target::wire(row, gate.wire_ith_multiplicand_0(copy)), x.0);
139 self.connect(Target::wire(row, gate.wire_ith_multiplicand_1(copy)), y.0);
140 self.connect(Target::wire(row, gate.wire_ith_addend(copy)), z.0);
141
142 let output_low = U32Target(Target::wire(row, gate.wire_ith_output_low_half(copy)));
143 let output_high = U32Target(Target::wire(row, gate.wire_ith_output_high_half(copy)));
144
145 (output_low, output_high)
146 }
147
148 fn add_u32(&mut self, a: U32Target, b: U32Target) -> (U32Target, U32Target) {
149 let one = self.one_u32();
150 self.mul_add_u32(a, one, b)
151 }
152
153 fn add_many_u32(&mut self, to_add: &[U32Target]) -> (U32Target, U32Target) {
154 match to_add.len() {
155 0 => (self.zero_u32(), self.zero_u32()),
156 1 => (to_add[0], self.zero_u32()),
157 2 => self.add_u32(to_add[0], to_add[1]),
158 _ => {
159 let num_addends = to_add.len();
160 let gate = U32AddManyGate::<F, D>::new_from_config(&self.config, num_addends);
161 let (row, copy) =
162 self.find_slot(gate, &[F::from_canonical_usize(num_addends)], &[]);
163
164 for j in 0..num_addends {
165 self.connect(
166 Target::wire(row, gate.wire_ith_op_jth_addend(copy, j)),
167 to_add[j].0,
168 );
169 }
170 let zero = self.zero();
171 self.connect(Target::wire(row, gate.wire_ith_carry(copy)), zero);
172
173 let output_low = U32Target(Target::wire(row, gate.wire_ith_output_result(copy)));
174 let output_high = U32Target(Target::wire(row, gate.wire_ith_output_carry(copy)));
175
176 (output_low, output_high)
177 }
178 }
179 }
180
181 fn add_u32s_with_carry(
182 &mut self,
183 to_add: &[U32Target],
184 carry: U32Target,
185 ) -> (U32Target, U32Target) {
186 if to_add.len() == 1 {
187 return self.add_u32(to_add[0], carry);
188 }
189
190 let num_addends = to_add.len();
191
192 let gate = U32AddManyGate::<F, D>::new_from_config(&self.config, num_addends);
193 let (row, copy) = self.find_slot(gate, &[F::from_canonical_usize(num_addends)], &[]);
194
195 for j in 0..num_addends {
196 self.connect(
197 Target::wire(row, gate.wire_ith_op_jth_addend(copy, j)),
198 to_add[j].0,
199 );
200 }
201 self.connect(Target::wire(row, gate.wire_ith_carry(copy)), carry.0);
202
203 let output = U32Target(Target::wire(row, gate.wire_ith_output_result(copy)));
204 let output_carry = U32Target(Target::wire(row, gate.wire_ith_output_carry(copy)));
205
206 (output, output_carry)
207 }
208
209 fn mul_u32(&mut self, a: U32Target, b: U32Target) -> (U32Target, U32Target) {
210 let zero = self.zero_u32();
211 self.mul_add_u32(a, b, zero)
212 }
213
214 // Returns x - y - borrow, as a pair (result, borrow), where borrow is 0 or 1 depending on whether borrowing from the next digit is required (iff y + borrow > x).
215 fn sub_u32(&mut self, x: U32Target, y: U32Target, borrow: U32Target) -> (U32Target, U32Target) {
216 let gate = U32SubtractionGate::<F, D>::new_from_config(&self.config);
217 let (row, copy) = self.find_slot(gate, &[], &[]);
218
219 self.connect(Target::wire(row, gate.wire_ith_input_x(copy)), x.0);
220 self.connect(Target::wire(row, gate.wire_ith_input_y(copy)), y.0);
221 self.connect(
222 Target::wire(row, gate.wire_ith_input_borrow(copy)),
223 borrow.0,
224 );
225
226 let output_result = U32Target(Target::wire(row, gate.wire_ith_output_result(copy)));
227 let output_borrow = U32Target(Target::wire(row, gate.wire_ith_output_borrow(copy)));
228
229 (output_result, output_borrow)
230 }
231}
232
233#[derive(Debug)]
234struct SplitToU32Generator<F: RichField + Extendable<D>, const D: usize> {
235 x: Target,
236 low: U32Target,
237 high: U32Target,
238 _phantom: PhantomData<F>,
239}
240
241impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
242 for SplitToU32Generator<F, D>
243{
244 fn dependencies(&self) -> Vec<Target> {
245 vec![self.x]
246 }
247
248 fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
249 let x = witness.get_target(self.x);
250 let x_u64 = x.to_canonical_u64();
251 let low = x_u64 as u32;
252 let high = (x_u64 >> 32) as u32;
253
254 out_buffer.set_u32_target(self.low, low);
255 out_buffer.set_u32_target(self.high, high);
256 }
257}
258
259#[cfg(test)]
260mod tests {
261 use anyhow::Result;
262 use plonky2::iop::witness::PartialWitness;
263 use plonky2::plonk::circuit_data::CircuitConfig;
264 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
265 use rand::rngs::OsRng;
266 use rand::Rng;
267
268 use super::*;
269
270 #[test]
271 pub fn test_add_many_u32s() -> Result<()> {
272 const D: usize = 2;
273 type C = PoseidonGoldilocksConfig;
274 type F = <C as GenericConfig<D>>::F;
275
276 const NUM_ADDENDS: usize = 15;
277
278 let config = CircuitConfig::standard_recursion_config();
279
280 let pw = PartialWitness::new();
281 let mut builder = CircuitBuilder::<F, D>::new(config);
282
283 let mut rng = OsRng;
284 let mut to_add = Vec::new();
285 let mut sum = 0u64;
286 for _ in 0..NUM_ADDENDS {
287 let x: u32 = rng.gen();
288 sum += x as u64;
289 to_add.push(builder.constant_u32(x));
290 }
291 let carry = builder.zero_u32();
292 let (result_low, result_high) = builder.add_u32s_with_carry(&to_add, carry);
293 let expected_low = builder.constant_u32((sum % (1 << 32)) as u32);
294 let expected_high = builder.constant_u32((sum >> 32) as u32);
295
296 builder.connect_u32(result_low, expected_low);
297 builder.connect_u32(result_high, expected_high);
298
299 let data = builder.build::<C>();
300 let proof = data.prove(pw).unwrap();
301 data.verify(proof)
302 }
303}