(v0.13.1) Move params out of prover/verifier & Additional sanity check (#1405)

* refactor(prover): move params out of prover/verifier

* clean

* remove redundant function 'degrees'

* add sanity check (decoded blob == batch bytes)

---------

Co-authored-by: Rohit Narurkar <[email protected]>

Author: lispc

aggregator/src/eip4844.rs

@@ -54,3 +54,36 @@ pub fn get_blob_bytes(batch_bytes: &[u8]) -> Vec<u8> {
 
     blob_bytes
 }
+
+/// Given the blob's bytes, take into account the first byte, i.e. enable_encoding? and either spit
+/// out the raw bytes or zstd decode them.
+pub fn decode_blob(blob_bytes: &[u8]) -> std::io::Result<Vec<u8>> {
+    let enable_encoding = blob_bytes[0].eq(&1);
+
+    // If not encoded, spit out the rest of the bytes, as it is.
+    if !enable_encoding {
+        return Ok(blob_bytes[1..].to_vec());
+    }
+
+    // The bytes following the first byte represent the zstd-encoded bytes.
+    let mut encoded_bytes = blob_bytes[1..].to_vec();
+    let mut encoded_len = encoded_bytes.len();
+    let mut decoded_bytes = Vec::with_capacity(5 * 4096 * 32);
+    loop {
+        let mut decoder = zstd_encoder::zstd::stream::read::Decoder::new(encoded_bytes.as_slice())?;
+        decoder.include_magicbytes(false)?;
+        decoder.window_log_max(30)?;
+
+        decoded_bytes.clear();
+
+        if std::io::copy(&mut decoder, &mut decoded_bytes).is_ok() {
+            break;
+        }
+
+        // The error above means we need to truncate the suffix 0-byte.
+        encoded_len -= 1;
+        encoded_bytes.truncate(encoded_len);
+    }
+
+    Ok(decoded_bytes)
+}

prover/src/aggregator/prover.rs

@@ -1,14 +1,17 @@
-use std::{env, iter::repeat};
+use std::{collections::BTreeMap, env, iter::repeat};
 
-use aggregator::{BatchHash, BatchHeader, ChunkInfo, MAX_AGG_SNARKS};
+use aggregator::{
+    eip4844::decode_blob, BatchData, BatchHash, BatchHeader, ChunkInfo, MAX_AGG_SNARKS,
+};
 use anyhow::{bail, Result};
 use eth_types::H256;
+use halo2_proofs::{halo2curves::bn256::Bn256, poly::kzg::commitment::ParamsKZG};
 use sha2::{Digest, Sha256};
 use snark_verifier_sdk::Snark;
 
 use crate::{
     common,
-    config::{LayerId, AGG_DEGREES},
+    config::LayerId,
     consts::{BATCH_KECCAK_ROW, BATCH_VK_FILENAME, BUNDLE_VK_FILENAME, CHUNK_PROTOCOL_FILENAME},
     io::{force_to_read, try_to_read},
     proof::BundleProof,
@@ -17,20 +20,23 @@ use crate::{
 };
 
 #[derive(Debug)]
-pub struct Prover {
+pub struct Prover<'params> {
     // Make it public for testing with inner functions (unnecessary for FFI).
-    pub prover_impl: common::Prover,
+    pub prover_impl: common::Prover<'params>,
     pub chunk_protocol: Vec<u8>,
     raw_vk_batch: Option<Vec<u8>>,
     raw_vk_bundle: Option<Vec<u8>>,
 }
 
-impl Prover {
-    pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
+impl<'params> Prover<'params> {
+    pub fn from_params_and_assets(
+        params_map: &'params BTreeMap<u32, ParamsKZG<Bn256>>,
+        assets_dir: &str,
+    ) -> Self {
         log::debug!("set env KECCAK_ROWS={}", BATCH_KECCAK_ROW.to_string());
         env::set_var("KECCAK_ROWS", BATCH_KECCAK_ROW.to_string());
 
-        let prover_impl = common::Prover::from_params_dir(params_dir, &AGG_DEGREES);
+        let prover_impl = common::Prover::from_params_map(params_map);
         let chunk_protocol = force_to_read(assets_dir, &CHUNK_PROTOCOL_FILENAME);
 
         let raw_vk_batch = try_to_read(assets_dir, &BATCH_VK_FILENAME);
@@ -208,6 +214,16 @@ impl Prover {
         let batch_hash = batch_header.batch_hash();
         let batch_info: BatchHash<N_SNARKS> =
             BatchHash::construct(&chunk_hashes, batch_header, &batch.blob_bytes);
+        let batch_data: BatchData<N_SNARKS> = BatchData::from(&batch_info);
+
+        // sanity check:
+        // - conditionally decoded blob should match batch data.
+        let batch_bytes = batch_data.get_batch_data_bytes();
+        let decoded_blob_bytes = decode_blob(&batch.blob_bytes)?;
+        assert_eq!(
+            batch_bytes, decoded_blob_bytes,
+            "BatchProvingTask(sanity) mismatch batch bytes and decoded blob bytes",
+        );
 
         let layer3_snark = self.prover_impl.load_or_gen_agg_snark(
             name,

prover/src/aggregator/verifier.rs

@@ -13,18 +13,18 @@ use halo2_proofs::{
 };
 use snark_verifier_sdk::verify_evm_calldata;
 use snark_verifier_sdk::Snark;
-use std::env;
+use std::{collections::BTreeMap, env};
 
 #[derive(Debug)]
-pub struct Verifier {
+pub struct Verifier<'params> {
     // Make it public for testing with inner functions (unnecessary for FFI).
-    pub inner: common::Verifier<CompressionCircuit>,
+    pub inner: common::Verifier<'params, CompressionCircuit>,
     deployment_code: Option<Vec<u8>>,
 }
 
-impl Verifier {
+impl<'params> Verifier<'params> {
     pub fn new(
-        params: ParamsKZG<Bn256>,
+        params: &'params ParamsKZG<Bn256>,
         vk: VerifyingKey<G1Affine>,
         deployment_code: Vec<u8>,
     ) -> Self {
@@ -36,12 +36,16 @@ impl Verifier {
         }
     }
 
-    pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
+    pub fn from_params_and_assets(
+        params_map: &'params BTreeMap<u32, ParamsKZG<Bn256>>,
+        assets_dir: &str,
+    ) -> Self {
         let raw_vk = force_to_read(assets_dir, &batch_vk_filename());
         let deployment_code = try_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);
+        let params = params_map.get(&*LAYER4_DEGREE).expect("should be loaded");
+        let inner = common::Verifier::from_params(params, &raw_vk);
 
         Self {
             inner,

prover/src/common/prover.rs

@@ -16,22 +16,22 @@ mod recursion;
 mod utils;
 
 #[derive(Debug)]
-pub struct Prover {
+pub struct Prover<'params> {
     // degree -> params (use BTreeMap to find proper degree for params downsize)
-    params_map: BTreeMap<u32, ParamsKZG<Bn256>>,
+    pub params_map: &'params BTreeMap<u32, ParamsKZG<Bn256>>,
     // Cached id -> pk
     pk_map: HashMap<String, ProvingKey<G1Affine>>,
 }
 
-impl Prover {
-    pub fn from_params(params_map: BTreeMap<u32, ParamsKZG<Bn256>>) -> Self {
+impl<'params> Prover<'params> {
+    pub fn from_params_map(params_map: &'params BTreeMap<u32, ParamsKZG<Bn256>>) -> Self {
         Self {
             params_map,
             pk_map: HashMap::new(),
         }
     }
 
-    pub fn from_params_dir(params_dir: &str, degrees: &[u32]) -> Self {
+    pub fn load_params_map(params_dir: &str, degrees: &[u32]) -> BTreeMap<u32, ParamsKZG<Bn256>> {
         let degrees = BTreeSet::from_iter(degrees);
         let max_degree = **degrees.last().unwrap();
 
@@ -63,9 +63,6 @@ impl Prover {
             params_map.insert(*d, params);
         }
 
-        Self {
-            params_map,
-            pk_map: HashMap::new(),
-        }
+        params_map
     }
 }

prover/src/common/prover/aggregation.rs

@@ -10,7 +10,7 @@ use rand::Rng;
 use snark_verifier_sdk::Snark;
 use std::env;
 
-impl Prover {
+impl<'params> Prover<'params> {
     pub fn gen_agg_snark<const N_SNARKS: usize>(
         &mut self,
         id: &str,

prover/src/common/prover/chunk.rs

@@ -5,7 +5,7 @@ use anyhow::{anyhow, Result};
 use snark_verifier_sdk::Snark;
 use zkevm_circuits::evm_circuit::witness::Block;
 
-impl Prover {
+impl<'params> Prover<'params> {
     pub fn load_or_gen_final_chunk_snark(
         &mut self,
         name: &str,

prover/src/common/prover/compression.rs

@@ -10,7 +10,7 @@ use rand::Rng;
 use snark_verifier_sdk::Snark;
 use std::env;
 
-impl Prover {
+impl<'params> Prover<'params> {
     pub fn gen_comp_snark(
         &mut self,
         id: &str,

prover/src/common/prover/evm.rs

@@ -11,7 +11,7 @@ use rand::Rng;
 use snark_verifier_sdk::{gen_evm_proof_shplonk, CircuitExt, Snark};
 use std::env;
 
-impl Prover {
+impl<'params> Prover<'params> {
     pub fn load_or_gen_comp_evm_proof(
         &mut self,
         name: &str,

prover/src/common/prover/inner.rs

@@ -10,7 +10,7 @@ use rand::Rng;
 use snark_verifier_sdk::{gen_snark_shplonk, Snark};
 use zkevm_circuits::evm_circuit::witness::Block;
 
-impl Prover {
+impl<'params> Prover<'params> {
     pub fn gen_inner_snark<C: TargetCircuit>(
         &mut self,
         id: &str,

prover/src/common/prover/mock.rs

@@ -6,7 +6,7 @@ use std::sync::LazyLock;
 
 pub static MOCK_PROVE: LazyLock<bool> = LazyLock::new(|| read_env_var("MOCK_PROVE", false));
 
-impl Prover {
+impl<'params> Prover<'params> {
     pub fn assert_if_mock_prover<C: CircuitExt<Fr>>(id: &str, degree: u32, circuit: &C) {
         if !*MOCK_PROVE {
             return;