Add bits to biguint struct

ax0 · ax0 · commit 1f8230b7227d · 2025-06-09T10:56:53.000+10:00
diff --git a/src/backends/plonky2/primitives/ec/bits.rs b/src/backends/plonky2/primitives/ec/bits.rs
@@ -14,7 +14,7 @@ use plonky2::{
         witness::{PartitionWitness, Witness, WitnessWrite},
     },
     plonk::{circuit_builder::CircuitBuilder, circuit_data::CommonCircuitData},
-    util::serialization::{IoResult, Read, Write},
+    util::serialization::{Buffer, IoResult, Read, Write},
 };
 
 use crate::backends::plonky2::basetypes::{D, F};
@@ -79,7 +79,10 @@ impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D>
 /// A big integer, represented in base `2^32` with 10 digits, in little endian
 /// form.
 #[derive(Clone, Debug)]
-pub struct BigUInt320Target(pub(super) [Target; 10]);
+pub struct BigUInt320Target {
+    pub limbs: [Target; 10],
+    pub bits: [BoolTarget; 320],
+}
 
 pub trait CircuitBuilderBits {
     /// Enforces the constraint that `then_zero` must be zero if `if_zero`
@@ -90,27 +93,26 @@ pub trait CircuitBuilderBits {
     /// are zero, then it chooses the solution `x = 0`.
     fn conditional_zero(&mut self, if_zero: Target, then_zero: Target);
 
-    /// Returns the binary representation of the target, in little-endian order.
-    fn biguint_bits(&mut self, x: &BigUInt320Target) -> [BoolTarget; 320];
-
     /// Decomposes the target x as `y + 2^32 z`, where `0 < y,z < 2**32`, and
     /// `y=0` if `z=2**32-1`.  Note that calling [`CircuitBuilder::split_le`]
     /// with `num_bits = 64` will not check the latter condition.
     fn split_32_bit(&mut self, x: Target) -> [Target; 2];
 
+    /// Like `split_low_high` except it doesn't discard the bit decompositions.
+    fn split_low_high_with_bits(
+        &mut self,
+        x: Target,
+        n_log: usize,
+        num_bits: usize,
+    ) -> ((Target, Vec<BoolTarget>), (Target, Vec<BoolTarget>));
+
     /// Interprets `arr` as an integer in base `[GoldilocksField::ORDER]`,
     /// with the digits in little endian order.  The length of `arr` must be at
     /// most 5.
     fn field_elements_to_biguint(&mut self, arr: &[Target]) -> BigUInt320Target;
 
-    fn normalize_bigint(
-        &mut self,
-        x: &BigUInt320Target,
-        max_digit_bits: usize,
-        max_num_carries: usize,
-    ) -> BigUInt320Target;
-
     fn constant_biguint320(&mut self, n: &BigUint) -> BigUInt320Target;
+    fn biguint320_target_from_limbs(&mut self, x: &[Target]) -> BigUInt320Target;
     fn add_virtual_biguint320_target(&mut self) -> BigUInt320Target;
     fn connect_biguint320(&mut self, x: &BigUInt320Target, y: &BigUInt320Target);
 }
@@ -128,20 +130,18 @@ impl CircuitBuilderBits for CircuitBuilder<GoldilocksField, 2> {
         self.connect(prod, then_zero);
     }
 
-    fn biguint_bits(&mut self, x: &BigUInt320Target) -> [BoolTarget; 320] {
-        let bits = x.0.map(|t| self.low_bits(t, 32, 32));
-        array::from_fn(|i| bits[i / 32][i % 32])
-    }
-
     fn field_elements_to_biguint(&mut self, arr: &[Target]) -> BigUInt320Target {
         assert!(arr.len() <= 5);
         let zero = self.zero();
         let neg_one = self.neg_one();
         let two_32 = self.constant(GoldilocksField::from_canonical_u64(1 << 32));
         // Apply Horner's method to Σarr[i]*p^i.
         // First map each target to its limbs.
-        let arr_limbs: Vec<_> = arr.iter().map(|x| self.split_32_bit(*x).to_vec()).collect();
-        let res_limbs = arr_limbs
+        let arr_limbs: Vec<_> = arr
+            .iter()
+            .map(|x| (self.split_32_bit(*x).to_vec(), vec![]))
+            .collect();
+        let (res_limbs, res_bits) = arr_limbs
             .into_iter()
             .rev()
             .enumerate()
@@ -153,25 +153,25 @@ impl CircuitBuilderBits for CircuitBuilder<GoldilocksField, 2> {
                     .map(|j| {
                         if j == 0 {
                             // x_0
-                            res[0]
+                            res.0[0]
                         } else if j == 1 {
                             // x_1 - x_0 + 2^32
-                            let diff = self.sub(res[1], res[0]);
+                            let diff = self.sub(res.0[1], res.0[0]);
                             self.add(diff, two_32)
                         } else if j < 2 * i {
                             // x_j + x_{j-2} - x_{j-1} + 2^32 - 1
-                            let diff = self.sub(res[j], res[j - 1]);
-                            let sum = self.add(diff, res[j - 2]);
+                            let diff = self.sub(res.0[j], res.0[j - 1]);
+                            let sum = self.add(diff, res.0[j - 2]);
                             let sum = self.add(sum, two_32);
                             self.add(sum, neg_one)
                         } else if j == 2 * i {
                             // x_{2*j - 2} - x_{2*j - 1} + 2^32
-                            let diff = self.sub(res[2 * i - 2], res[2 * i - 1]);
+                            let diff = self.sub(res.0[2 * i - 2], res.0[2 * i - 1]);
                             let sum = self.add(diff, two_32);
                             self.add(sum, neg_one)
                         } else {
                             // x_{2*i - 1} - 1
-                            self.add(res[2 * i - 1], neg_one)
+                            self.add(res.0[2 * i - 1], neg_one)
                         }
                     })
                     .collect::<Vec<_>>();
@@ -180,8 +180,8 @@ impl CircuitBuilderBits for CircuitBuilder<GoldilocksField, 2> {
                     .into_iter()
                     .enumerate()
                     .map(|(i, x)| match i {
-                        0 => self.add(a[0], x),
-                        1 => self.add(a[1], x),
+                        0 => self.add(a.0[0], x),
+                        1 => self.add(a.0[1], x),
                         _ => x,
                     })
                     .collect::<Vec<_>>();
@@ -192,30 +192,24 @@ impl CircuitBuilderBits for CircuitBuilder<GoldilocksField, 2> {
                 )
             })
             .map(|(_, v)| v)
-            .unwrap_or(vec![]);
+            .unwrap_or((vec![], vec![]));
         // Collect limbs, padding with 0s if necessary.
-        BigUInt320Target(array::from_fn(|i| {
+        let limbs: [Target; 10] = array::from_fn(|i| {
             if i < res_limbs.len() {
                 res_limbs[i]
             } else {
                 zero
             }
-        }))
-    }
-
-    fn normalize_bigint(
-        &mut self,
-        x: &BigUInt320Target,
-        max_digit_bits: usize,
-        max_num_carries: usize,
-    ) -> BigUInt320Target {
-        let mut x = x.clone();
-        for i in 0..max_num_carries {
-            let (low, high) = self.split_low_high(x.0[i], 32, max_digit_bits);
-            x.0[i] = low;
-            x.0[i + 1] = self.add(x.0[i + 1], high);
-        }
-        x
+        });
+        // Collect bits, padding with 0s if necessary.
+        let bits: [BoolTarget; 320] = array::from_fn(|i| {
+            if i < res_bits.len() {
+                res_bits[i]
+            } else {
+                self._false()
+            }
+        });
+        BigUInt320Target { limbs, bits }
     }
 
     fn split_32_bit(&mut self, x: Target) -> [Target; 2] {
@@ -226,44 +220,146 @@ impl CircuitBuilderBits for CircuitBuilder<GoldilocksField, 2> {
         [low, high]
     }
 
+    fn split_low_high_with_bits(
+        &mut self,
+        x: Target,
+        n_log: usize,
+        num_bits: usize,
+    ) -> ((Target, Vec<BoolTarget>), (Target, Vec<BoolTarget>)) {
+        let low = self.add_virtual_target();
+        let high = self.add_virtual_target();
+
+        self.add_simple_generator(LowHighGenerator {
+            integer: x,
+            n_log,
+            low,
+            high,
+        });
+
+        let low_bits = self.split_le(low, n_log);
+        let high_bits = self.split_le(high, num_bits - n_log);
+
+        let pow2 = self.constant(F::from_canonical_u64(1 << n_log));
+        let comp_x = self.mul_add(high, pow2, low);
+        self.connect(x, comp_x);
+
+        ((low, low_bits), (high, high_bits))
+    }
+
     fn constant_biguint320(&mut self, n: &BigUint) -> BigUInt320Target {
         assert!(n.bits() <= 320);
         let digits = n.to_u32_digits();
-        let targets = array::from_fn(|i| {
+        let limbs: [Target; 10] = array::from_fn(|i| {
             let d = digits.get(i).copied().unwrap_or(0);
             self.constant(GoldilocksField::from_canonical_u32(d))
         });
-        BigUInt320Target(targets)
+        self.biguint320_target_from_limbs(&limbs)
     }
 
-    fn add_virtual_biguint320_target(&mut self) -> BigUInt320Target {
-        let targets = self.add_virtual_target_arr();
-        for t in targets {
-            self.range_check(t, 32);
+    fn biguint320_target_from_limbs(&mut self, x: &[Target]) -> BigUInt320Target {
+        assert!(x.len() == 10);
+        let limbs = array::from_fn(|i| x[i]);
+        let bit_vec = biguint_limbs_to_bits(self, x);
+        BigUInt320Target {
+            limbs,
+            bits: array::from_fn(|i| bit_vec[i]),
         }
-        BigUInt320Target(targets)
+    }
+
+    fn add_virtual_biguint320_target(&mut self) -> BigUInt320Target {
+        let limbs: [Target; 10] = self.add_virtual_target_arr();
+        self.biguint320_target_from_limbs(&limbs)
     }
 
     fn connect_biguint320(&mut self, x: &BigUInt320Target, y: &BigUInt320Target) {
         for i in 0..10 {
-            self.connect(x.0[i], y.0[i]);
+            self.connect(x.limbs[i], y.limbs[i]);
         }
     }
 }
 
 /// Normalises the limbs of a biguint assuming no overflow in the
-/// field.
+/// field. Returns the limbs together with their bit decomposition.
 fn normalize_biguint_limbs(
     builder: &mut CircuitBuilder<F, D>,
     x: &[Target],
     max_digit_bits: usize,
     max_num_carries: usize,
-) -> Vec<Target> {
+) -> (Vec<Target>, Vec<BoolTarget>) {
     let mut x = x.to_vec();
+    let mut bits = Vec::with_capacity(32 * (max_num_carries + 1));
     for i in 0..max_num_carries {
-        let (low, high) = builder.split_low_high(x[i], 32, max_digit_bits);
+        let ((low, mut low_bits), (high, _)) =
+            builder.split_low_high_with_bits(x[i], 32, max_digit_bits);
         x[i] = low;
         x[i + 1] = builder.add(x[i + 1], high);
+        bits.append(&mut low_bits);
+    }
+    let mut final_bits = builder.split_le(x[max_num_carries], 32);
+    bits.append(&mut final_bits);
+    (x, bits)
+}
+
+/// Converts biguint limbs to bits, checking that each limb is 32-bits
+/// long.
+fn biguint_limbs_to_bits(builder: &mut CircuitBuilder<F, D>, limbs: &[Target]) -> Vec<BoolTarget> {
+    limbs
+        .iter()
+        .flat_map(|t| builder.split_le(*t, 32))
+        .collect()
+}
+
+/*
+Copied from https://github.com/0xPolygonZero/plonky2/blob/82791c4809d6275682c34b926390ecdbdc2a5297/plonky2/src/gadgets/range_check.rs#L62
+ */
+
+#[derive(Debug, Default)]
+pub struct LowHighGenerator {
+    integer: Target,
+    n_log: usize,
+    low: Target,
+    high: Target,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D> for LowHighGenerator {
+    fn id(&self) -> String {
+        "LowHighGenerator".to_string()
+    }
+
+    fn dependencies(&self) -> Vec<Target> {
+        vec![self.integer]
+    }
+
+    fn run_once(
+        &self,
+        witness: &PartitionWitness<F>,
+        out_buffer: &mut GeneratedValues<F>,
+    ) -> anyhow::Result<()> {
+        let integer_value = witness.get_target(self.integer).to_canonical_u64();
+        let low = integer_value & ((1 << self.n_log) - 1);
+        let high = integer_value >> self.n_log;
+
+        out_buffer.set_target(self.low, F::from_canonical_u64(low))?;
+        out_buffer.set_target(self.high, F::from_canonical_u64(high))
+    }
+
+    fn serialize(&self, dst: &mut Vec<u8>, _common_data: &CommonCircuitData<F, D>) -> IoResult<()> {
+        dst.write_target(self.integer)?;
+        dst.write_usize(self.n_log)?;
+        dst.write_target(self.low)?;
+        dst.write_target(self.high)
+    }
+
+    fn deserialize(src: &mut Buffer, _common_data: &CommonCircuitData<F, D>) -> IoResult<Self> {
+        let integer = src.read_target()?;
+        let n_log = src.read_usize()?;
+        let low = src.read_target()?;
+        let high = src.read_target()?;
+        Ok(Self {
+            integer,
+            n_log,
+            low,
+            high,
+        })
     }
-    x
 }
diff --git a/src/backends/plonky2/primitives/ec/curve.rs b/src/backends/plonky2/primitives/ec/curve.rs
@@ -630,7 +630,7 @@ pub trait WitnessWriteCurve: WitnessWrite<GoldilocksField> {
         let digits = value.to_u32_digits();
         for i in 0..10 {
             let d = digits.get(i).copied().unwrap_or(0);
-            self.set_target(target.0[i], GoldilocksField::from_canonical_u32(d))?;
+            self.set_target(target.limbs[i], GoldilocksField::from_canonical_u32(d))?;
         }
         Ok(())
     }
@@ -775,8 +775,8 @@ mod test {
         let p_tgt = builder.constant_point(p);
         let g_scalar_bigint = builder.constant_biguint320(&n2);
         let p_scalar_bigint = builder.constant_biguint320(&n3);
-        let g_scalar_bits = builder.biguint_bits(&g_scalar_bigint);
-        let p_scalar_bits = builder.biguint_bits(&p_scalar_bigint);
+        let g_scalar_bits = g_scalar_bigint.bits;
+        let p_scalar_bits = p_scalar_bigint.bits;
         let e = builder.constant_point((&n2) * g + (&n3) * p);
         let f = builder.linear_combination_points(&g_scalar_bits, &p_scalar_bits, &g_tgt, &p_tgt);
         builder.connect_point(&e, &f);
diff --git a/src/backends/plonky2/primitives/ec/schnorr.rs b/src/backends/plonky2/primitives/ec/schnorr.rs
@@ -117,14 +117,14 @@ impl SignatureTarget {
         public_key: &PointTarget,
     ) -> BoolTarget {
         let g = builder.constant_point(Point::generator());
-        let sig1_bits = builder.biguint_bits(&self.s);
-        let sig2_bits = builder.biguint_bits(&self.e);
+        let sig1_bits = self.s.bits;
+        let sig2_bits = self.e.bits;
         let r = builder.linear_combination_points(&sig1_bits, &sig2_bits, &g, public_key);
         let u_arr = r.u.components;
         let inputs = u_arr.into_iter().chain(msg.elements).collect::<Vec<_>>();
         let e_hash = hash_array_circuit(builder, &inputs);
         let e = builder.field_elements_to_biguint(&e_hash);
-        builder.is_equal_slice(&self.e.0, &e.0)
+        builder.is_equal_slice(&self.e.limbs, &e.limbs)
     }
 }
 
@@ -316,6 +316,7 @@ mod test {
         let int_const = builder.constant_biguint320(&hash_int);
         let int_circuit = builder.field_elements_to_biguint(&hash_const);
         builder.connect_biguint320(&int_const, &int_circuit);
+        println!("{}", builder.num_gates());
         let pw = PartialWitness::new();
         let data = builder.build::<PoseidonGoldilocksConfig>();
         let proof = data.prove(pw)?;
