Read VKs from files when init.

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, false, None);
+    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/src/aggregator/prover.rs

@@ -1,40 +1,38 @@
 use crate::{
     common,
-    config::{AGG_DEGREES, LAYER3_DEGREE, LAYER4_DEGREE},
-    io::{read_all, serialize_vk},
-    utils::read_env_var,
+    config::{LayerId, AGG_DEGREES, LAYER3_DEGREE, LAYER4_DEGREE},
+    consts::{AGG_VK_FILENAME, CHUNK_PROTOCOL_FILENAME},
+    io::{force_to_read, serialize_vk, 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>,
+    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 = force_to_read(assets_dir, &*CHUNK_PROTOCOL_FILENAME);
 
-        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 vk = try_to_read(assets_dir, &*AGG_VK_FILENAME);
+        if vk.is_none() {
+            log::warn!("{} doesn't exist in {}", *AGG_VK_FILENAME, assets_dir);
         }
-        let chunk_protocol = read_all(&chunk_protocol_path);
 
         Self {
             inner,
             chunk_protocol,
+            vk,
         }
     }
 
@@ -56,8 +54,10 @@ 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
+            .pk(LayerId::Layer4.id())
+            .map(|pk| serialize_vk(pk.get_vk()))
+            .or_else(|| self.vk.clone())
     }
 
     // Return the EVM proof for verification.

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/consts.rs

@@ -0,0 +1,11 @@
+use crate::utils::read_env_var;
+use once_cell::sync::Lazy;
+
+pub static AGG_VK_FILENAME: Lazy<String> =
+    Lazy::new(|| read_env_var("AGG_VK_FILENAME", "agg_vk.vkey".to_string()));
+pub static CHUNK_PROTOCOL_FILENAME: Lazy<String> =
+    Lazy::new(|| read_env_var("CHUNK_PROTOCOL_FILENAME", "chunk.protocol".to_string()));
+pub static CHUNK_VK_FILENAME: Lazy<String> =
+    Lazy::new(|| read_env_var("CHUNK_VK_FILENAME", "chunk_vk.vkey".to_string()));
+pub static DEPLOYMENT_CODE_FILENAME: Lazy<String> =
+    Lazy::new(|| read_env_var("DEPLOYMENT_CODE_FILENAME", "evm_verifier.bin".to_string()));

prover/src/io.rs

@@ -74,6 +74,21 @@ pub fn read_file(folder: &mut PathBuf, filename: &str) -> Vec<u8> {
     buf
 }
 
+pub fn try_to_read(dir: &str, filename: &str) -> Option<Vec<u8>> {
+    let mut path = PathBuf::from(dir);
+    path.push(filename);
+
+    if path.exists() {
+        Some(read_all(&path.to_string_lossy()))
+    } else {
+        None
+    }
+}
+
+pub fn force_to_read(dir: &str, filename: &str) -> Vec<u8> {
+    try_to_read(dir, filename).unwrap_or_else(|| panic!("File {filename} must exist in {dir}"))
+}
+
 pub fn write_file(folder: &mut PathBuf, filename: &str, buf: &[u8]) {
     folder.push(filename);
     let mut fd = std::fs::File::create(folder.as_path()).unwrap();

prover/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod aggregator;
 pub mod common;
 pub mod config;
+pub mod consts;
 mod evm_verifier;
 pub mod inner;
 pub mod io;

prover/src/test_util.rs

@@ -1,6 +1,5 @@
 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;
@@ -10,10 +9,9 @@ pub use proof::{
     gen_and_verify_batch_proofs, gen_and_verify_chunk_proofs, gen_and_verify_normal_and_evm_proofs,
 };
 
+pub const ASSETS_DIR: &str = "./test_assets";
 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

@@ -1,6 +1,11 @@
 use crate::{
-    common, config::ZKEVM_DEGREES, io::serialize_vk, utils::chunk_trace_to_witness_block,
-    zkevm::circuit::normalize_withdraw_proof, ChunkHash, ChunkProof,
+    common,
+    config::{LayerId, ZKEVM_DEGREES},
+    consts::CHUNK_VK_FILENAME,
+    io::{serialize_vk, try_to_read},
+    utils::chunk_trace_to_witness_block,
+    zkevm::circuit::normalize_withdraw_proof,
+    ChunkHash, ChunkProof,
 };
 use anyhow::Result;
 use types::eth::BlockTrace;
@@ -9,27 +14,26 @@ use types::eth::BlockTrace;
 pub struct Prover {
     // Make it public for testing with inner functions (unnecessary for FFI).
     pub inner: common::Prover,
-}
-
-impl From<common::Prover> for Prover {
-    fn from(inner: common::Prover) -> Self {
-        Self { inner }
-    }
+    vk: Option<Vec<u8>>,
 }
 
 impl Prover {
-    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();
+    pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
+        let inner = common::Prover::from_params_dir(params_dir, &ZKEVM_DEGREES);
+
+        let vk = try_to_read(assets_dir, &*CHUNK_VK_FILENAME);
+        if vk.is_none() {
+            log::warn!("{} doesn't exist in {}", *CHUNK_VK_FILENAME, assets_dir);
         }
 
-        inner.into()
+        Self { inner, vk }
     }
 
     pub fn get_vk(&self) -> Option<Vec<u8>> {
-        // TODO: replace `layer2` string with an enum value.
-        self.inner.pk("layer2").map(|pk| serialize_vk(pk.get_vk()))
+        self.inner
+            .pk(LayerId::Layer2.id())
+            .map(|pk| serialize_vk(pk.get_vk()))
+            .or_else(|| self.vk.clone())
     }
 
     pub fn gen_chunk_proof(

prover/src/zkevm/verifier.rs

@@ -1,8 +1,8 @@
 use crate::{
     common,
     config::{LAYER2_CONFIG_PATH, LAYER2_DEGREE},
-    io::read_all,
-    utils::read_env_var,
+    consts::CHUNK_VK_FILENAME,
+    io::force_to_read,
     ChunkProof,
 };
 use aggregator::CompressionCircuit;
@@ -11,11 +11,7 @@ use halo2_proofs::{
     plonk::VerifyingKey,
     poly::kzg::commitment::ParamsKZG,
 };
-use once_cell::sync::Lazy;
-use std::{env, path::Path};
-
-static CHUNK_VK_FILENAME: Lazy<String> =
-    Lazy::new(|| read_env_var("CHUNK_VK_FILENAME", "chunk_vk.vkey".to_string()));
+use std::env;
 
 #[derive(Debug)]
 pub struct Verifier {
@@ -35,12 +31,7 @@ impl Verifier {
     }
 
     pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
-        let vk_path = format!("{assets_dir}/{}", *CHUNK_VK_FILENAME);
-        if !Path::new(&vk_path).exists() {
-            panic!("File {vk_path} must exist");
-        }
-
-        let raw_vk = read_all(&vk_path);
+        let raw_vk = force_to_read(assets_dir, &*CHUNK_VK_FILENAME);
 
         env::set_var("COMPRESSION_CONFIG", &*LAYER2_CONFIG_PATH);
         common::Verifier::from_params_dir(params_dir, *LAYER2_DEGREE, &raw_vk).into()

prover/tests/aggregation_tests.rs

@@ -1,8 +1,6 @@
 use prover::{
     aggregator::Prover,
-    test_util::{
-        gen_and_verify_batch_proofs, load_block_traces_for_test, PARAMS_DIR, TEST_INIT_GEN_PK,
-    },
+    test_util::{gen_and_verify_batch_proofs, load_block_traces_for_test, ASSETS_DIR, PARAMS_DIR},
     utils::{chunk_trace_to_witness_block, init_env_and_log},
     zkevm, ChunkHash, ChunkProof,
 };
@@ -35,8 +33,7 @@ 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, *TEST_INIT_GEN_PK, Some(&output_dir));
+    let mut zkevm_prover = zkevm::Prover::from_dirs(PARAMS_DIR, ASSETS_DIR);
     log::info!("Constructed zkevm prover");
 
     let chunk_traces: Vec<_> = trace_paths

prover/tests/chunk_tests.rs

@@ -1,7 +1,5 @@
 use prover::{
-    test_util::{
-        gen_and_verify_chunk_proofs, load_block_traces_for_test, PARAMS_DIR, TEST_INIT_GEN_PK,
-    },
+    test_util::{gen_and_verify_chunk_proofs, load_block_traces_for_test, ASSETS_DIR, PARAMS_DIR},
     utils::{chunk_trace_to_witness_block, init_env_and_log},
     zkevm::Prover,
 };
@@ -20,15 +18,9 @@ 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, *TEST_INIT_GEN_PK, Some(&output_dir));
+    let mut zkevm_prover = Prover::from_dirs(PARAMS_DIR, ASSETS_DIR);
     log::info!("Constructed zkevm prover");
 
-    if *TEST_INIT_GEN_PK {
-        assert!(!zkevm_prover.get_vk().unwrap().is_empty());
-        log::info!("Chunk VK is available after init");
-    }
-
     // Load or generate compression wide snark (layer-1).
     let layer1_snark = zkevm_prover
         .inner