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