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