Fix PI prover nil reference bug (#312)

* fix get setup first
* don't "overbatch"

---------

Co-authored-by: Arya Tabaie <[email protected]>

prover/backend/aggregation/prove.go

@@ -114,13 +114,13 @@ func makePiProof(cfg *config.Config, cf *CollectedFields) (plonk.Proof, witness.
 		return nil, nil, fmt.Errorf("could not extract interconnection circuit public witness: %w", err)
 	}
 
-	proverOpts := emPlonk.GetNativeProverOptions(ecc.BW6_761.ScalarField(), setup.Circuit.Field())
-	verifierOpts := emPlonk.GetNativeVerifierOptions(ecc.BW6_761.ScalarField(), setup.Circuit.Field())
-
 	if err = <-setupErr; err != nil { // wait for setup to load and check for errors
 		return nil, nil, fmt.Errorf("could not load the setup: %w", err)
 	}
 
+	proverOpts := emPlonk.GetNativeProverOptions(ecc.BW6_761.ScalarField(), setup.Circuit.Field())
+	verifierOpts := emPlonk.GetNativeVerifierOptions(ecc.BW6_761.ScalarField(), setup.Circuit.Field())
+
 	proof, err := circuits.ProveCheck(&setup, &assignment, proverOpts, verifierOpts)
 
 	return proof, w, err

prover/circuits/aggregation/circuit.go

@@ -64,6 +64,7 @@ func (c *Circuit) Define(api frontend.API) error {
 
 	assertSlicesEqualZEXT(api, piBits[:16*8], field.ToBitsCanonical(&c.PublicInputWitness.Public[1]))
 	assertSlicesEqualZEXT(api, piBits[16*8:], field.ToBitsCanonical(&c.PublicInputWitness.Public[0]))
+	api.AssertIsDifferent(c.PublicInput, 0) // making sure at least one element of the PI circuit's public input is nonzero to justify using incomplete arithmetic
 
 	vks := append(slices.Clone(c.verifyingKeys), c.publicInputVerifyingKey)
 	piVkIndex := len(vks) - 1
@@ -161,27 +162,44 @@ func verifyClaimBatch(api frontend.API, vks []emVkey, claims []proofClaim) error
 
 	var (
 		bvk       = vks[0].BaseVerifyingKey
-		cvks      = make([]emCircVKey, len(vks))
+		cvks      = make([]emCircVKey, len(vks)-1)
 		switches  = make([]frontend.Variable, len(claims))
 		proofs    = make([]emProof, len(claims))
 		witnesses = make([]emWitness, len(claims))
 	)
 
-	for i := range vks {
+	for i := range cvks {
 		cvks[i] = vks[i].CircuitVerifyingKey
 	}
 
+	for i := 1; i < len(vks)-1; i++ { // TODO @Tabaie make sure these don't generate any constraints
+		fmt.Printf("checking base vk #%d/%d\n", i+1, len(vks)-1)
+		assertBaseKeyEquals(api, bvk, vks[i].BaseVerifyingKey)
+	}
+
 	for i := range claims {
 		proofs[i] = claims[i].Proof
 		switches[i] = claims[i].CircuitID
 		witnesses[i] = claims[i].PublicInput
 	}
 
-	err = verifier.AssertDifferentProofs(bvk, cvks, switches, proofs, witnesses, emPlonk.WithCompleteArithmetic())
-	if err != nil {
+	lastProofI := len(proofs) - 1
+	if err = verifier.AssertDifferentProofs(
+		bvk, cvks,
+		switches[:lastProofI],
+		proofs[:lastProofI],
+		witnesses[:lastProofI],
+		emPlonk.WithCompleteArithmetic(),
+	); err != nil {
 		return fmt.Errorf("AssertDifferentProofs returned an error: %w", err)
 	}
 
+	// The PI proof cannot be batched with the rest because it has more than one public input
+	// complete arithmetic is not necessary here because the circuit is nontrivial and at least one element of the public input is nonzero
+	if err = verifier.AssertProof(vks[len(cvks)], proofs[lastProofI], witnesses[lastProofI]); err != nil {
+		return fmt.Errorf("AssertProof returned an error: %w", err)
+	}
+
 	return nil
 }
 
@@ -198,3 +216,25 @@ func assertSlicesEqualZEXT(api frontend.API, a, b []frontend.Variable) {
 		api.AssertIsEqual(b[i], 0)
 	}
 }
+
+// assertBaseKeyEquals is very aggressive in equality testing between emulated elements. The representations have to be exactly equal, not only equal modulo the group size
+func assertBaseKeyEquals(api frontend.API, a, b emPlonk.BaseVerifyingKey[emFr, emG1, emG2]) {
+
+	internal.AssertSliceEquals(api, a.CosetShift.Limbs, b.CosetShift.Limbs)
+
+	assertG2AffEquals := func(a, b sw_bls12377.G2Affine) {
+		api.AssertIsEqual(a.P.X.A0, b.P.X.A0)
+		api.AssertIsEqual(a.P.X.A1, b.P.X.A1)
+		api.AssertIsEqual(a.P.Y.A0, b.P.Y.A0)
+		api.AssertIsEqual(a.P.Y.A1, b.P.Y.A1)
+	}
+
+	api.AssertIsEqual(a.Kzg.G1.X, b.Kzg.G1.X)
+	api.AssertIsEqual(a.Kzg.G1.Y, b.Kzg.G1.Y)
+	assertG2AffEquals(a.Kzg.G2[0], b.Kzg.G2[0])
+	assertG2AffEquals(a.Kzg.G2[1], b.Kzg.G2[1])
+
+	// NOT CHECKING THE LINE EVALUATIONS
+
+	api.AssertIsEqual(a.NbPublicVariables, b.NbPublicVariables)
+}

prover/circuits/setup_manifest.go

@@ -51,12 +51,16 @@ func (m *SetupManifest) GetInt(key string) (int, error) {
 		return 0, err
 	}
 
-	i, ok := v.(int)
-	if !ok {
-		return 0, fmt.Errorf("flag `%s` is not an int", key)
+	switch x := v.(type) {
+	case int:
+		return x, nil
+	case float64:
+		if i := int(x); float64(i) == x {
+			return i, nil
+		}
 	}
 
-	return i, nil
+	return 0, fmt.Errorf("flag `%s` is not an int", key)
 }
 
 func (m *SetupManifest) GetString(key string) (string, error) {