From 5b94df914877aa7f20f7c37fabf80d73f1c8cf2c Mon Sep 17 00:00:00 2001
From: Steven Gu <asongala@163.com>
Date: Thu, 7 Sep 2023 08:26:35 +0800
Subject: [PATCH] Set `KECCAK_ROW=50` when init batch-prover.

---
 prover/src/aggregator/prover.rs | 6 ++++--
 prover/src/consts.rs            | 1 +
 prover/tests/batch_tests.rs     | 2 --
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/prover/src/aggregator/prover.rs b/prover/src/aggregator/prover.rs
index f0917d175..7477e5aa2 100644
--- a/prover/src/aggregator/prover.rs
+++ b/prover/src/aggregator/prover.rs
@@ -1,7 +1,7 @@
 use crate::{
     common,
     config::{LayerId, AGG_DEGREES},
-    consts::{AGG_VK_FILENAME, CHUNK_PROTOCOL_FILENAME},
+    consts::{AGG_KECCAK_ROW, AGG_VK_FILENAME, CHUNK_PROTOCOL_FILENAME},
     io::{force_to_read, try_to_read},
     BatchProof, ChunkProof,
 };
@@ -9,7 +9,7 @@ use aggregator::{ChunkHash, MAX_AGG_SNARKS};
 use anyhow::{bail, Result};
 use sha2::{Digest, Sha256};
 use snark_verifier_sdk::Snark;
-use std::iter::repeat;
+use std::{env, iter::repeat};
 
 #[derive(Debug)]
 pub struct Prover {
@@ -21,6 +21,8 @@ pub struct Prover {
 
 impl Prover {
     pub fn from_dirs(params_dir: &str, assets_dir: &str) -> Self {
+        env::set_var("KECCAK_ROW", AGG_KECCAK_ROW.to_string());
+
         let inner = common::Prover::from_params_dir(params_dir, &AGG_DEGREES);
         let chunk_protocol = force_to_read(assets_dir, &CHUNK_PROTOCOL_FILENAME);
 
diff --git a/prover/src/consts.rs b/prover/src/consts.rs
index 2300606ee..91390301d 100644
--- a/prover/src/consts.rs
+++ b/prover/src/consts.rs
@@ -1,6 +1,7 @@
 use crate::utils::read_env_var;
 use once_cell::sync::Lazy;
 
+pub static AGG_KECCAK_ROW: Lazy<usize> = Lazy::new(|| read_env_var("AGG_KECCAK_ROW", 50));
 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> =
diff --git a/prover/tests/batch_tests.rs b/prover/tests/batch_tests.rs
index 1dd8acdbb..b79417c51 100644
--- a/prover/tests/batch_tests.rs
+++ b/prover/tests/batch_tests.rs
@@ -11,8 +11,6 @@ use std::env;
 #[cfg(feature = "prove_verify")]
 #[test]
 fn test_batch_prove_verify() {
-    env::set_var("KECCAK_ROW", "50");
-
     let output_dir = init_env_and_log("batch_tests");
     log::info!("Initialized ENV and created output-dir {output_dir}");