From d1034b05754219b603508ef79c114d908c94c1e9 Mon Sep 17 00:00:00 2001
From: Akash <112477155+akash-chandrakar@users.noreply.github.com>
Date: Thu, 26 Oct 2023 13:54:04 +0400
Subject: [PATCH] fix(prover-fri): Update setup loading for node agg circuit
 (#323)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# What ❔

## Update setup loading for node agg circuit

## Why ❔

One one setup key exist for all node agg

## Checklist


- [ *] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [* ] Tests for the changes have been added / updated.
- [* ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zk fmt` and `zk lint`.
---
 .../src/gpu_prover_job_processor.rs           |  5 +-
 prover/prover_fri/src/prover_job_processor.rs |  5 +-
 prover/prover_fri/src/utils.rs                | 54 ++++++++++++++++++-
 3 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/prover/prover_fri/src/gpu_prover_job_processor.rs b/prover/prover_fri/src/gpu_prover_job_processor.rs
index 270f94998c61..79fef4bcc4bc 100644
--- a/prover/prover_fri/src/gpu_prover_job_processor.rs
+++ b/prover/prover_fri/src/gpu_prover_job_processor.rs
@@ -30,8 +30,8 @@ pub mod gpu_prover {
     };
 
     use crate::utils::{
-        save_proof, setup_metadata_to_setup_data_key, verify_proof, GpuProverJob, ProverArtifacts,
-        SharedWitnessVectorQueue,
+        get_setup_data_key, save_proof, setup_metadata_to_setup_data_key, verify_proof,
+        GpuProverJob, ProverArtifacts, SharedWitnessVectorQueue,
     };
 
     type DefaultTranscript = GoldilocksPoisedon2Transcript;
@@ -90,6 +90,7 @@ pub mod gpu_prover {
             &self,
             key: ProverServiceDataKey,
         ) -> anyhow::Result<Arc<GoldilocksGpuProverSetupData>> {
+            let key = get_setup_data_key(key);
             Ok(match &self.setup_load_mode {
                 SetupLoadMode::FromMemory(cache) => cache
                     .get(&key)
diff --git a/prover/prover_fri/src/prover_job_processor.rs b/prover/prover_fri/src/prover_job_processor.rs
index e4354c697113..bedee075b83d 100644
--- a/prover/prover_fri/src/prover_job_processor.rs
+++ b/prover/prover_fri/src/prover_job_processor.rs
@@ -31,7 +31,9 @@ use zksync_vk_setup_data_server_fri::{
     get_cpu_setup_data_for_circuit_type, GoldilocksProverSetupData,
 };
 
-use crate::utils::{save_proof, setup_metadata_to_setup_data_key, verify_proof, ProverArtifacts};
+use crate::utils::{
+    get_setup_data_key, save_proof, setup_metadata_to_setup_data_key, verify_proof, ProverArtifacts,
+};
 
 pub enum SetupLoadMode {
     FromMemory(HashMap<ProverServiceDataKey, Arc<GoldilocksProverSetupData>>),
@@ -76,6 +78,7 @@ impl Prover {
         &self,
         key: ProverServiceDataKey,
     ) -> anyhow::Result<Arc<GoldilocksProverSetupData>> {
+        let key = get_setup_data_key(key);
         Ok(match &self.setup_load_mode {
             SetupLoadMode::FromMemory(cache) => cache
                 .get(&key)
diff --git a/prover/prover_fri/src/utils.rs b/prover/prover_fri/src/utils.rs
index de8791c2916b..5b0f8be04a47 100644
--- a/prover/prover_fri/src/utils.rs
+++ b/prover/prover_fri/src/utils.rs
@@ -17,14 +17,16 @@ use zksync_prover_fri_types::circuit_definitions::boojum::cs::implementations::v
 use zksync_prover_fri_types::circuit_definitions::boojum::field::goldilocks::{
     GoldilocksExt2, GoldilocksField,
 };
-use zksync_prover_fri_types::circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerProof;
+use zksync_prover_fri_types::circuit_definitions::circuit_definitions::recursion_layer::{
+    ZkSyncRecursionLayerProof, ZkSyncRecursionLayerStorageType,
+};
 use zksync_prover_fri_types::queue::FixedSizeQueue;
 use zksync_prover_fri_types::{
     CircuitWrapper, FriProofWrapper, ProverServiceDataKey, WitnessVectorArtifacts,
 };
 use zksync_prover_fri_utils::get_base_layer_circuit_id_for_recursive_layer;
 
-use zksync_types::{basic_fri_types::CircuitIdRoundTuple, L1BatchNumber};
+use zksync_types::{basic_fri_types::CircuitIdRoundTuple, proofs::AggregationRound, L1BatchNumber};
 
 pub type F = GoldilocksField;
 pub type H = GoldilocksPoseidon2Sponge<AbsorptionModeOverwrite>;
@@ -161,3 +163,51 @@ pub fn setup_metadata_to_setup_data_key(
         round: setup_metadata.aggregation_round.into(),
     }
 }
+
+pub fn get_setup_data_key(key: ProverServiceDataKey) -> ProverServiceDataKey {
+    match key.round {
+        AggregationRound::NodeAggregation => {
+            // For node aggregation only one key exist for all circuit types
+            ProverServiceDataKey {
+                circuit_id: ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8,
+                round: key.round,
+            }
+        }
+        _ => key,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_get_setup_data_key_for_node_agg_key() {
+        let key = ProverServiceDataKey {
+            circuit_id: 10,
+            round: AggregationRound::NodeAggregation,
+        };
+        let expected = ProverServiceDataKey {
+            circuit_id: ZkSyncRecursionLayerStorageType::NodeLayerCircuit as u8,
+            round: AggregationRound::NodeAggregation,
+        };
+
+        let result = get_setup_data_key(key);
+
+        // Check if the circuit_id has been changed to NodeLayerCircuit's id
+        assert_eq!(expected, result);
+    }
+
+    #[test]
+    fn test_get_setup_data_key_for_non_node_agg_key() {
+        let key = ProverServiceDataKey {
+            circuit_id: 10,
+            round: AggregationRound::BasicCircuits,
+        };
+
+        let result = get_setup_data_key(key.clone());
+
+        // Check if the key has remained same
+        assert_eq!(key, result);
+    }
+}