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