init: check invertible inputs for finexp and to_affine functions

Author: manishbista28

g16ckt/src/gadgets/bn254/final_exponentiation.rs

@@ -5,10 +5,14 @@
 //! the structure used in the main branch.
 
 use ark_ec::bn::BnConfig;
-use ark_ff::{BitIteratorBE, Field};
+use ark_ff::{AdditiveGroup, BitIteratorBE, Field};
 use circuit_component_macro::component;
 
-use crate::{CircuitContext, gadgets::bn254::fq12::Fq12};
+use crate::{
+    CircuitContext,
+    circuit::TRUE_WIRE,
+    gadgets::bn254::fq12::{Fq12, ValidFq12},
+};
 
 pub fn conjugate_native(f: ark_bn254::Fq12) -> ark_bn254::Fq12 {
     ark_bn254::Fq12::new(f.c0, -f.c1)
@@ -98,7 +102,16 @@ pub fn exp_by_neg_x_montgomery<C: CircuitContext>(circuit: &mut C, f: &Fq12) ->
 }
 
 #[component]
-pub fn final_exponentiation_montgomery<C: CircuitContext>(circuit: &mut C, f: &Fq12) -> Fq12 {
+pub fn final_exponentiation_montgomery<C: CircuitContext>(circuit: &mut C, f: &Fq12) -> ValidFq12 {
+    let is_zero = Fq12::equal_constant(circuit, f, &ark_bn254::Fq12::ZERO);
+    let is_valid = circuit.issue_wire();
+    circuit.add_gate(crate::Gate {
+        wire_a: is_zero,
+        wire_b: TRUE_WIRE,
+        wire_c: is_valid,
+        gate_type: crate::GateType::Xor,
+    });
+
     let f_inv = Fq12::inverse_montgomery(circuit, f);
     let f_conjugate = Fq12::conjugate(circuit, f);
     let u = Fq12::mul_montgomery(circuit, &f_inv, &f_conjugate);
@@ -127,7 +140,9 @@ pub fn final_exponentiation_montgomery<C: CircuitContext>(circuit: &mut C, f: &F
     let y18 = Fq12::mul_montgomery(circuit, &r2, &y11);
     let y19 = Fq12::frobenius_montgomery(circuit, &y18, 3);
 
-    Fq12::mul_montgomery(circuit, &y19, &y17)
+    let f = Fq12::mul_montgomery(circuit, &y19, &y17);
+
+    ValidFq12 { f, is_valid }
 }
 
 #[cfg(test)]
@@ -171,6 +186,7 @@ mod tests {
         }
         struct Out {
             value: ark_bn254::Fq12,
+            valid: bool,
         }
         fn encode_fq6_to_wires<M: CircuitMode<WireValue = bool>>(
             val: &ark_bn254::Fq6,
@@ -260,7 +276,7 @@ mod tests {
             }
         }
         impl CircuitOutput<ExecuteMode> for Out {
-            type WireRepr = Fq12Wires;
+            type WireRepr = ValidFq12;
             fn decode(wires: Self::WireRepr, cache: &mut ExecuteMode) -> Self {
                 fn decode_fq6_from_wires(
                     wires: &Fq6Wires,
@@ -298,10 +314,12 @@ mod tests {
                         ark_bn254::Fq2::new(ark_bn254::Fq::from(c2_c0), ark_bn254::Fq::from(c2_c1));
                     ark_bn254::Fq6::new(c0, c1, c2)
                 }
-                let c0 = decode_fq6_from_wires(&wires.0[0], cache);
-                let c1 = decode_fq6_from_wires(&wires.0[1], cache);
+                let c0 = decode_fq6_from_wires(&wires.f.0[0], cache);
+                let c1 = decode_fq6_from_wires(&wires.f.0[1], cache);
+                let valid = cache.lookup_wire(wires.is_valid).expect("missing wire");
                 Self {
                     value: ark_bn254::Fq12::new(c0, c1),
+                    valid,
                 }
             }
         }
@@ -312,6 +330,27 @@ mod tests {
                 final_exponentiation_montgomery(ctx, &input.f)
             });
 
-        assert_eq!(result.output_value.value, expected_m);
+        assert!(
+            result.output_value.valid,
+            "final_exponentiation_montgomery input should be valid"
+        );
+        assert_eq!(
+            result.output_value.value, expected_m,
+            "final_exponentiation_montgomery output should be valid"
+        );
+
+        // Test for non-invertible element
+        let input = In {
+            f: ark_bn254::Fq12::ZERO,
+        };
+        let result: StreamingResult<_, _, Out> =
+            CircuitBuilder::streaming_execute(input, 10_000, |ctx, input| {
+                final_exponentiation_montgomery(ctx, &input.f)
+            });
+
+        assert!(
+            !result.output_value.valid,
+            "final_exponentiation_montgomery input should be invalid"
+        );
     }
 }

g16ckt/src/gadgets/bn254/fq12.rs

@@ -5,7 +5,7 @@ use rand::Rng;
 use super::fq6::Fq6Components;
 use crate::{
     CircuitContext, Gate, WireId,
-    circuit::{FromWires, WiresObject},
+    circuit::{FromWires, WiresArity, WiresObject},
     gadgets::bn254::{fq::Fq, fq2::Fq2, fq6::Fq6},
 };
 
@@ -447,6 +447,58 @@ impl Fq12 {
     }
 }
 
+/// Analogous to Option<Fq12> where `is_valid` carries `false` if variable is None
+#[derive(Clone, Debug)]
+pub struct ValidFq12 {
+    pub f: Fq12,
+    pub is_valid: WireId,
+}
+
+impl WiresObject for ValidFq12 {
+    fn to_wires_vec(&self) -> Vec<WireId> {
+        let mut wires: Vec<WireId> = self.f.0[0]
+            .to_wires_vec()
+            .into_iter()
+            .chain(self.f.0[1].to_wires_vec())
+            .collect();
+        wires.push(self.is_valid);
+        wires
+    }
+
+    fn clone_from(&self, mut wire_gen: &mut impl FnMut() -> WireId) -> Self {
+        ValidFq12 {
+            f: Fq12([
+                self.f.0[0].clone_from(&mut wire_gen),
+                self.f.0[1].clone_from(&mut wire_gen),
+            ]),
+            is_valid: wire_gen(),
+        }
+    }
+}
+
+impl FromWires for ValidFq12 {
+    fn from_wires(wires: &[WireId]) -> Option<Self> {
+        if wires.len() == ValidFq12::ARITY {
+            let mid = Fq6::N_BITS;
+            let fq6_1 = Fq6::from_wires(&wires[..mid])?;
+            let fq6_2 = Fq6::from_wires(&wires[mid..2 * mid])?;
+            let is_valid_wires = &wires[2 * mid..];
+            assert_eq!(is_valid_wires.len(), 1); // single is valid wire
+            let res = ValidFq12 {
+                f: Fq12([fq6_1, fq6_2]),
+                is_valid: is_valid_wires[0],
+            };
+            Some(res)
+        } else {
+            None
+        }
+    }
+}
+
+impl WiresArity for ValidFq12 {
+    const ARITY: usize = Fq12::N_BITS + 1;
+}
+
 #[cfg(test)]
 mod tests {
     use std::{array, str::FromStr};

g16ckt/src/gadgets/bn254/pairing.rs

@@ -22,8 +22,13 @@ use crate::{
     CircuitContext, Fp254Impl,
     circuit::{FromWires, OffCircuitParam, WiresArity, WiresObject},
     gadgets::bn254::{
-        final_exponentiation::final_exponentiation_montgomery, fq::Fq, fq2::Fq2, fq6::Fq6,
-        fq12::Fq12, g1::G1Projective, g2::G2Projective,
+        final_exponentiation::final_exponentiation_montgomery,
+        fq::Fq,
+        fq2::Fq2,
+        fq6::Fq6,
+        fq12::{Fq12, ValidFq12},
+        g1::G1Projective,
+        g2::G2Projective,
     },
 };
 
@@ -883,7 +888,7 @@ pub fn pairing_const_q<C: CircuitContext>(
     circuit: &mut C,
     p: &G1Projective,
     q: &ark_bn254::G2Affine,
-) -> Fq12 {
+) -> ValidFq12 {
     let f = miller_loop_const_q(circuit, p, q);
     final_exponentiation_montgomery(circuit, &f)
 }
@@ -894,7 +899,7 @@ pub fn multi_pairing_const_q<C: CircuitContext>(
     circuit: &mut C,
     ps: &[G1Projective],
     qs: &[ark_bn254::G2Affine],
-) -> Fq12 {
+) -> ValidFq12 {
     let f = multi_miller_loop_const_q(circuit, ps, qs);
     final_exponentiation_montgomery(circuit, &f)
 }
@@ -1146,7 +1151,10 @@ mod tests {
         };
         let result =
             CircuitBuilder::streaming_execute::<_, _, Fq12Output>(input, 10_000, |ctx, w| {
-                ell_montgomery(ctx, &w.f, &w.c, &w.p)
+                ValidFq12 {
+                    f: ell_montgomery(ctx, &w.f, &w.c, &w.p),
+                    is_valid: TRUE_WIRE,
+                }
             });
 
         assert_eq!(result.output_value.value, expected_m);
@@ -1720,6 +1728,7 @@ mod tests {
 
         struct FinalExpOutput {
             value: ark_bn254::Fq12,
+            is_valid: bool,
         }
         impl CircuitInput for FEInput {
             type WireRepr = FEWires;
@@ -1740,7 +1749,7 @@ mod tests {
             }
         }
         impl CircuitOutput<ExecuteMode> for FinalExpOutput {
-            type WireRepr = Fq12;
+            type WireRepr = ValidFq12;
             fn decode(wires: Self::WireRepr, cache: &mut ExecuteMode) -> Self {
                 // Reuse local decoder helpers
                 fn decode_fq6_from_wires(
@@ -1779,10 +1788,14 @@ mod tests {
                         ark_bn254::Fq2::new(ark_bn254::Fq::from(c2_c0), ark_bn254::Fq::from(c2_c1));
                     ark_bn254::Fq6::new(c0, c1, c2)
                 }
-                let c0 = decode_fq6_from_wires(&wires.0[0], cache);
-                let c1 = decode_fq6_from_wires(&wires.0[1], cache);
+                let c0 = decode_fq6_from_wires(&wires.f.0[0], cache);
+                let c1 = decode_fq6_from_wires(&wires.f.0[1], cache);
+                let is_valid = cache
+                    .lookup_wire(wires.is_valid)
+                    .expect("missing wire value");
                 Self {
                     value: ark_bn254::Fq12::new(c0, c1),
+                    is_valid,
                 }
             }
         }
@@ -1794,7 +1807,14 @@ mod tests {
             |ctx, input| final_exponentiation_montgomery(ctx, &input.f),
         );
 
-        assert_eq!(result.output_value.value, expected_m);
+        assert_eq!(
+            result.output_value.value, expected_m,
+            "final_exponentiation_montgomery output should be valid"
+        );
+        assert!(
+            result.output_value.is_valid,
+            "final_exponentiation_montgomery input should be valid"
+        );
     }
 
     #[test]
@@ -1914,14 +1934,19 @@ mod tests {
 
     struct Fq12Output {
         value: ark_bn254::Fq12,
+        valid: bool,
     }
     impl CircuitOutput<ExecuteMode> for Fq12Output {
-        type WireRepr = Fq12;
+        type WireRepr = ValidFq12;
         fn decode(wires: Self::WireRepr, cache: &mut ExecuteMode) -> Self {
-            let c0 = decode_fq6_from_wires(&wires.0[0], cache);
-            let c1 = decode_fq6_from_wires(&wires.0[1], cache);
+            let c0 = decode_fq6_from_wires(&wires.f.0[0], cache);
+            let c1 = decode_fq6_from_wires(&wires.f.0[1], cache);
+            let is_valid = cache
+                .lookup_wire(wires.is_valid)
+                .expect("missing wire value");
             Self {
                 value: ark_bn254::Fq12::new(c0, c1),
+                valid: is_valid,
             }
         }
     }
@@ -2018,7 +2043,10 @@ mod tests {
         };
         let result =
             CircuitBuilder::streaming_execute::<_, _, Fq12Output>(input, 10_000, |ctx, input| {
-                ell_eval_const(ctx, &input.f, &coeff, &input.p)
+                ValidFq12 {
+                    f: ell_eval_const(ctx, &input.f, &coeff, &input.p),
+                    is_valid: TRUE_WIRE,
+                }
             });
 
         assert_eq!(result.output_value.value, expected_m);
@@ -2095,7 +2123,10 @@ mod tests {
         let result = CircuitBuilder::streaming_execute::<_, _, Fq12Output>(
             In { p },
             10_000,
-            |ctx, input| miller_loop_const_q(ctx, &input.p, &q),
+            |ctx, input| ValidFq12 {
+                f: miller_loop_const_q(ctx, &input.p, &q),
+                is_valid: TRUE_WIRE,
+            },
         );
 
         assert_eq!(result.output_value.value, expected_m);
@@ -2175,7 +2206,8 @@ mod tests {
             |ctx, input| pairing_const_q(ctx, &input.p, &q),
         );
 
-        assert_eq!(result.output_value.value, expected_m);
+        assert!(result.output_value.valid, "input should be valid");
+        assert_eq!(result.output_value.value, expected_m, "output should match");
     }
 
     #[test]
@@ -2511,7 +2543,10 @@ mod tests {
         let result = CircuitBuilder::streaming_execute::<_, _, Fq12Output>(
             input,
             40_000,
-            |circuit, wires| multi_miller_loop_montgomery_fast(circuit, &wires.ps, &wires.qs),
+            |circuit, wires| ValidFq12 {
+                f: multi_miller_loop_montgomery_fast(circuit, &wires.ps, &wires.qs),
+                is_valid: TRUE_WIRE,
+            },
         );
 
         assert_eq!(result.output_value.value, expected_m);
@@ -2730,9 +2765,12 @@ mod tests {
             let input = In { p1, p2, p3, q3 };
             let result =
                 CircuitBuilder::streaming_execute::<_, _, Fq12Output>(input, 80_000, |ctx, w| {
-                    multi_miller_loop_groth16_evaluate_montgomery_fast(
-                        ctx, &w.p1, &w.p2, &w.p3, q1, q2, &w.q3,
-                    )
+                    ValidFq12 {
+                        f: multi_miller_loop_groth16_evaluate_montgomery_fast(
+                            ctx, &w.p1, &w.p2, &w.p3, q1, q2, &w.q3,
+                        ),
+                        is_valid: TRUE_WIRE,
+                    }
                 });
 
             assert_eq!(