Add gen_pk flag when init chunk-prover.

bin/src/zkevm_prove.rs

@@ -27,7 +27,7 @@ fn main() {
     log::info!("Initialized ENV and created output-dir {output_dir}");
 
     let args = Args::parse();
-    let mut prover = Prover::from_params_dir(&args.params_path);
+    let mut prover = Prover::from_params_dir(&args.params_path, false, None);
 
     let mut traces = Vec::new();
     let trace_path = PathBuf::from(&args.trace_path);

prover/src/common/prover/chunk.rs

@@ -1,11 +1,21 @@
 use super::Prover;
-use crate::config::{LAYER1_DEGREE, LAYER2_DEGREE};
+use crate::{
+    config::{LayerId, LAYER1_DEGREE, LAYER2_DEGREE},
+    zkevm::circuit::block_traces_to_witness_block,
+};
 use anyhow::Result;
 use halo2_proofs::halo2curves::bn256::Fr;
 use snark_verifier_sdk::Snark;
 use zkevm_circuits::evm_circuit::witness::Block;
 
 impl Prover {
+    pub fn gen_chunk_pk(&mut self, output_dir: Option<&str>) -> Result<()> {
+        let empty_block = block_traces_to_witness_block(&[])?;
+        let layer1_snark = self.load_or_gen_last_chunk_snark("empty", &empty_block, output_dir)?;
+
+        self.gen_comp_pk(LayerId::Layer2, layer1_snark)
+    }
+
     pub fn load_or_gen_final_chunk_snark(
         &mut self,
         name: &str,

prover/src/common/prover/compression.rs

@@ -1,6 +1,6 @@
 use super::Prover;
 use crate::{
-    config::layer_config_path,
+    config::{layer_config_path, LayerId},
     io::{load_snark, write_snark},
     utils::gen_rng,
 };
@@ -11,6 +11,22 @@ use snark_verifier_sdk::Snark;
 use std::env;
 
 impl Prover {
+    pub fn gen_comp_pk(&mut self, layer: LayerId, prev_snark: Snark) -> Result<()> {
+        let id = layer.id();
+        let degree = layer.degree();
+        let config_path = layer.config_path();
+
+        let rng = gen_rng();
+        env::set_var("COMPRESSION_CONFIG", config_path);
+        let circuit = CompressionCircuit::new(self.params(degree), prev_snark, true, rng)
+            .map_err(|err| anyhow!("Failed to construct compression circuit: {err:?}"))?;
+        Self::assert_if_mock_prover(id, degree, &circuit);
+
+        self.params_and_pk(id, degree, &circuit)?;
+
+        Ok(())
+    }
+
     pub fn gen_comp_snark(
         &mut self,
         id: &str,

prover/src/test_util.rs

@@ -1,5 +1,6 @@
 use crate::utils::{get_block_trace_from_file, read_env_var};
 use glob::glob;
+use once_cell::sync::Lazy;
 use types::eth::BlockTrace;
 
 pub mod mock_plonk;
@@ -9,6 +10,8 @@ pub use proof::{gen_and_verify_batch_proofs, gen_and_verify_chunk_proofs};
 
 pub const PARAMS_DIR: &str = "./test_params";
 
+pub static TEST_INIT_GEN_PK: Lazy<bool> = Lazy::new(|| read_env_var("TEST_INIT_GEN_PK", false));
+
 pub fn parse_trace_path_from_mode(mode: &str) -> &'static str {
     let trace_path = match mode {
         "empty" => "./tests/traces/bridge/01.json",

prover/src/zkevm/prover.rs

@@ -18,8 +18,13 @@ impl From<common::Prover> for Prover {
 }
 
 impl Prover {
-    pub fn from_params_dir(params_dir: &str) -> Self {
-        common::Prover::from_params_dir(params_dir, &ZKEVM_DEGREES).into()
+    pub fn from_params_dir(params_dir: &str, gen_pk: bool, output_dir: Option<&str>) -> Self {
+        let mut inner = common::Prover::from_params_dir(params_dir, &ZKEVM_DEGREES);
+        if gen_pk {
+            inner.gen_chunk_pk(output_dir).unwrap();
+        }
+
+        inner.into()
     }
 
     pub fn get_vk(&self) -> Option<Vec<u8>> {

prover/tests/aggregation_tests.rs

@@ -1,6 +1,8 @@
 use prover::{
     aggregator::Prover,
-    test_util::{gen_and_verify_batch_proofs, load_block_traces_for_test, PARAMS_DIR},
+    test_util::{
+        gen_and_verify_batch_proofs, load_block_traces_for_test, PARAMS_DIR, TEST_INIT_GEN_PK,
+    },
     utils::{chunk_trace_to_witness_block, init_env_and_log},
     zkevm, ChunkHash, ChunkProof,
 };
@@ -12,10 +14,7 @@ fn test_agg_prove_verify() {
     let output_dir = init_env_and_log("agg_tests");
     log::info!("Initialized ENV and created output-dir {output_dir}");
 
-    let trace_paths = vec![
-        "./tests/traces/erc20/1_transfer.json".to_string(),
-        "./tests/traces/erc20/10_transfer.json".to_string(),
-    ];
+    let trace_paths = vec!["./tests/extra_traces/simple.json".to_string()];
 
     let chunk_hashes_proofs = gen_chunk_hashes_and_proofs(&output_dir, &trace_paths);
     log::info!("Generated chunk hashes and proofs");
@@ -36,7 +35,8 @@ fn gen_chunk_hashes_and_proofs(
     output_dir: &str,
     trace_paths: &[String],
 ) -> Vec<(ChunkHash, ChunkProof)> {
-    let mut zkevm_prover = zkevm::Prover::from_params_dir(PARAMS_DIR);
+    let mut zkevm_prover =
+        zkevm::Prover::from_params_dir(PARAMS_DIR, *TEST_INIT_GEN_PK, Some(&output_dir));
     log::info!("Constructed zkevm prover");
 
     let chunk_traces: Vec<_> = trace_paths

prover/tests/chunk_tests.rs

@@ -1,5 +1,7 @@
 use prover::{
-    test_util::{gen_and_verify_chunk_proofs, load_block_traces_for_test, PARAMS_DIR},
+    test_util::{
+        gen_and_verify_chunk_proofs, load_block_traces_for_test, PARAMS_DIR, TEST_INIT_GEN_PK,
+    },
     utils::{chunk_trace_to_witness_block, init_env_and_log},
     zkevm::Prover,
 };
@@ -16,7 +18,8 @@ fn test_chunk_prove_verify() {
     let witness_block = chunk_trace_to_witness_block(chunk_trace).unwrap();
     log::info!("Got witness block");
 
-    let mut zkevm_prover = Prover::from_params_dir(PARAMS_DIR);
+    let mut zkevm_prover =
+        Prover::from_params_dir(PARAMS_DIR, *TEST_INIT_GEN_PK, Some(&output_dir));
     log::info!("Constructed zkevm prover");
 
     // Load or generate compression wide snark (layer-1).