Update Prover functions.

Author: silathdiir

prover/src/aggregator/prover.rs

@@ -7,11 +7,12 @@ use std::collections::HashMap;
 
 mod aggregation;
 mod chunk;
+mod common;
 mod compression;
 mod utils;
 
 #[derive(Debug)]
 pub struct Prover {
-    inner_params: ParamsKZG<Bn256>,
-    inner_pks: HashMap<String, ProvingKey<G1Affine>>,
+    params: ParamsKZG<Bn256>,
+    pks: HashMap<String, ProvingKey<G1Affine>>,
 }

prover/src/aggregator/prover/aggregation.rs

@@ -1,15 +1,10 @@
 use super::Prover;
-use crate::proof::Proof;
 use aggregator::AggregationCircuit;
 use anyhow::Result;
 use snark_verifier_sdk::Snark;
 
 impl Prover {
-    pub fn build_agg_circuit(&self, _snarks: Vec<Snark>) -> Result<AggregationCircuit> {
-        todo!()
-    }
-
-    pub fn gen_agg_proof(&self, _snarks: Vec<Snark>) -> Result<Proof> {
+    pub fn gen_agg_proof(&self, _snarks: Vec<Snark>) -> Result<(AggregationCircuit, Snark)> {
         todo!()
     }
 }

prover/src/aggregator/prover/chunk.rs

@@ -1,13 +1,11 @@
-use super::{utils::tick, Prover};
+use super::{utils::gen_rng, Prover};
 use crate::{
     utils::{metric_of_witness_block, read_env_var},
     zkevm::circuit::{block_traces_to_witness_block, check_batch_capacity, TargetCircuit, DEGREE},
 };
 use anyhow::{bail, Result};
-use halo2_proofs::{dev::MockProver, halo2curves::bn256::Fr, plonk::keygen_pk2};
+use halo2_proofs::{dev::MockProver, halo2curves::bn256::Fr};
 use once_cell::sync::Lazy;
-use rand::SeedableRng;
-use rand_xorshift::XorShiftRng;
 use snark_verifier_sdk::{gen_snark_shplonk, Snark};
 use types::eth::BlockTrace;
 
@@ -34,7 +32,7 @@ impl Prover {
 
         let (circuit, instance) = C::from_witness_block(&witness_block)?;
 
-        // generate the proof for the inner circuit
+        // Generate the proof for the inner circuit.
         log::info!(
             "Create {} proof of block {} ... block {}, batch len {}",
             C::name(),
@@ -43,9 +41,6 @@ impl Prover {
             chunk_trace.len()
         );
 
-        let seed = [0u8; 16];
-        let mut rng = XorShiftRng::from_seed(seed);
-
         if *MOCK_PROVE {
             log::info!("Mock prove {} start", C::name());
             let prover = MockProver::<Fr>::run(*DEGREE as u32, &circuit, instance)?;
@@ -59,25 +54,17 @@ impl Prover {
             log::info!("Mock prove {} done", C::name());
         }
 
-        if !self.inner_pks.contains_key(&C::name()) {
-            self.gen_inner_pk::<C>(&C::dummy_inner_circuit());
+        // Reuse pk.
+        let id = C::name();
+        if !self.pks.contains_key(&id) {
+            self.gen_inner_pk::<C>(&circuit)?;
         }
-        let pk = &self.inner_pks[&C::name()];
+        let pk = &self.pks[&id];
 
         // Generate the SNARK proof for inner circuit.
         let snark_proof =
-            gen_snark_shplonk(&self.inner_params, pk, circuit, &mut rng, None::<String>);
+            gen_snark_shplonk(&self.params, pk, circuit, &mut gen_rng(), None::<String>);
 
         Ok(snark_proof)
     }
-
-    fn gen_inner_pk<C: TargetCircuit>(&mut self, circuit: &<C as TargetCircuit>::Inner) {
-        tick(&format!("Before init pk of {}", C::name()));
-
-        let pk = keygen_pk2(&self.inner_params, circuit)
-            .unwrap_or_else(|e| panic!("Failed to generate {} pk: {:?}", C::name(), e));
-        self.inner_pks.insert(C::name(), pk);
-
-        tick(&format!("After init pk of {}", C::name()));
-    }
 }

prover/src/aggregator/prover/common.rs

@@ -0,0 +1,63 @@
+use super::{utils::tick, Prover};
+use crate::zkevm::circuit::TargetCircuit;
+use anyhow::Result;
+use halo2_proofs::{
+    halo2curves::bn256::{Bn256, Fr},
+    plonk::keygen_pk2,
+    poly::kzg::commitment::ParamsKZG,
+};
+use rand::Rng;
+use snark_verifier_sdk::{gen_pk, gen_snark_shplonk, verify_snark_shplonk, CircuitExt, Snark};
+
+impl Prover {
+    pub(crate) fn gen_snark(
+        &mut self,
+        id: &str,
+        rng: &mut (impl Rng + Send),
+        params: &ParamsKZG<Bn256>,
+        circuit: impl CircuitExt<Fr> + Clone,
+    ) -> Snark {
+        // Reuse pk.
+        if !self.pks.contains_key(id) {
+            self.gen_outer_pk(id, params, circuit.clone());
+        }
+        let pk = &self.pks[id];
+
+        gen_snark_shplonk(params, pk, circuit, rng, None::<&str>)
+    }
+
+    pub(crate) fn verify_snark<C: CircuitExt<Fr>>(
+        &self,
+        id: &str,
+        params: &ParamsKZG<Bn256>,
+        snark: Snark,
+    ) -> bool {
+        // Must have cached pk.
+        let pk = &self.pks[id];
+
+        verify_snark_shplonk::<C>(params, snark, pk.get_vk())
+    }
+
+    pub(crate) fn gen_inner_pk<C: TargetCircuit>(
+        &mut self,
+        circuit: &<C as TargetCircuit>::Inner,
+    ) -> Result<()> {
+        let id = C::name();
+
+        tick(&format!("Before generate inner pk of {}", &id));
+        let pk = keygen_pk2(&self.params, circuit)?;
+        tick(&format!("After generate inner pk of {}", &id));
+
+        self.pks.insert(id, pk);
+
+        Ok(())
+    }
+
+    fn gen_outer_pk(&mut self, id: &str, params: &ParamsKZG<Bn256>, circuit: impl CircuitExt<Fr>) {
+        tick(&format!("Before generate outer pk of {}", &id));
+        let pk = gen_pk(params, &circuit, None);
+        tick(&format!("After generate outer pk of {}", &id));
+
+        self.pks.insert(id.to_string(), pk);
+    }
+}

prover/src/aggregator/prover/compression.rs

@@ -1,10 +1,23 @@
 use super::Prover;
-use crate::proof::Proof;
-use anyhow::Result;
+use aggregator::CompressionCircuit;
+use halo2_proofs::poly::commitment::Params;
+use rand::Rng;
 use snark_verifier_sdk::Snark;
 
 impl Prover {
-    pub fn gen_comp_proof(&self, _snarks: Vec<Snark>) -> Result<Proof> {
-        todo!()
+    pub fn gen_comp_proof(
+        &mut self,
+        id: &str,
+        is_fresh: bool,
+        degree: u32,
+        mut rng: impl Rng + Send,
+        prev_snark: Snark,
+    ) -> Snark {
+        let mut params = self.params.clone();
+        params.downsize(degree);
+
+        let circuit = CompressionCircuit::new(&params, prev_snark, is_fresh, &mut rng);
+
+        self.gen_snark(id, &mut rng, &params, circuit)
     }
 }

prover/src/aggregator/prover/utils.rs

@@ -1,3 +1,12 @@
+use rand::Rng;
+use rand::SeedableRng;
+use rand_xorshift::XorShiftRng;
+
+pub(crate) fn gen_rng() -> impl Rng + Send {
+    let seed = [0u8; 16];
+    XorShiftRng::from_seed(seed)
+}
+
 pub(crate) fn tick(desc: &str) {
     #[cfg(target_os = "linux")]
     let memory = match procfs::Meminfo::new() {