From 4d48233ca774d20988d51a832e5ebc52dfd824e4 Mon Sep 17 00:00:00 2001
From: Rohit Narurkar <rohit.narurkar@proton.me>
Date: Tue, 2 Jul 2024 12:14:26 +0100
Subject: [PATCH] compilation fix and batch header's hash (test)

---
 aggregator/src/blob.rs                  | 23 ++++++++++++-
 aggregator/src/chunk.rs                 | 43 +++++++++++++++++++++++++
 aggregator/src/lib.rs                   |  2 +-
 prover/src/common/prover/aggregation.rs |  4 +--
 prover/src/recursion.rs                 |  4 +--
 5 files changed, 70 insertions(+), 6 deletions(-)

diff --git a/aggregator/src/blob.rs b/aggregator/src/blob.rs
index aa6871c90d..359b12de0f 100644
--- a/aggregator/src/blob.rs
+++ b/aggregator/src/blob.rs
@@ -708,7 +708,7 @@ mod tests {
                 vec![vec![]; MAX_AGG_SNARKS],
             ),
             (
-                "max number of chunkks all non-empty",
+                "max number of chunks all non-empty",
                 (0..MAX_AGG_SNARKS)
                     .map(|i| (10u8..11 + u8::try_from(i).unwrap()).collect())
                     .collect(),
@@ -742,6 +742,27 @@ mod tests {
         ]
         .iter()
         {
+            // batch header
+            let batch_header = crate::batch::BatchHeader {
+                version: 3,
+                batch_index: 6789,
+                l1_message_popped: 101,
+                total_l1_message_popped: 10101,
+                parent_batch_hash: H256::repeat_byte(1),
+                last_block_timestamp: 192837,
+                ..Default::default()
+            };
+            let chunks_without_padding = crate::chunk::ChunkInfo::mock_chunk_infos(&tcase);
+            let batch_hash = BatchHash::<MAX_AGG_SNARKS>::construct_with_unpadded(
+                &chunks_without_padding,
+                batch_header,
+            );
+            println!(
+                "[[ {:60} ]] batch_hash = {:0>64x}\n\n",
+                annotation, batch_hash.current_batch_hash,
+            );
+
+            // blob data
             let batch_data: BatchData<MAX_AGG_SNARKS> = tcase.into();
             let point_evaluation_assignments = PointEvaluationAssignments::from(&batch_data);
             let versioned_hash = batch_data.get_versioned_hash();
diff --git a/aggregator/src/chunk.rs b/aggregator/src/chunk.rs
index 697d35237c..4aad73da65 100644
--- a/aggregator/src/chunk.rs
+++ b/aggregator/src/chunk.rs
@@ -175,6 +175,49 @@ impl ChunkInfo {
         H256(keccak256(&self.tx_bytes))
     }
 
+    #[cfg(test)]
+    pub(crate) fn mock_chunk_infos(txs_data: &[Vec<u8>]) -> Vec<Self> {
+        use crate::MAX_AGG_SNARKS;
+
+        assert!(txs_data.len() <= MAX_AGG_SNARKS);
+        let state_roots: [H256; MAX_AGG_SNARKS + 1] = (0..=MAX_AGG_SNARKS)
+            .map(|i| {
+                let i = i as u8;
+                let mut state_root = [0u8; 32];
+                state_root[31] = i;
+                state_root.into()
+            })
+            .collect::<Vec<_>>()
+            .try_into()
+            .expect("should not fail");
+
+        txs_data
+            .iter()
+            .enumerate()
+            .map(|(i, tx_data)| {
+                let withdraw_root = {
+                    let mut root = [0u8; 32];
+                    root[31] = 255 - (i as u8);
+                    root.into()
+                };
+                let data_hash = {
+                    let mut root = [0u8; 32];
+                    root[0] = 255 - (i as u8);
+                    root.into()
+                };
+                ChunkInfo {
+                    chain_id: 123456,
+                    prev_state_root: state_roots[i],
+                    post_state_root: state_roots[i + 1],
+                    withdraw_root,
+                    data_hash,
+                    tx_bytes: tx_data.to_vec(),
+                    is_padding: false,
+                }
+            })
+            .collect::<Vec<_>>()
+    }
+
     /// Sample a chunk info from random (for testing)
     #[cfg(test)]
     pub(crate) fn mock_random_chunk_info_for_testing<R: rand::RngCore>(r: &mut R) -> Self {
diff --git a/aggregator/src/lib.rs b/aggregator/src/lib.rs
index 7dedb40094..e31abe4279 100644
--- a/aggregator/src/lib.rs
+++ b/aggregator/src/lib.rs
@@ -27,7 +27,7 @@ mod tests;
 
 pub use self::core::extract_proof_and_instances_with_pairing_check;
 pub use aggregation::*;
-pub use batch::BatchHash;
+pub use batch::{BatchHash, BatchHeader};
 pub use blob::BatchData;
 pub use chunk::ChunkInfo;
 pub use compression::*;
diff --git a/prover/src/common/prover/aggregation.rs b/prover/src/common/prover/aggregation.rs
index 5f962fd6b5..2411ff7987 100644
--- a/prover/src/common/prover/aggregation.rs
+++ b/prover/src/common/prover/aggregation.rs
@@ -4,7 +4,7 @@ use crate::{
     io::{load_snark, write_snark},
     utils::gen_rng,
 };
-use aggregator::{BatchCircuit, BatchHash, ChunkInfo, MAX_AGG_SNARKS};
+use aggregator::{BatchCircuit, BatchHash, BatchHeader, ChunkInfo, MAX_AGG_SNARKS};
 use anyhow::{anyhow, Result};
 use rand::Rng;
 use snark_verifier_sdk::Snark;
@@ -21,7 +21,7 @@ impl Prover {
     ) -> Result<Snark> {
         env::set_var("AGGREGATION_CONFIG", layer_config_path(id));
 
-        let batch_hash = BatchHash::construct(chunk_hashes);
+        let batch_hash = BatchHash::construct(chunk_hashes, BatchHeader::default());
 
         let circuit: BatchCircuit<MAX_AGG_SNARKS> =
             BatchCircuit::new(self.params(degree), previous_snarks, &mut rng, batch_hash)
diff --git a/prover/src/recursion.rs b/prover/src/recursion.rs
index f391ecfe24..780782f8fd 100644
--- a/prover/src/recursion.rs
+++ b/prover/src/recursion.rs
@@ -1,6 +1,6 @@
 use halo2_proofs::halo2curves::bn256::Fr;
 
-use aggregator::{AggregationCircuit, StateTransition};
+use aggregator::{BatchCircuit, StateTransition};
 use snark_verifier_sdk::Snark;
 
 mod prover;
@@ -37,7 +37,7 @@ impl<const N_SNARK: usize> AggregatedBatchProvingTask<'_, N_SNARK> {
 
 impl<'a, const N_SNARK: usize> StateTransition for AggregatedBatchProvingTask<'a, N_SNARK> {
     type Input = &'a [Snark];
-    type Circuit = AggregationCircuit<N_SNARK>;
+    type Circuit = BatchCircuit<N_SNARK>;
 
     fn new(state: Self::Input) -> Self {
         assert!(!state.is_empty());