refactor: remove keccak constants from aggregator

aggregator/src/aggregation/config.rs

@@ -83,18 +83,10 @@ impl AggregationConfig {
             params.degree as usize,
         );
 
-        // The current code base is hardcoded for KeccakCircuit configured
-        // with 300 rows and 87 columns per hash call.
         let columns = keccak_circuit_config.cell_manager.columns();
 
-        assert_eq!(
-            columns.len(),
-            87,
-            "cell manager configuration does not match the hard coded setup"
-        );
-
         // enabling equality for preimage column
-        meta.enable_equality(columns[6].advice);
+        meta.enable_equality(columns[keccak_circuit_config.preimage_column].advice);
         // enable equality for the digest column
         meta.enable_equality(columns.last().unwrap().advice);
         // enable equality for the data RLC column

aggregator/src/constants.rs

@@ -11,15 +11,6 @@ pub(crate) const DIGEST_LEN: usize = 32;
 /// Input length per round
 pub(crate) const INPUT_LEN_PER_ROUND: usize = 136;
 
-// Each round requires (NUM_ROUNDS+1) * DEFAULT_KECCAK_ROWS = 300 rows.
-// This library is hard coded for this parameter.
-// Modifying the following parameters may result into bugs.
-// Adopted from keccak circuit
-pub(crate) const DEFAULT_KECCAK_ROWS: usize = 12;
-// Adopted from keccak circuit
-pub(crate) const NUM_ROUNDS: usize = 24;
-pub(crate) const ROWS_PER_ROUND: usize = (NUM_ROUNDS + 1) * DEFAULT_KECCAK_ROWS;
-
 // TODO(ZZ): update to the right degree
 #[allow(dead_code)]
 pub(crate) const LOG_DEGREE: u32 = 19;
@@ -58,9 +49,3 @@ pub(crate) const BITS: usize = 88;
 /// will be padded.
 // TODO: update me(?)
 pub const MAX_AGG_SNARKS: usize = 10;
-
-/// The number of keccak rounds is the sum of
-/// - batch public input hash: 2 rounds
-/// - chunk's public input hash: 2 * MAX_AGG_SNARKS
-/// - batch data hash: (32 * MAX_AGG_SNARKS)/136 = 3
-pub(crate) const MAX_KECCAK_ROUNDS: usize = 2 * MAX_AGG_SNARKS + 5;

aggregator/src/core.rs

@@ -23,18 +23,18 @@ use snark_verifier_sdk::{
     Snark,
 };
 use zkevm_circuits::{
-    keccak_circuit::{keccak_packed_multi::multi_keccak, KeccakCircuitConfig},
+    keccak_circuit::{
+        keccak_packed_multi::{self, multi_keccak},
+        KeccakCircuit, KeccakCircuitConfig,
+    },
     table::LookupTable,
     util::Challenges,
 };
 
 use crate::{
-    constants::{
-        CHAIN_ID_LEN, DIGEST_LEN, INPUT_LEN_PER_ROUND, LOG_DEGREE, MAX_AGG_SNARKS,
-        MAX_KECCAK_ROUNDS, ROWS_PER_ROUND,
-    },
+    constants::{CHAIN_ID_LEN, DIGEST_LEN, INPUT_LEN_PER_ROUND, LOG_DEGREE, MAX_AGG_SNARKS},
     util::{
-        assert_conditional_equal, assert_equal, assert_exist, get_indices, keccak_round_capacity,
+        assert_conditional_equal, assert_equal, assert_exist, get_indices, get_max_keccak_updates,
         parse_hash_digest_cells, parse_hash_preimage_cells, parse_pi_hash_rlc_cells,
     },
     AggregationConfig, RlcConfig, CHUNK_DATA_HASH_INDEX, POST_STATE_ROOT_INDEX,
@@ -185,7 +185,9 @@ pub(crate) fn extract_hash_cells(
     preimages: &[Vec<u8>],
 ) -> Result<ExtractedHashCells, Error> {
     let mut is_first_time = true;
-    let num_rows = 1 << LOG_DEGREE;
+    let keccak_capacity = KeccakCircuit::<Fr>::capacity_for_row(1 << LOG_DEGREE);
+    let max_keccak_updates = get_max_keccak_updates(MAX_AGG_SNARKS);
+    let keccak_f_rows = keccak_packed_multi::get_num_rows_per_update();
 
     let timer = start_timer!(|| ("multi keccak").to_string());
     // preimages consists of the following parts
@@ -202,7 +204,7 @@ pub(crate) fn extract_hash_cells(
     // (3) batchDataHash preimage =
     //      (chunk[0].dataHash || ... || chunk[k-1].dataHash)
     // each part of the preimage is mapped to image by Keccak256
-    let witness = multi_keccak(preimages, challenges, keccak_round_capacity(num_rows))
+    let witness = multi_keccak(preimages, challenges, keccak_capacity)
         .map_err(|e| Error::AssertionFailure(format!("multi keccak assignment failed: {e:?}")))?;
     end_timer!(timer);
 
@@ -241,17 +243,18 @@ pub(crate) fn extract_hash_cells(
                     let row = keccak_config.set_row(&mut region, offset, keccak_row)?;
 
                     if cur_preimage_index.is_some() && *cur_preimage_index.unwrap() == offset {
-                        // 10-th column is Keccak input in Keccak circuit
-                        hash_input_cells.push(row[10].clone());
+                        hash_input_cells.push(
+                            row[keccak_packed_multi::get_input_bytes_col_cell_manager() + 4]
+                                .clone(),
+                        );
                         cur_preimage_index = preimage_indices_iter.next();
                     }
                     if cur_digest_index.is_some() && *cur_digest_index.unwrap() == offset {
                         // last column is Keccak output in Keccak circuit
                         hash_output_cells.push(row.last().unwrap().clone()); // sage unwrap
                         cur_digest_index = digest_indices_iter.next();
                     }
-                    if offset % ROWS_PER_ROUND == 0 && offset / ROWS_PER_ROUND <= MAX_KECCAK_ROUNDS
-                    {
+                    if offset % keccak_f_rows == 0 && offset / keccak_f_rows <= max_keccak_updates {
                         // first column is is_final
                         is_final_cells.push(row[0].clone());
                         // second column is data rlc
@@ -268,7 +271,7 @@ pub(crate) fn extract_hash_cells(
                 // sanity
                 assert_eq!(
                     hash_input_cells.len(),
-                    MAX_KECCAK_ROUNDS * INPUT_LEN_PER_ROUND
+                    max_keccak_updates * INPUT_LEN_PER_ROUND
                 );
                 assert_eq!(hash_output_cells.len(), (MAX_AGG_SNARKS + 4) * DIGEST_LEN);
 

aggregator/src/tests/rlc/dynamic_hashes.rs

@@ -9,17 +9,14 @@ use snark_verifier::loader::halo2::halo2_ecc::halo2_base::utils::fs::gen_srs;
 use snark_verifier_sdk::{gen_pk, gen_snark_shplonk, verify_snark_shplonk, CircuitExt};
 use zkevm_circuits::{
     keccak_circuit::{
-        keccak_packed_multi::multi_keccak, KeccakCircuitConfig, KeccakCircuitConfigArgs,
+        keccak_packed_multi::{self, multi_keccak},
+        KeccakCircuit, KeccakCircuitConfig, KeccakCircuitConfigArgs,
     },
     table::{KeccakTable, LookupTable},
     util::{Challenges, SubCircuitConfig},
 };
 
-use crate::{
-    aggregation::RlcConfig,
-    constants::{LOG_DEGREE, ROWS_PER_ROUND},
-    util::keccak_round_capacity,
-};
+use crate::{aggregation::RlcConfig, constants::LOG_DEGREE};
 
 #[derive(Default, Debug, Clone)]
 struct DynamicHashCircuit {
@@ -78,6 +75,7 @@ impl Circuit<Fr> for DynamicHashCircuit {
         mut layouter: impl Layouter<Fr>,
     ) -> Result<(), Error> {
         let (config, challenges) = config;
+        let keccak_f_rows = keccak_packed_multi::get_num_rows_per_update();
 
         config
             .keccak_circuit_config
@@ -95,7 +93,7 @@ impl Circuit<Fr> for DynamicHashCircuit {
         let witness = multi_keccak(
             &[hash_preimage.clone()],
             challenge,
-            keccak_round_capacity(1 << LOG_DEGREE),
+            KeccakCircuit::<Fr>::capacity_for_row(1 << LOG_DEGREE),
         )
         .unwrap();
 
@@ -113,7 +111,7 @@ impl Circuit<Fr> for DynamicHashCircuit {
                         config
                             .keccak_circuit_config
                             .set_row(&mut region, offset, keccak_row)?;
-                    if offset % ROWS_PER_ROUND == 0 && data_rlc_cells.len() < 4 {
+                    if offset % keccak_f_rows == 0 && data_rlc_cells.len() < 4 {
                         // second element is data rlc
                         data_rlc_cells.push(row[1].clone());
                     }

aggregator/src/util.rs

@@ -1,3 +1,7 @@
+use crate::{
+    aggregation::RlcConfig,
+    constants::{DIGEST_LEN, INPUT_LEN_PER_ROUND, MAX_AGG_SNARKS},
+};
 use eth_types::Field;
 use halo2_proofs::{
     circuit::{AssignedCell, Region},
@@ -9,29 +13,28 @@ use snark_verifier::loader::halo2::halo2_ecc::halo2_base::{
     gates::{flex_gate::FlexGateConfig, GateInstructions},
     AssignedValue, Context,
 };
-
-use crate::{
-    aggregation::RlcConfig,
-    constants::{DIGEST_LEN, INPUT_LEN_PER_ROUND, MAX_AGG_SNARKS},
-    DEFAULT_KECCAK_ROWS, NUM_ROUNDS,
+use zkevm_circuits::keccak_circuit::keccak_packed_multi::{
+    get_num_rows_per_round, get_num_rows_per_update,
 };
 
-use std::env::var;
+// Calculates the maximum keccak updates (1 absorb, or 1 f-box invoke)
+// needed for the number of snarks
+pub(crate) fn get_max_keccak_updates(max_snarks: usize) -> usize {
+    let pi_rounds = 2;
+    let chunk_hash_rounds = 2 * max_snarks;
+    let data_hash_rounds = get_data_hash_keccak_updates(max_snarks);
 
-pub(crate) fn keccak_round_capacity(num_rows: usize) -> Option<usize> {
-    if num_rows > 0 {
-        // Subtract two for unusable rows
-        Some(num_rows / ((NUM_ROUNDS + 1) * get_num_rows_per_round()) - 2)
-    } else {
-        None
-    }
+    pi_rounds + chunk_hash_rounds + data_hash_rounds
 }
+pub(crate) fn get_data_hash_keccak_updates(max_snarks: usize) -> usize {
+    let data_hash_rounds = (32 * max_snarks) / INPUT_LEN_PER_ROUND;
+    let padding_round = if data_hash_rounds * INPUT_LEN_PER_ROUND < 32 * max_snarks {
+        1
+    } else {
+        0
+    };
 
-pub(crate) fn get_num_rows_per_round() -> usize {
-    var("KECCAK_ROWS")
-        .unwrap_or_else(|_| format!("{DEFAULT_KECCAK_ROWS}"))
-        .parse()
-        .expect("Cannot parse KECCAK_ROWS env var as usize")
+    data_hash_rounds + padding_round
 }
 
 /// Return
@@ -42,9 +45,14 @@ pub(crate) fn get_indices(preimages: &[Vec<u8>]) -> (Vec<usize>, Vec<usize>) {
     let mut digest_indices = vec![];
     let mut round_ctr = 0;
 
+    let keccak_f_rows = get_num_rows_per_update();
+    let inner_round_rows = get_num_rows_per_round();
+
     for preimage in preimages.iter().take(MAX_AGG_SNARKS + 1) {
         //  136 = 17 * 8 is the size in bytes of each
         //  input chunk that can be processed by Keccak circuit using absorb
+
+        //  For example, if num_rows_per_inner_round for Keccak is 12, then
         //  each chunk of size 136 needs 300 Keccak circuit rows to prove
         //  which consists of 12 Keccak rows for each of 24 + 1 Keccak circuit rounds
         //  digest only happens at the end of the last input chunk with
@@ -53,38 +61,47 @@ pub(crate) fn get_indices(preimages: &[Vec<u8>]) -> (Vec<usize>, Vec<usize>) {
         let mut preimage_padded = preimage.clone();
         preimage_padded.resize(INPUT_LEN_PER_ROUND * num_rounds, 0);
         for (i, round) in preimage_padded.chunks(INPUT_LEN_PER_ROUND).enumerate() {
+            let f_round_offset = round_ctr * keccak_f_rows;
             // indices for preimages
             for (j, _chunk) in round.chunks(8).into_iter().enumerate() {
+                let inner_offset = f_round_offset + (j + 1) * inner_round_rows;
                 for k in 0..8 {
-                    preimage_indices.push(round_ctr * 300 + j * 12 + k + 12)
+                    preimage_indices.push(inner_offset + k);
                 }
             }
             // indices for digests
             if i == num_rounds - 1 {
                 for j in 0..4 {
+                    let inner_offset = f_round_offset
+                        + j * inner_round_rows
+                        + (keccak_f_rows - inner_round_rows * (DIGEST_LEN / 8));
                     for k in 0..8 {
-                        digest_indices.push(round_ctr * 300 + j * 12 + k + 252)
+                        digest_indices.push(inner_offset + k);
                     }
                 }
             }
             round_ctr += 1;
         }
     }
     // last hash is for data_hash and has various length, so we output all the possible cells
-    for _i in 0..3 {
+    for _i in 0..get_data_hash_keccak_updates(MAX_AGG_SNARKS) {
         for (j, _) in (0..INPUT_LEN_PER_ROUND)
             .into_iter()
             .chunks(8)
             .into_iter()
             .enumerate()
         {
+            let inner_offset = round_ctr * keccak_f_rows + (j + 1) * inner_round_rows;
             for k in 0..8 {
-                preimage_indices.push(round_ctr * 300 + j * 12 + k + 12)
+                preimage_indices.push(inner_offset + k);
             }
         }
         for j in 0..4 {
+            let inner_offset = round_ctr * keccak_f_rows
+                + j * inner_round_rows
+                + (keccak_f_rows - inner_round_rows * (DIGEST_LEN / 8));
             for k in 0..8 {
-                digest_indices.push(round_ctr * 300 + j * 12 + k + 252)
+                digest_indices.push(inner_offset + k);
             }
         }
         round_ctr += 1;

zkevm-circuits/src/keccak_circuit.rs

@@ -65,6 +65,8 @@ pub struct KeccakCircuitConfig<F> {
     normalize_6: [TableColumn; 2],
     chi_base_table: [TableColumn; 2],
     pack_table: [TableColumn; 2],
+    /// The column for enabling copy constraints in aggregator
+    pub preimage_column: usize,
     _marker: PhantomData<F>,
 }
 
@@ -138,6 +140,8 @@ impl<F: Field> SubCircuitConfig<F> for KeccakCircuitConfig<F> {
                 s_next[i][j] = cell.at_offset(meta, get_num_rows_per_round() as i32).expr();
             }
         }
+        log::debug!("- Post states:");
+        log::debug!("Columns: {}", cell_manager.get_width());
         // Absorb data
         let absorb_from = cell_manager.query_cell(meta);
         let absorb_data = cell_manager.query_cell(meta);
@@ -185,6 +189,7 @@ impl<F: Field> SubCircuitConfig<F> for KeccakCircuitConfig<F> {
         log::debug!("- Post absorb:");
         log::debug!("Lookups: {}", lookup_counter);
         log::debug!("Columns: {}", cell_manager.get_width());
+        let preimage_column: usize = cell_manager.get_width() + 1;
         total_lookup_counter += lookup_counter;
 
         // Process inputs.
@@ -863,6 +868,7 @@ impl<F: Field> SubCircuitConfig<F> for KeccakCircuitConfig<F> {
             normalize_6,
             chi_base_table,
             pack_table,
+            preimage_column,
             _marker: PhantomData,
         }
     }
@@ -1073,6 +1079,17 @@ impl<F: Field> KeccakCircuit<F> {
         }
     }
 
+    /// The number of keccak_f's that can be done for
+    /// a particular row number depending on current Keccak params
+    pub fn capacity_for_row(num_rows: usize) -> Option<usize> {
+        if num_rows > 0 {
+            // Subtract two for unusable rows
+            Some(num_rows / ((NUM_ROUNDS + 1) * get_num_rows_per_round()) - 2)
+        } else {
+            None
+        }
+    }
+
     /// Sets the witness using the data to be hashed
     pub(crate) fn generate_witness(&self, challenges: Challenges<Value<F>>) -> Vec<KeccakRow<F>> {
         multi_keccak(self.inputs.as_slice(), challenges, self.capacity())