fix(prover-fri): Update setup loading for node agg circuit (#323)

# 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`.

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)

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)

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);
+    }
+}