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