WIP: HyperKZG verifier: make pairings work

a16z · imikushin · Jul 12, 2024 · Jul 25, 2024 · Jul 28, 2024 · Jul 29, 2024
commit 4ab444d56edb52cb2fac6d4387f42f65f41aff38
diff --git a/jolt-core/src/circuits/groups/curves/mod.rs b/jolt-core/src/circuits/groups/curves/mod.rs
@@ -121,16 +121,9 @@ mod tests {
     {
         type Circuit = DelayedOpsCircuit<E, G1Var>;
 
-        fn offloaded_setup(
-            circuit: Self::Circuit,
-            snark_vk: S::ProcessedVerifyingKey,
-        ) -> Result<OffloadedSNARKVerifyingKey<E, S>, OffloadedSNARKError<S::Error>> {
-            Ok(OffloadedSNARKVerifyingKey {
-                snark_pvk: snark_vk,
-                delayed_pairings: vec![], // TODO none yet
-                g2_elements: vec![],
-            })
-        }
+        // fn pairing_setup(circuit: Self::Circuit) -> Vec<Vec<E::G2Affine>> {
+        //     vec![]
+        // }
     }
 
     #[test]

diff --git a/jolt-core/src/circuits/offloaded/mod.rs b/jolt-core/src/circuits/offloaded/mod.rs
@@ -3,15 +3,32 @@ use ark_ec::pairing::Pairing;
 use ark_ec::{CurveGroup, VariableBaseMSM};
 use ark_ff::{One, PrimeField};
 use ark_r1cs_std::alloc::AllocVar;
+use ark_r1cs_std::boolean::Boolean;
 use ark_r1cs_std::eq::EqGadget;
 use ark_r1cs_std::fields::fp::FpVar;
 use ark_r1cs_std::fields::FieldVar;
 use ark_r1cs_std::groups::CurveVar;
 use ark_r1cs_std::{R1CSVar, ToConstraintFieldGadget};
 use ark_relations::ns;
 use ark_relations::r1cs::{ConstraintSystemRef, Namespace, SynthesisError};
+use ark_serialize::Valid;
+use ark_std::Zero;
 use std::marker::PhantomData;
 
+pub trait MSMGadget<FVar, G, GVar>
+where
+    FVar: FieldVar<G::ScalarField, G::ScalarField> + ToConstraintFieldGadget<G::ScalarField>,
+    G: CurveGroup,
+    GVar: CurveVar<G, G::ScalarField> + ToConstraintFieldGadget<G::ScalarField>,
+{
+    fn msm(
+        &self,
+        cs: impl Into<Namespace<G::ScalarField>>,
+        g1s: &[GVar],
+        scalars: &[FVar],
+    ) -> Result<GVar, SynthesisError>;
+}
+
 pub struct OffloadedMSMGadget<'a, FVar, G, GVar, Circuit>
 where
     Circuit: OffloadedDataCircuit<G>,
@@ -98,10 +115,13 @@ where
                     scalar_input.enforce_equal(&x)?;
                 }
 
+                let mut offsets = vec![];
+
                 // write g1s to public_input
                 for g1 in g1s {
                     let f_vec = g1.to_constraint_field()?;
 
+                    offsets.push(cs.num_instance_variables() - 1);
                     for f in f_vec.iter() {
                         let f_input = FpVar::new_input(ns!(cs, "g1s"), || f.value())?;
                         f_input.enforce_equal(f)?;
@@ -110,8 +130,10 @@ where
 
                 // write msm_g1 to public_input
                 {
+                    dbg!(cs.num_instance_variables() - 1);
                     let f_vec = msm_g1_var.to_constraint_field()?;
 
+                    offsets.push(cs.num_instance_variables() - 1);
                     for f in f_vec.iter() {
                         let f_input = FpVar::new_input(ns!(cs, "msm_g1"), || f.value())?;
                         f_input.enforce_equal(f)?;
@@ -120,7 +142,7 @@ where
                 dbg!(cs.num_constraints());
                 dbg!(cs.num_instance_variables());
 
-                Ok(full_msm_value)
+                Ok((full_msm_value, offsets))
             })
         };
         dbg!(cs.num_constraints());
@@ -129,16 +151,114 @@ where
     }
 }
 
-pub trait MSMGadget<FVar, G, GVar>
+pub trait PairingGadget<E, G1Var>
 where
-    FVar: FieldVar<G::ScalarField, G::ScalarField> + ToConstraintFieldGadget<G::ScalarField>,
-    G: CurveGroup,
-    GVar: CurveVar<G, G::ScalarField> + ToConstraintFieldGadget<G::ScalarField>,
+    E: Pairing,
+    G1Var: CurveVar<E::G1, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
 {
-    fn msm(
+    fn multi_pairing_is_zero(
         &self,
-        cs: impl Into<Namespace<G::ScalarField>>,
+        cs: impl Into<Namespace<E::ScalarField>>,
+        g1s: &[G1Var],
+        g2s: &[E::G2Affine],
+    ) -> Result<(), SynthesisError>;
+}
+
+pub struct OffloadedPairingGadget<'a, E, FVar, GVar, Circuit>
+where
+    E: Pairing,
+    Circuit: OffloadedDataCircuit<E::G1>,
+    FVar: FieldVar<E::ScalarField, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
+    GVar: CurveVar<E::G1, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
+{
+    _params: PhantomData<(E, FVar, GVar)>,
+    circuit: &'a Circuit,
+}
+
+impl<'a, E, FVar, GVar, Circuit> OffloadedPairingGadget<'a, E, FVar, GVar, Circuit>
+where
+    E: Pairing,
+    Circuit: OffloadedDataCircuit<E::G1>,
+    FVar: FieldVar<E::ScalarField, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
+    GVar: CurveVar<E::G1, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
+{
+    pub(crate) fn new(circuit: &'a Circuit) -> Self {
+        Self {
+            _params: PhantomData,
+            circuit,
+        }
+    }
+}
+
+impl<'a, E, FVar, GVar, Circuit> PairingGadget<E, GVar>
+    for OffloadedPairingGadget<'a, E, FVar, GVar, Circuit>
+where
+    E: Pairing,
+    Circuit: OffloadedDataCircuit<E::G1>,
+    FVar: FieldVar<E::ScalarField, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
+    GVar: CurveVar<E::G1, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
+{
+    fn multi_pairing_is_zero(
+        &self,
+        cs: impl Into<Namespace<E::ScalarField>>,
         g1s: &[GVar],
-        scalars: &[FVar],
-    ) -> Result<GVar, SynthesisError>;
+        g2s: &[E::G2Affine],
+    ) -> Result<(), SynthesisError> {
+        let ns = cs.into();
+        let cs = ns.cs();
+
+        let g1_values_opt = g1s
+            .iter()
+            .map(|g1| g1.value().ok())
+            .collect::<Option<Vec<_>>>();
+
+        let g2_values = g2s;
+
+        let is_zero_opt =
+            g1_values_opt.map(|g1_values| E::multi_pairing(dbg!(&g1_values), g2_values).is_zero());
+        if let Some(false) = is_zero_opt {
+            dbg!("multi_pairing_is_zero: false");
+            return Err(SynthesisError::Unsatisfiable);
+        }
+
+        // {
+        //     let g1s = g1s.to_vec();
+        //     let ns = ns!(cs, "deferred_pairing");
+        //     let cs = ns.cs();
+        //
+        //     self.circuit.defer_msm(move || {
+        //         let mut offsets = vec![];
+        //
+        //         // write g1s to public_input
+        //         for g1 in g1s {
+        //             let f_vec = g1.to_constraint_field()?;
+        //
+        //             offsets.push(cs.num_instance_variables() - 1);
+        //             for f in f_vec.iter() {
+        //                 let f_input = FpVar::new_input(ns!(cs, "g1s"), || f.value())?;
+        //                 f_input.enforce_equal(f)?;
+        //             }
+        //         }
+        //
+        //         // write g2s to public_input
+        //         for g2 in g2s {
+        //             let f_vec = g2.to_constraint_field()?;
+        //
+        //             offsets.push(cs.num_instance_variables() - 1);
+        //             for f in f_vec.iter() {
+        //                 let f_input = FpVar::new_input(ns!(cs, "g2s"), || f.value())?;
+        //                 f_input.enforce_equal(f)?;
+        //             }
+        //         }
+        //
+        //         dbg!(cs.num_constraints());
+        //         dbg!(cs.num_instance_variables());
+        //
+        //         Ok(())
+        //     })
+        // }
+
+        dbg!("multi_pairing_is_zero: success");
+        Ok(())
+    }
 }
diff --git a/jolt-core/src/circuits/poly/commitment/hyperkzg.rs b/jolt-core/src/circuits/poly/commitment/hyperkzg.rs
@@ -1,4 +1,6 @@
-use crate::circuits::offloaded::{MSMGadget, OffloadedMSMGadget};
+use crate::circuits::offloaded::{
+    MSMGadget, OffloadedMSMGadget, OffloadedPairingGadget, PairingGadget,
+};
 use crate::circuits::poly::commitment::commitment_scheme::CommitmentVerifierGadget;
 use crate::circuits::transcript::ImplAbsorb;
 use crate::field::JoltField;
@@ -129,6 +131,7 @@ where
 {
     _params: PhantomData<(E, S, G1Var)>,
     circuit: &'a Circuit,
+    g2_elements: Vec<E::G2Affine>,
 }
 
 impl<'a, E, S, G1Var, Circuit> HyperKZGVerifierGadget<'a, E, S, G1Var, Circuit>
@@ -138,10 +141,11 @@ where
     G1Var: CurveVar<E::G1, E::ScalarField> + ToConstraintFieldGadget<E::ScalarField>,
     Circuit: OffloadedDataCircuit<E::G1>,
 {
-    pub fn new(circuit: &'a Circuit) -> Self {
+    pub fn new(circuit: &'a Circuit, g2_elements: Vec<E::G2Affine>) -> Self {
         Self {
             _params: PhantomData,
             circuit,
+            g2_elements,
         }
     }
 }
@@ -187,6 +191,8 @@ where
             .next()
             .unwrap();
 
+        dbg!(r.value());
+
         let u = vec![r.clone(), r.negate()?, r.clone() * &r];
 
         let com = [vec![c.clone()], com.clone()].concat();
@@ -215,22 +221,25 @@ where
 
         // kzg_verify_batch
 
-        transcript.absorb(&v.iter().flatten().cloned().collect::<Vec<_>>())?;
-        let q_powers = q_powers::<E, S>(transcript, ell)?;
-
         transcript.absorb(
-            &proof
-                .w
-                .iter()
-                .map(|g| ImplAbsorb::wrap(g))
+            &v.iter()
+                .flatten()
+                .map(|v_ij| ImplAbsorb::wrap(v_ij))
                 .collect::<Vec<_>>(),
         )?;
+        let q_powers = q_powers::<E, S>(transcript, ell)?;
+
+        dbg!(q_powers.value());
+
+        transcript.absorb(&w.iter().map(|g| ImplAbsorb::wrap(g)).collect::<Vec<_>>())?;
         let d = transcript
             .squeeze_field_elements(1)?
             .into_iter()
             .next()
             .unwrap();
 
+        dbg!(d.value());
+
         let d_square = d.square()?;
         let q_power_multiplier = one + &d + &d_square;
         let q_powers_multiplied = q_powers
@@ -250,9 +259,11 @@ where
             .collect::<Vec<_>>();
 
         let msm_gadget = OffloadedMSMGadget::<FpVar<F>, E::G1, G1Var, Circuit>::new(self.circuit);
+        let pairing_gadget =
+            OffloadedPairingGadget::<E, FpVar<F>, G1Var, Circuit>::new(self.circuit);
 
-        let g1s = &[com.as_slice(), w.as_slice(), &[g1.clone()]].concat();
-        let scalars = &[
+        let l_g1s = &[com.as_slice(), w.as_slice(), &[g1.clone()]].concat();
+        let l_scalars = &[
             q_powers_multiplied.as_slice(),
             &[
                 u[0].clone(),
@@ -262,17 +273,29 @@ where
             ],
         ]
         .concat();
+        debug_assert_eq!(l_g1s.len(), l_scalars.len());
 
-        let l_g1 = msm_gadget.msm(ns!(transcript.cs(), "l_g1"), g1s, scalars)?;
+        dbg!(transcript.cs().num_instance_variables() - 1);
+        let l_g1 = msm_gadget.msm(ns!(transcript.cs(), "l_g1"), l_g1s, l_scalars)?;
 
-        let g1s = w.as_slice();
-        let scalars = &[FpVar::one(), d, d_square];
-        debug_assert_eq!(g1s.len(), scalars.len());
+        dbg!(w.as_slice().value());
 
-        let r_g1 = msm_gadget.msm(ns!(transcript.cs(), "r_g1"), g1s, scalars)?;
+        let r_g1s = w.as_slice();
+        let r_scalars = &[FpVar::one().negate()?, d.negate()?, d_square.negate()?];
+        debug_assert_eq!(r_g1s.len(), r_scalars.len());
+
+        dbg!(transcript.cs().num_instance_variables() - 1);
+        let r_g1 = msm_gadget.msm(ns!(transcript.cs(), "r_g1"), r_g1s, r_scalars)?;
+
+        // (dbg!(l_g1.value()), dbg!(r_g1.value()));
+
+        pairing_gadget.multi_pairing_is_zero(
+            ns!(transcript.cs(), "multi_pairing"),
+            &[l_g1, r_g1],
+            self.g2_elements.as_slice(),
+        )?;
         dbg!();
 
-        // TODO implement
         Ok(Boolean::TRUE)
     }
 }
@@ -311,8 +334,8 @@ mod tests {
     use crate::poly::commitment::kzg::KZGVerifierKey;
     use crate::poly::dense_mlpoly::DensePolynomial;
     use crate::snark::{
-        DeferredFnsRef, OffloadedDataCircuit, OffloadedSNARK, OffloadedSNARKError,
-        OffloadedSNARKVerifyingKey,
+        DeferredFnsRef, OffloadedDataCircuit, OffloadedPairingDef, OffloadedSNARK,
+        OffloadedSNARKError, OffloadedSNARKVerifyingKey,
     };
     use crate::utils::errors::ProofVerifyError;
     use crate::utils::transcript::ProofTranscript;
@@ -409,14 +432,14 @@ mod tests {
             let mut transcript_var =
                 MockSpongeVar::new(ns!(cs, "transcript").cs(), &(b"TestEval".as_slice()));
 
-            let h_kzg = HyperKZGVerifierGadget::<
-                E,
-                MockSponge<E::ScalarField>,
-                G1Var,
-                HyperKZGVerifierCircuit<E, G1Var>,
-            >::new(&self);
+            let kzg_vk = self.pcs_pk_vk.1.kzg_vk;
+            let hyper_kzg =
+                HyperKZGVerifierGadget::<E, MockSponge<E::ScalarField>, G1Var, Self>::new(
+                    &self,
+                    vec![kzg_vk.g2, kzg_vk.beta_g2],
+                );
 
-            let r = h_kzg.verify(
+            let r = hyper_kzg.verify(
                 &proof_var,
                 &vk_var,
                 &mut transcript_var,
@@ -455,18 +478,11 @@ mod tests {
     {
         type Circuit = HyperKZGVerifierCircuit<E, G1Var>;
 
-        fn offloaded_setup(
-            circuit: Self::Circuit,
-            snark_vk: S::ProcessedVerifyingKey,
-        ) -> Result<OffloadedSNARKVerifyingKey<E, S>, OffloadedSNARKError<S::Error>> {
-            let KZGVerifierKey { g1, g2, beta_g2 } = circuit.pcs_pk_vk.1.kzg_vk;
-
-            Ok(OffloadedSNARKVerifyingKey {
-                snark_pvk: snark_vk,
-                delayed_pairings: vec![],
-                g2_elements: vec![vec![g2, beta_g2]],
-            })
-        }
+        // fn pairing_setup(circuit: Self::Circuit) -> Vec<Vec<E::G2Affine>> {
+        //     let KZGVerifierKey { g1, g2, beta_g2 } = circuit.pcs_pk_vk.1.kzg_vk;
+        //
+        //     vec![vec![g2, beta_g2]]
+        // }
     }
 
     #[test]