feat: return VKs from chunk and batch provers

bin/src/zkevm_prove.rs

@@ -8,9 +8,12 @@ use std::{env, fs, path::PathBuf, time::Instant};
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
 struct Args {
-    /// Get params and write into file.
+    /// Get params dir path.
     #[clap(short, long = "params", default_value = "test_params")]
     params_path: String,
+    /// Get asserts dir path.
+    #[clap(short, long = "assets", default_value = "test_assets")]
+    assets_path: String,
     /// Get BlockTrace from file or dir.
     #[clap(
         short,
@@ -27,7 +30,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_dirs(&args.params_path, &args.assets_path);
 
     let mut traces = Vec::new();
     let trace_path = PathBuf::from(&args.trace_path);

prover/Cargo.toml

@@ -19,6 +19,7 @@ snark-verifier = { git = "https://github.com/scroll-tech/snark-verifier", tag =
 snark-verifier-sdk = { git = "https://github.com/scroll-tech/snark-verifier", tag = "v0.1.3", default-features = false, features = ["loader_halo2", "loader_evm", "halo2-pse"] }
 
 anyhow = "1.0"
+base64 = "0.13.0"
 blake2 = "0.10.3"
 chrono = "0.4.19"
 dotenv = "0.15.0"

prover/src/aggregator/prover.rs

@@ -1,40 +1,42 @@
 use crate::{
     common,
-    config::{AGG_DEGREES, LAYER3_DEGREE, LAYER4_DEGREE},
-    io::{read_all, serialize_vk},
-    utils::read_env_var,
+    config::{LayerId, AGG_DEGREES},
+    consts::{AGG_VK_FILENAME, CHUNK_PROTOCOL_FILENAME},
+    io::{force_to_read, try_to_read},
     BatchProof, ChunkProof,
 };
 use aggregator::{ChunkHash, MAX_AGG_SNARKS};
 use anyhow::{bail, Result};
-use once_cell::sync::Lazy;
 use sha2::{Digest, Sha256};
 use snark_verifier_sdk::Snark;
-use std::{iter::repeat, path::Path};
-
-static CHUNK_PROTOCOL_FILENAME: Lazy<String> =
-    Lazy::new(|| read_env_var("CHUNK_PROTOCOL_FILENAME", "chunk.protocol".to_string()));
+use std::iter::repeat;
 
 #[derive(Debug)]
 pub struct Prover {
     // Make it public for testing with inner functions (unnecessary for FFI).
     pub inner: common::Prover,
     pub chunk_protocol: Vec<u8>,
+    raw_vk: Option<Vec<u8>>,
 }
 
 impl Prover {
     pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
         let inner = common::Prover::from_params_dir(params_dir, &AGG_DEGREES);
-
-        let chunk_protocol_path = format!("{assets_dir}/{}", *CHUNK_PROTOCOL_FILENAME);
-        if !Path::new(&chunk_protocol_path).exists() {
-            panic!("File {chunk_protocol_path} must exist");
+        let chunk_protocol = force_to_read(assets_dir, &CHUNK_PROTOCOL_FILENAME);
+
+        let raw_vk = try_to_read(assets_dir, &AGG_VK_FILENAME);
+        if raw_vk.is_none() {
+            log::warn!(
+                "agg-prover: {} doesn't exist in {}",
+                *AGG_VK_FILENAME,
+                assets_dir
+            );
         }
-        let chunk_protocol = read_all(&chunk_protocol_path);
 
         Self {
             inner,
             chunk_protocol,
+            raw_vk,
         }
     }
 
@@ -56,8 +58,9 @@ impl Prover {
     }
 
     pub fn get_vk(&self) -> Option<Vec<u8>> {
-        // TODO: replace `layer4` string with an enum value.
-        self.inner.pk("layer4").map(|pk| serialize_vk(pk.get_vk()))
+        self.inner
+            .raw_vk(LayerId::Layer4.id())
+            .or_else(|| self.raw_vk.clone())
     }
 
     // Return the EVM proof for verification.
@@ -86,14 +89,16 @@ impl Prover {
         // Load or generate final compression thin EVM proof (layer-4).
         let evm_proof = self.inner.load_or_gen_comp_evm_proof(
             &name,
-            "layer4",
+            LayerId::Layer4.id(),
             true,
-            *LAYER4_DEGREE,
+            LayerId::Layer4.degree(),
             layer3_snark,
             output_dir,
         )?;
         log::info!("Got final compression thin EVM proof (layer-4): {name}");
 
+        self.check_and_clear_raw_vk();
+
         let batch_proof = BatchProof::from(evm_proof.proof);
         if let Some(output_dir) = output_dir {
             batch_proof.dump(output_dir, "agg")?;
@@ -135,8 +140,8 @@ impl Prover {
         // Load or generate aggregation snark (layer-3).
         let layer3_snark = self.inner.load_or_gen_agg_snark(
             name,
-            "layer3",
-            *LAYER3_DEGREE,
+            LayerId::Layer3.id(),
+            LayerId::Layer3.degree(),
             &chunk_hashes,
             &layer2_snarks,
             output_dir,
@@ -145,4 +150,20 @@ impl Prover {
 
         Ok(layer3_snark)
     }
+
+    fn check_and_clear_raw_vk(&mut self) {
+        if self.raw_vk.is_some() {
+            // Check VK is same with the init one, and take (clear) init VK.
+            let gen_vk = self.inner.raw_vk(LayerId::Layer4.id()).unwrap_or_default();
+            let init_vk = self.raw_vk.take().unwrap_or_default();
+
+            if gen_vk != init_vk {
+                log::error!(
+                    "agg-prover: generated VK is different with init one - gen_vk = {}, init_vk = {}",
+                    base64::encode(gen_vk),
+                    base64::encode(init_vk),
+                );
+            }
+        }
+    }
 }

prover/src/aggregator/verifier.rs

@@ -1,8 +1,8 @@
 use crate::{
     common,
     config::{LAYER4_CONFIG_PATH, LAYER4_DEGREE},
-    io::read_all,
-    utils::read_env_var,
+    consts::{AGG_VK_FILENAME, DEPLOYMENT_CODE_FILENAME},
+    io::force_to_read,
     BatchProof,
 };
 use aggregator::CompressionCircuit;
@@ -11,14 +11,8 @@ use halo2_proofs::{
     plonk::VerifyingKey,
     poly::kzg::commitment::ParamsKZG,
 };
-use once_cell::sync::Lazy;
 use snark_verifier_sdk::verify_evm_calldata;
-use std::{env, path::Path};
-
-static AGG_VK_FILENAME: Lazy<String> =
-    Lazy::new(|| read_env_var("AGG_VK_FILENAME", "agg_vk.vkey".to_string()));
-static DEPLOYMENT_CODE_FILENAME: Lazy<String> =
-    Lazy::new(|| read_env_var("DEPLOYMENT_CODE_FILENAME", "evm_verifier.bin".to_string()));
+use std::env;
 
 #[derive(Debug)]
 pub struct Verifier {
@@ -42,14 +36,8 @@ impl Verifier {
     }
 
     pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
-        let vk_path = format!("{assets_dir}/{}", *AGG_VK_FILENAME);
-        let deployment_code_path = format!("{assets_dir}/{}", *DEPLOYMENT_CODE_FILENAME);
-        if !(Path::new(&vk_path).exists() && Path::new(&deployment_code_path).exists()) {
-            panic!("File {vk_path} and {deployment_code_path} must exist");
-        }
-
-        let raw_vk = read_all(&vk_path);
-        let deployment_code = read_all(&deployment_code_path);
+        let raw_vk = force_to_read(assets_dir, &AGG_VK_FILENAME);
+        let deployment_code = force_to_read(assets_dir, &DEPLOYMENT_CODE_FILENAME);
 
         env::set_var("COMPRESSION_CONFIG", &*LAYER4_CONFIG_PATH);
         let inner = common::Verifier::from_params_dir(params_dir, *LAYER4_DEGREE, &raw_vk);

prover/src/common/prover/chunk.rs

@@ -1,15 +1,37 @@
 use super::Prover;
 use crate::{
-    config::{LAYER1_DEGREE, LAYER2_DEGREE},
-    utils::gen_rng,
+    config::{asset_file_path, LayerId, LAYER1_DEGREE, LAYER2_DEGREE},
+    utils::{chunk_trace_to_witness_block, gen_rng, get_block_trace_from_file, read_env_var},
 };
 use aggregator::extract_proof_and_instances_with_pairing_check;
 use anyhow::{anyhow, Result};
 use halo2_proofs::halo2curves::bn256::Fr;
+use once_cell::sync::Lazy;
 use snark_verifier_sdk::Snark;
+use std::path::Path;
 use zkevm_circuits::evm_circuit::witness::Block;
 
+static SIMPLE_TRACE_FILENAME: Lazy<String> =
+    Lazy::new(|| read_env_var("SIMPLE_TRACE_FILENAME", "simple_trace.json".to_string()));
+
 impl Prover {
+    pub fn gen_chunk_pk(&mut self, output_dir: Option<&str>) -> Result<()> {
+        if self.pk(LayerId::Layer2.id()).is_some() {
+            return Ok(());
+        }
+
+        let simple_trace_path = asset_file_path(&SIMPLE_TRACE_FILENAME);
+        if !Path::new(&simple_trace_path).exists() {
+            panic!("File {simple_trace_path} must exist");
+        }
+        let simple_trace = get_block_trace_from_file(simple_trace_path);
+        let witness_block = chunk_trace_to_witness_block(vec![simple_trace])?;
+        let layer1_snark =
+            self.load_or_gen_last_chunk_snark("empty", &witness_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,