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/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/common/prover/utils.rs

@@ -1,4 +1,5 @@
 use super::Prover;
+use crate::io::serialize_vk;
 use anyhow::Result;
 use halo2_proofs::{
     halo2curves::bn256::{Bn256, Fr, G1Affine},
@@ -68,4 +69,8 @@ impl Prover {
 
         Ok((&self.params_map[&degree], &self.pk_map[id]))
     }
+
+    pub fn raw_vk(&self, id: &str) -> Option<Vec<u8>> {
+        self.pk_map.get(id).map(|pk| serialize_vk(pk.get_vk()))
+    }
 }

prover/src/config.rs

@@ -76,6 +76,13 @@ impl LayerId {
     }
 }
 
+pub fn asset_file_path(filename: &str) -> String {
+    Path::new(&*ASSETS_DIR)
+        .join(filename)
+        .to_string_lossy()
+        .into_owned()
+}
+
 pub fn layer_config_path(id: &str) -> &str {
     match id {
         "layer1" => &LAYER1_CONFIG_PATH,
@@ -86,13 +93,6 @@ pub fn layer_config_path(id: &str) -> &str {
     }
 }
 
-fn asset_file_path(filename: &str) -> String {
-    Path::new(&*ASSETS_DIR)
-        .join(filename)
-        .to_string_lossy()
-        .into_owned()
-}
-
 fn layer_degree(config_file: &str) -> u32 {
     let f = File::open(config_file).unwrap_or_else(|_| panic!("Failed to open {config_file}"));