Feat/hybrid snark agg (#39)

* load preprocessed as witness

* change fn name

* return assigned values (so that loader's rc is not still borrowed)

* chore: rename function

---------

Co-authored-by: Rohit Narurkar <[email protected]>

snark-verifier-sdk/src/aggregation.rs

@@ -10,13 +10,16 @@ use ark_std::start_timer;
 use halo2_base::{
     halo2_proofs::{
         circuit::Value,
-        halo2curves::bn256::{Fr, G1Affine},
+        halo2curves::bn256::{Fq, Fr, G1Affine},
     },
     AssignedValue,
 };
 use itertools::Itertools;
 use snark_verifier::{
-    loader::halo2::EccInstructions,
+    loader::halo2::{
+        halo2_ecc::{ecc::EccChip, fields::fp::FpConfig},
+        EccInstructions,
+    },
     pcs::{
         kzg::{KzgAccumulator, KzgAs},
         AccumulationScheme, MultiOpenScheme, PolynomialCommitmentScheme,
@@ -41,8 +44,8 @@ pub fn load_verify_circuit_degree() -> u32 {
     params.degree
 }
 
-pub fn flatten_accumulator<'a>(
-    accumulator: KzgAccumulator<G1Affine, Rc<Halo2Loader<'a>>>,
+pub fn flatten_accumulator(
+    accumulator: KzgAccumulator<G1Affine, Rc<Halo2Loader>>,
 ) -> Vec<AssignedValue<Fr>> {
     let KzgAccumulator { lhs, rhs } = accumulator;
     let lhs = lhs.into_assigned();
@@ -52,27 +55,120 @@ pub fn flatten_accumulator<'a>(
         .truncation
         .limbs
         .into_iter()
-        .chain(lhs.y.truncation.limbs.into_iter())
-        .chain(rhs.x.truncation.limbs.into_iter())
-        .chain(rhs.y.truncation.limbs.into_iter())
+        .chain(lhs.y.truncation.limbs)
+        .chain(rhs.x.truncation.limbs)
+        .chain(rhs.y.truncation.limbs)
         .collect()
 }
 
+type AssignedInstances<'a> =
+    Vec<Vec<<BaseFieldEccChip as EccInstructions<'a, G1Affine>>::AssignedScalar>>;
+type KzgAcc<'a> = KzgAccumulator<G1Affine, Rc<Halo2Loader<'a>>>;
+type AssignedEcPoints<'a> =
+    Vec<Vec<<EccChip<Fr, FpConfig<Fr, Fq>> as EccInstructions<'a, G1Affine>>::AssignedEcPoint>>;
+type AssignedScalars<'a> =
+    Vec<Option<<EccChip<Fr, FpConfig<Fr, Fq>> as EccInstructions<'a, G1Affine>>::AssignedScalar>>;
+
+/// Core function used in `synthesize` to aggregate multiple `snarks`.
+///  
+/// Returns the assigned instances of previous snarks and the new final pair that needs to be verified in a pairing check.
+/// For each previous snark, we concatenate all instances into a single vector. We return a vector of vectors,
+/// one vector per snark, for convenience.
+///
+/// The difference between [`aggregate_as_witness`] and [`aggregate`] is that `aggregate` loads
+/// the [`Protocol`][snark_verifier::Protocol] of all SNARKs as constants. Whereas in
+/// [`aggregate_as_witness`] we load them as witness and return the preprocessed polynomials and
+/// the transcript initial states.
 #[allow(clippy::type_complexity)]
+pub fn aggregate_as_witness<'a, PCS>(
+    svk: &PCS::SuccinctVerifyingKey,
+    loader: &Rc<Halo2Loader<'a>>,
+    snarks: &[SnarkWitness],
+    as_proof: Value<&'_ [u8]>,
+) -> (AssignedInstances<'a>, KzgAcc<'a>, AssignedEcPoints<'a>, AssignedScalars<'a>)
+where
+    PCS: PolynomialCommitmentScheme<
+            G1Affine,
+            Rc<Halo2Loader<'a>>,
+            Accumulator = KzgAccumulator<G1Affine, Rc<Halo2Loader<'a>>>,
+        > + MultiOpenScheme<G1Affine, Rc<Halo2Loader<'a>>>,
+{
+    let assign_instances = |instances: &[Vec<Value<Fr>>]| {
+        instances
+            .iter()
+            .map(|instances| {
+                instances.iter().map(|instance| loader.assign_scalar(*instance)).collect_vec()
+            })
+            .collect_vec()
+    };
+
+    let mut transcript = PoseidonTranscript::<Rc<Halo2Loader<'a>>, _>::from_spec(
+        loader,
+        Value::unknown(),
+        POSEIDON_SPEC.clone(),
+    );
+
+    let mut previous_instances = Vec::with_capacity(snarks.len());
+    let mut preprocessed_polys = Vec::with_capacity(snarks.len());
+    let mut transcript_init_states = Vec::with_capacity(snarks.len());
+    let mut accumulators = snarks
+        .iter()
+        .flat_map(|snark| {
+            // The SNARK protocol's preprocessed polynomials and the initial state of the
+            // prover/verifier's transcripts are loaded as witnesses.
+            //
+            // This allows us to aggregate SNARKs that may have been generated using one or more
+            // circuits.
+            //
+            // Note: Its important to further constrain the assigned preprocessed polynomial commitments
+            // and the assigned transcript initial state to belong to a fixed expected set.
+            let protocol = snark.protocol.loaded_preprocessed_as_witness(loader);
+
+            // TODO use 1d vector
+            let instances = assign_instances(&snark.instances);
+
+            // read the transcript and perform Fiat-Shamir
+            // run through verification computation and produce the final pair `succinct`
+            transcript.new_stream(snark.proof());
+            let proof = Plonk::<PCS>::read_proof(svk, &protocol, &instances, &mut transcript);
+            let accumulator = Plonk::<PCS>::succinct_verify(svk, &protocol, &instances, &proof);
+
+            previous_instances.push(
+                instances.into_iter().flatten().map(|scalar| scalar.into_assigned()).collect(),
+            );
+            preprocessed_polys
+                .push(protocol.preprocessed.into_iter().map(|ec| ec.into_assigned()).collect());
+            transcript_init_states
+                .push(protocol.transcript_initial_state.map(|s| s.into_assigned()));
+
+            accumulator
+        })
+        .collect_vec();
+
+    let accumulator = if accumulators.len() > 1 {
+        transcript.new_stream(as_proof);
+        let proof =
+            KzgAs::<PCS>::read_proof(&Default::default(), &accumulators, &mut transcript).unwrap();
+        KzgAs::<PCS>::verify(&Default::default(), &accumulators, &proof).unwrap()
+    } else {
+        accumulators.pop().unwrap()
+    };
+
+    (previous_instances, accumulator, preprocessed_polys, transcript_init_states)
+}
+
 /// Core function used in `synthesize` to aggregate multiple `snarks`.
 ///  
 /// Returns the assigned instances of previous snarks and the new final pair that needs to be verified in a pairing check.
 /// For each previous snark, we concatenate all instances into a single vector. We return a vector of vectors,
 /// one vector per snark, for convenience.
+#[allow(clippy::type_complexity)]
 pub fn aggregate<'a, PCS>(
     svk: &PCS::SuccinctVerifyingKey,
     loader: &Rc<Halo2Loader<'a>>,
     snarks: &[SnarkWitness],
     as_proof: Value<&'_ [u8]>,
-) -> (
-    Vec<Vec<<BaseFieldEccChip as EccInstructions<'a, G1Affine>>::AssignedScalar>>,
-    KzgAccumulator<G1Affine, Rc<Halo2Loader<'a>>>,
-)
+) -> (AssignedInstances<'a>, KzgAcc<'a>)
 where
     PCS: PolynomialCommitmentScheme<
             G1Affine,

snark-verifier-sdk/src/lib.rs

@@ -18,7 +18,7 @@ pub mod types;
 pub use aggregation::aggregation_circuit::AggregationCircuit;
 pub use aggregation::load_verify_circuit_degree;
 pub use aggregation::multi_aggregation_circuit::PublicAggregationCircuit;
-pub use aggregation::{aggregate, flatten_accumulator};
+pub use aggregation::{aggregate, aggregate_as_witness, flatten_accumulator};
 pub use circuit_ext::CircuitExt;
 pub use param::{BITS, LIMBS};
 pub use snark::gen_dummy_snark;