arnaucube
arnaucube/plonky2-ecdsaPublic Self-hosted repository
plonky2-ecdsa/src/gadgets/biguint.rs
src/gadgets/biguint.rsRaw
  1use alloc::vec;
  2use alloc::vec::Vec;
  3use core::marker::PhantomData;
  4
  5use num::{BigUint, Integer, Zero};
  6use plonky2::field::extension::Extendable;
  7use plonky2::field::types::{PrimeField, PrimeField64};
  8use plonky2::hash::hash_types::RichField;
  9use plonky2::iop::generator::{GeneratedValues, SimpleGenerator};
 10use plonky2::iop::target::{BoolTarget, Target};
 11use plonky2::iop::witness::{PartitionWitness, Witness};
 12use plonky2::plonk::circuit_builder::CircuitBuilder;
 13use plonky2_u32::gadgets::arithmetic_u32::{CircuitBuilderU32, U32Target};
 14use plonky2_u32::gadgets::multiple_comparison::list_le_u32_circuit;
 15use plonky2_u32::witness::{GeneratedValuesU32, WitnessU32};
 16
 17#[derive(Clone, Debug)]
 18pub struct BigUintTarget {
 19    pub limbs: Vec<U32Target>,
 20}
 21
 22impl BigUintTarget {
 23    pub fn num_limbs(&self) -> usize {
 24        self.limbs.len()
 25    }
 26
 27    pub fn get_limb(&self, i: usize) -> U32Target {
 28        self.limbs[i]
 29    }
 30}
 31
 32pub trait CircuitBuilderBiguint<F: RichField + Extendable<D>, const D: usize> {
 33    fn constant_biguint(&mut self, value: &BigUint) -> BigUintTarget;
 34
 35    fn zero_biguint(&mut self) -> BigUintTarget;
 36
 37    fn connect_biguint(&mut self, lhs: &BigUintTarget, rhs: &BigUintTarget);
 38
 39    fn pad_biguints(
 40        &mut self,
 41        a: &BigUintTarget,
 42        b: &BigUintTarget,
 43    ) -> (BigUintTarget, BigUintTarget);
 44
 45    fn cmp_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BoolTarget;
 46
 47    fn add_virtual_biguint_target(&mut self, num_limbs: usize) -> BigUintTarget;
 48
 49    /// Add two `BigUintTarget`s.
 50    fn add_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget;
 51
 52    /// Subtract two `BigUintTarget`s. We assume that the first is larger than the second.
 53    fn sub_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget;
 54
 55    fn mul_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget;
 56
 57    fn mul_biguint_by_bool(&mut self, a: &BigUintTarget, b: BoolTarget) -> BigUintTarget;
 58
 59    /// Returns x * y + z. This is no more efficient than mul-then-add; it's purely for convenience (only need to call one CircuitBuilder function).
 60    fn mul_add_biguint(
 61        &mut self,
 62        x: &BigUintTarget,
 63        y: &BigUintTarget,
 64        z: &BigUintTarget,
 65    ) -> BigUintTarget;
 66
 67    fn div_rem_biguint(
 68        &mut self,
 69        a: &BigUintTarget,
 70        b: &BigUintTarget,
 71    ) -> (BigUintTarget, BigUintTarget);
 72
 73    fn div_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget;
 74
 75    fn rem_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget;
 76}
 77
 78impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilderBiguint<F, D>
 79    for CircuitBuilder<F, D>
 80{
 81    fn constant_biguint(&mut self, value: &BigUint) -> BigUintTarget {
 82        let limb_values = value.to_u32_digits();
 83        let limbs = limb_values.iter().map(|&l| self.constant_u32(l)).collect();
 84
 85        BigUintTarget { limbs }
 86    }
 87
 88    fn zero_biguint(&mut self) -> BigUintTarget {
 89        self.constant_biguint(&BigUint::zero())
 90    }
 91
 92    fn connect_biguint(&mut self, lhs: &BigUintTarget, rhs: &BigUintTarget) {
 93        let min_limbs = lhs.num_limbs().min(rhs.num_limbs());
 94        for i in 0..min_limbs {
 95            self.connect_u32(lhs.get_limb(i), rhs.get_limb(i));
 96        }
 97
 98        for i in min_limbs..lhs.num_limbs() {
 99            self.assert_zero_u32(lhs.get_limb(i));
100        }
101        for i in min_limbs..rhs.num_limbs() {
102            self.assert_zero_u32(rhs.get_limb(i));
103        }
104    }
105
106    fn pad_biguints(
107        &mut self,
108        a: &BigUintTarget,
109        b: &BigUintTarget,
110    ) -> (BigUintTarget, BigUintTarget) {
111        if a.num_limbs() > b.num_limbs() {
112            let mut padded_b = b.clone();
113            for _ in b.num_limbs()..a.num_limbs() {
114                padded_b.limbs.push(self.zero_u32());
115            }
116
117            (a.clone(), padded_b)
118        } else {
119            let mut padded_a = a.clone();
120            for _ in a.num_limbs()..b.num_limbs() {
121                padded_a.limbs.push(self.zero_u32());
122            }
123
124            (padded_a, b.clone())
125        }
126    }
127
128    fn cmp_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BoolTarget {
129        let (a, b) = self.pad_biguints(a, b);
130
131        list_le_u32_circuit(self, a.limbs, b.limbs)
132    }
133
134    fn add_virtual_biguint_target(&mut self, num_limbs: usize) -> BigUintTarget {
135        let limbs = self.add_virtual_u32_targets(num_limbs);
136
137        BigUintTarget { limbs }
138    }
139
140    fn add_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget {
141        let num_limbs = a.num_limbs().max(b.num_limbs());
142
143        let mut combined_limbs = vec![];
144        let mut carry = self.zero_u32();
145        for i in 0..num_limbs {
146            let a_limb = (i < a.num_limbs())
147                .then(|| a.limbs[i])
148                .unwrap_or_else(|| self.zero_u32());
149            let b_limb = (i < b.num_limbs())
150                .then(|| b.limbs[i])
151                .unwrap_or_else(|| self.zero_u32());
152
153            let (new_limb, new_carry) = self.add_many_u32(&[carry, a_limb, b_limb]);
154            carry = new_carry;
155            combined_limbs.push(new_limb);
156        }
157        combined_limbs.push(carry);
158
159        BigUintTarget {
160            limbs: combined_limbs,
161        }
162    }
163
164    fn sub_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget {
165        let (a, b) = self.pad_biguints(a, b);
166        let num_limbs = a.limbs.len();
167
168        let mut result_limbs = vec![];
169
170        let mut borrow = self.zero_u32();
171        for i in 0..num_limbs {
172            let (result, new_borrow) = self.sub_u32(a.limbs[i], b.limbs[i], borrow);
173            result_limbs.push(result);
174            borrow = new_borrow;
175        }
176        // Borrow should be zero here.
177
178        BigUintTarget {
179            limbs: result_limbs,
180        }
181    }
182
183    fn mul_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget {
184        let total_limbs = a.limbs.len() + b.limbs.len();
185
186        let mut to_add = vec![vec![]; total_limbs];
187        for i in 0..a.limbs.len() {
188            for j in 0..b.limbs.len() {
189                let (product, carry) = self.mul_u32(a.limbs[i], b.limbs[j]);
190                to_add[i + j].push(product);
191                to_add[i + j + 1].push(carry);
192            }
193        }
194
195        let mut combined_limbs = vec![];
196        let mut carry = self.zero_u32();
197        for summands in &mut to_add {
198            let (new_result, new_carry) = self.add_u32s_with_carry(summands, carry);
199            combined_limbs.push(new_result);
200            carry = new_carry;
201        }
202        combined_limbs.push(carry);
203
204        BigUintTarget {
205            limbs: combined_limbs,
206        }
207    }
208
209    fn mul_biguint_by_bool(&mut self, a: &BigUintTarget, b: BoolTarget) -> BigUintTarget {
210        let t = b.target;
211
212        BigUintTarget {
213            limbs: a
214                .limbs
215                .iter()
216                .map(|&l| U32Target(self.mul(l.0, t)))
217                .collect(),
218        }
219    }
220
221    fn mul_add_biguint(
222        &mut self,
223        x: &BigUintTarget,
224        y: &BigUintTarget,
225        z: &BigUintTarget,
226    ) -> BigUintTarget {
227        let prod = self.mul_biguint(x, y);
228        self.add_biguint(&prod, z)
229    }
230
231    fn div_rem_biguint(
232        &mut self,
233        a: &BigUintTarget,
234        b: &BigUintTarget,
235    ) -> (BigUintTarget, BigUintTarget) {
236        let a_len = a.limbs.len();
237        let b_len = b.limbs.len();
238        let div_num_limbs = if b_len > a_len + 1 {
239            0
240        } else {
241            a_len - b_len + 1
242        };
243        let div = self.add_virtual_biguint_target(div_num_limbs);
244        let rem = self.add_virtual_biguint_target(b_len);
245
246        self.add_simple_generator(BigUintDivRemGenerator::<F, D> {
247            a: a.clone(),
248            b: b.clone(),
249            div: div.clone(),
250            rem: rem.clone(),
251            _phantom: PhantomData,
252        });
253
254        let div_b = self.mul_biguint(&div, b);
255        let div_b_plus_rem = self.add_biguint(&div_b, &rem);
256        self.connect_biguint(a, &div_b_plus_rem);
257
258        let cmp_rem_b = self.cmp_biguint(&rem, b);
259        self.assert_one(cmp_rem_b.target);
260
261        (div, rem)
262    }
263
264    fn div_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget {
265        let (div, _rem) = self.div_rem_biguint(a, b);
266        div
267    }
268
269    fn rem_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget {
270        let (_div, rem) = self.div_rem_biguint(a, b);
271        rem
272    }
273}
274
275pub trait WitnessBigUint<F: PrimeField64>: Witness<F> {
276    fn get_biguint_target(&self, target: BigUintTarget) -> BigUint;
277    fn set_biguint_target(&mut self, target: &BigUintTarget, value: &BigUint);
278}
279
280impl<T: Witness<F>, F: PrimeField64> WitnessBigUint<F> for T {
281    fn get_biguint_target(&self, target: BigUintTarget) -> BigUint {
282        target
283            .limbs
284            .into_iter()
285            .rev()
286            .fold(BigUint::zero(), |acc, limb| {
287                (acc << 32) + self.get_target(limb.0).to_canonical_biguint()
288            })
289    }
290
291    fn set_biguint_target(&mut self, target: &BigUintTarget, value: &BigUint) {
292        let mut limbs = value.to_u32_digits();
293        assert!(target.num_limbs() >= limbs.len());
294        limbs.resize(target.num_limbs(), 0);
295        for i in 0..target.num_limbs() {
296            self.set_u32_target(target.limbs[i], limbs[i]);
297        }
298    }
299}
300
301pub trait GeneratedValuesBigUint<F: PrimeField> {
302    fn set_biguint_target(&mut self, target: &BigUintTarget, value: &BigUint);
303}
304
305impl<F: PrimeField> GeneratedValuesBigUint<F> for GeneratedValues<F> {
306    fn set_biguint_target(&mut self, target: &BigUintTarget, value: &BigUint) {
307        let mut limbs = value.to_u32_digits();
308        assert!(target.num_limbs() >= limbs.len());
309        limbs.resize(target.num_limbs(), 0);
310        for i in 0..target.num_limbs() {
311            self.set_u32_target(target.get_limb(i), limbs[i]);
312        }
313    }
314}
315
316#[derive(Debug)]
317struct BigUintDivRemGenerator<F: RichField + Extendable<D>, const D: usize> {
318    a: BigUintTarget,
319    b: BigUintTarget,
320    div: BigUintTarget,
321    rem: BigUintTarget,
322    _phantom: PhantomData<F>,
323}
324
325impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F>
326    for BigUintDivRemGenerator<F, D>
327{
328    fn dependencies(&self) -> Vec<Target> {
329        self.a
330            .limbs
331            .iter()
332            .chain(&self.b.limbs)
333            .map(|&l| l.0)
334            .collect()
335    }
336
337    fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
338        let a = witness.get_biguint_target(self.a.clone());
339        let b = witness.get_biguint_target(self.b.clone());
340        let (div, rem) = a.div_rem(&b);
341
342        out_buffer.set_biguint_target(&self.div, &div);
343        out_buffer.set_biguint_target(&self.rem, &rem);
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use anyhow::Result;
350    use num::{BigUint, FromPrimitive, Integer};
351    use plonky2::iop::witness::PartialWitness;
352    use plonky2::plonk::circuit_builder::CircuitBuilder;
353    use plonky2::plonk::circuit_data::CircuitConfig;
354    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
355    use rand::rngs::OsRng;
356    use rand::Rng;
357
358    use crate::gadgets::biguint::{CircuitBuilderBiguint, WitnessBigUint};
359
360    #[test]
361    fn test_biguint_add() -> Result<()> {
362        const D: usize = 2;
363        type C = PoseidonGoldilocksConfig;
364        type F = <C as GenericConfig<D>>::F;
365        let mut rng = OsRng;
366
367        let x_value = BigUint::from_u128(rng.gen()).unwrap();
368        let y_value = BigUint::from_u128(rng.gen()).unwrap();
369        let expected_z_value = &x_value + &y_value;
370
371        let config = CircuitConfig::standard_recursion_config();
372        let mut pw = PartialWitness::new();
373        let mut builder = CircuitBuilder::<F, D>::new(config);
374
375        let x = builder.add_virtual_biguint_target(x_value.to_u32_digits().len());
376        let y = builder.add_virtual_biguint_target(y_value.to_u32_digits().len());
377        let z = builder.add_biguint(&x, &y);
378        let expected_z = builder.add_virtual_biguint_target(expected_z_value.to_u32_digits().len());
379        builder.connect_biguint(&z, &expected_z);
380
381        pw.set_biguint_target(&x, &x_value);
382        pw.set_biguint_target(&y, &y_value);
383        pw.set_biguint_target(&expected_z, &expected_z_value);
384
385        let data = builder.build::<C>();
386        let proof = data.prove(pw).unwrap();
387        data.verify(proof)
388    }
389
390    #[test]
391    fn test_biguint_sub() -> Result<()> {
392        const D: usize = 2;
393        type C = PoseidonGoldilocksConfig;
394        type F = <C as GenericConfig<D>>::F;
395        let mut rng = OsRng;
396
397        let mut x_value = BigUint::from_u128(rng.gen()).unwrap();
398        let mut y_value = BigUint::from_u128(rng.gen()).unwrap();
399        if y_value > x_value {
400            (x_value, y_value) = (y_value, x_value);
401        }
402        let expected_z_value = &x_value - &y_value;
403
404        let config = CircuitConfig::standard_recursion_config();
405        let pw = PartialWitness::new();
406        let mut builder = CircuitBuilder::<F, D>::new(config);
407
408        let x = builder.constant_biguint(&x_value);
409        let y = builder.constant_biguint(&y_value);
410        let z = builder.sub_biguint(&x, &y);
411        let expected_z = builder.constant_biguint(&expected_z_value);
412
413        builder.connect_biguint(&z, &expected_z);
414
415        let data = builder.build::<C>();
416        let proof = data.prove(pw).unwrap();
417        data.verify(proof)
418    }
419
420    #[test]
421    fn test_biguint_mul() -> Result<()> {
422        const D: usize = 2;
423        type C = PoseidonGoldilocksConfig;
424        type F = <C as GenericConfig<D>>::F;
425        let mut rng = OsRng;
426
427        let x_value = BigUint::from_u128(rng.gen()).unwrap();
428        let y_value = BigUint::from_u128(rng.gen()).unwrap();
429        let expected_z_value = &x_value * &y_value;
430
431        let config = CircuitConfig::standard_recursion_config();
432        let mut pw = PartialWitness::new();
433        let mut builder = CircuitBuilder::<F, D>::new(config);
434
435        let x = builder.add_virtual_biguint_target(x_value.to_u32_digits().len());
436        let y = builder.add_virtual_biguint_target(y_value.to_u32_digits().len());
437        let z = builder.mul_biguint(&x, &y);
438        let expected_z = builder.add_virtual_biguint_target(expected_z_value.to_u32_digits().len());
439        builder.connect_biguint(&z, &expected_z);
440
441        pw.set_biguint_target(&x, &x_value);
442        pw.set_biguint_target(&y, &y_value);
443        pw.set_biguint_target(&expected_z, &expected_z_value);
444
445        let data = builder.build::<C>();
446        let proof = data.prove(pw).unwrap();
447        data.verify(proof)
448    }
449
450    #[test]
451    fn test_biguint_cmp() -> Result<()> {
452        const D: usize = 2;
453        type C = PoseidonGoldilocksConfig;
454        type F = <C as GenericConfig<D>>::F;
455        let mut rng = OsRng;
456
457        let x_value = BigUint::from_u128(rng.gen()).unwrap();
458        let y_value = BigUint::from_u128(rng.gen()).unwrap();
459
460        let config = CircuitConfig::standard_recursion_config();
461        let pw = PartialWitness::new();
462        let mut builder = CircuitBuilder::<F, D>::new(config);
463
464        let x = builder.constant_biguint(&x_value);
465        let y = builder.constant_biguint(&y_value);
466        let cmp = builder.cmp_biguint(&x, &y);
467        let expected_cmp = builder.constant_bool(x_value <= y_value);
468
469        builder.connect(cmp.target, expected_cmp.target);
470
471        let data = builder.build::<C>();
472        let proof = data.prove(pw).unwrap();
473        data.verify(proof)
474    }
475
476    #[test]
477    fn test_biguint_div_rem() -> Result<()> {
478        const D: usize = 2;
479        type C = PoseidonGoldilocksConfig;
480        type F = <C as GenericConfig<D>>::F;
481        let mut rng = OsRng;
482
483        let mut x_value = BigUint::from_u128(rng.gen()).unwrap();
484        let mut y_value = BigUint::from_u128(rng.gen()).unwrap();
485        if y_value > x_value {
486            (x_value, y_value) = (y_value, x_value);
487        }
488        let (expected_div_value, expected_rem_value) = x_value.div_rem(&y_value);
489
490        let config = CircuitConfig::standard_recursion_config();
491        let pw = PartialWitness::new();
492        let mut builder = CircuitBuilder::<F, D>::new(config);
493
494        let x = builder.constant_biguint(&x_value);
495        let y = builder.constant_biguint(&y_value);
496        let (div, rem) = builder.div_rem_biguint(&x, &y);
497
498        let expected_div = builder.constant_biguint(&expected_div_value);
499        let expected_rem = builder.constant_biguint(&expected_rem_value);
500
501        builder.connect_biguint(&div, &expected_div);
502        builder.connect_biguint(&rem, &expected_rem);
503
504        let data = builder.build::<C>();
505        let proof = data.prove(pw).unwrap();
506        data.verify(proof)
507    }
508}