From 9b58ae97875455d58d42fe203cfb1f51cb270f62 Mon Sep 17 00:00:00 2001
From: Vladislav Volosnikov <Volosnikov.apmath@gmail.com>
Date: Wed, 7 Aug 2024 10:06:32 +0200
Subject: [PATCH] fix(prover): Fix NWG (#2590)

Invalid repacking of proofs for processing and invalid assertion do not
allow to complete NWG.
---
 .../witness_generator/src/leaf_aggregation.rs |  2 +
 .../witness_generator/src/node_aggregation.rs | 42 +++++++++++--------
 .../crates/bin/witness_generator/src/utils.rs | 26 ++++++++++++
 3 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/prover/crates/bin/witness_generator/src/leaf_aggregation.rs b/prover/crates/bin/witness_generator/src/leaf_aggregation.rs
index 00546a11407e..d8cad84e777d 100644
--- a/prover/crates/bin/witness_generator/src/leaf_aggregation.rs
+++ b/prover/crates/bin/witness_generator/src/leaf_aggregation.rs
@@ -267,6 +267,8 @@ pub async fn process_leaf_aggregation_job(
     let circuit_id = job.circuit_id;
     let queues = split_recursion_queue(job.closed_form_inputs.1);
 
+    assert_eq!(circuit_id, job.base_vk.numeric_circuit_type());
+
     let aggregations = queues
         .iter()
         .cloned()
diff --git a/prover/crates/bin/witness_generator/src/node_aggregation.rs b/prover/crates/bin/witness_generator/src/node_aggregation.rs
index ec6d011f8644..c9d5ab32bc5f 100644
--- a/prover/crates/bin/witness_generator/src/node_aggregation.rs
+++ b/prover/crates/bin/witness_generator/src/node_aggregation.rs
@@ -8,7 +8,7 @@ use zkevm_test_harness::witness::recursive_aggregation::{
     compute_node_vk_commitment, create_node_witness,
 };
 use zksync_config::configs::FriWitnessGeneratorConfig;
-use zksync_object_store::ObjectStore;
+use zksync_object_store::{ObjectStore, ObjectStoreError};
 use zksync_prover_dal::{ConnectionPool, Prover, ProverDal};
 use zksync_prover_fri_types::{
     circuit_definitions::{
@@ -34,7 +34,7 @@ use crate::{
     metrics::WITNESS_GENERATOR_METRICS,
     utils::{
         load_proofs_for_job_ids, save_node_aggregations_artifacts,
-        save_recursive_layer_prover_input_artifacts, AggregationWrapper,
+        save_recursive_layer_prover_input_artifacts, AggregationWrapper, AggregationWrapperLegacy,
     },
 };
 
@@ -113,14 +113,7 @@ impl NodeAggregationWitnessGenerator {
         let mut proof_ids_iter = job.proofs_ids.into_iter();
         let mut proofs_ids = vec![];
         for queues in job.aggregations.chunks(RECURSION_ARITY) {
-            let mut proofs_for_chunk = vec![];
-            for (_, queue) in queues {
-                let proofs_ids_for_queue: Vec<_> = (&mut proof_ids_iter)
-                    .take(queue.num_items as usize)
-                    .collect();
-                assert_eq!(queue.num_items as usize, proofs_ids_for_queue.len());
-                proofs_for_chunk.extend(proofs_ids_for_queue);
-            }
+            let proofs_for_chunk: Vec<_> = (&mut proof_ids_iter).take(queues.len()).collect();
             proofs_ids.push(proofs_for_chunk);
         }
 
@@ -175,14 +168,12 @@ impl NodeAggregationWitnessGenerator {
                     &all_leafs_layer_params,
                 );
 
-                assert_eq!(job.circuit_id as u64, result_circuit_id);
-
                 let recursive_circuit_id_and_url = save_recursive_layer_prover_input_artifacts(
                     job.block_number,
                     circuit_idx,
                     vec![recursive_circuit],
                     AggregationRound::NodeAggregation,
-                    job.depth,
+                    job.depth + 1,
                     &*object_store,
                     Some(job.circuit_id),
                 )
@@ -453,10 +444,27 @@ async fn get_artifacts(
         circuit_id: metadata.circuit_id,
         depth: metadata.depth,
     };
-    object_store
-        .get(key)
-        .await
-        .unwrap_or_else(|_| panic!("node aggregation job artifacts missing: {:?}", key))
+    let result = object_store.get(key).await;
+
+    // TODO: remove after transition
+    return match result {
+        Ok(aggregation_wrapper) => aggregation_wrapper,
+        Err(error) => {
+            // probably legacy struct is saved in GCS
+            if let ObjectStoreError::Serialization(serialization_error) = error {
+                let legacy_wrapper: AggregationWrapperLegacy =
+                    object_store.get(key).await.unwrap_or_else(|inner_error| {
+                        panic!(
+                            "node aggregation job artifacts getting error. Key: {:?}, errors: {:?} {:?}",
+                            key, serialization_error, inner_error
+                        )
+                    });
+                AggregationWrapper(legacy_wrapper.0.into_iter().map(|x| (x.0, x.1)).collect())
+            } else {
+                panic!("node aggregation job artifacts missing: {:?}", key)
+            }
+        }
+    };
 }
 
 #[tracing::instrument(
diff --git a/prover/crates/bin/witness_generator/src/utils.rs b/prover/crates/bin/witness_generator/src/utils.rs
index c1fd5dd3b407..a21aabc5d6d1 100644
--- a/prover/crates/bin/witness_generator/src/utils.rs
+++ b/prover/crates/bin/witness_generator/src/utils.rs
@@ -97,6 +97,32 @@ impl StoredObject for AggregationWrapper {
     serialize_using_bincode!();
 }
 
+/// TODO: remove after transition
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct AggregationWrapperLegacy(
+    pub  Vec<(
+        u64,
+        RecursionQueueSimulator<GoldilocksField>,
+        ZkSyncRecursiveLayerCircuit,
+    )>,
+);
+
+impl StoredObject for AggregationWrapperLegacy {
+    const BUCKET: Bucket = Bucket::NodeAggregationWitnessJobsFri;
+    type Key<'a> = AggregationsKey;
+
+    fn encode_key(key: Self::Key<'_>) -> String {
+        let AggregationsKey {
+            block_number,
+            circuit_id,
+            depth,
+        } = key;
+        format!("aggregations_{block_number}_{circuit_id}_{depth}.bin")
+    }
+
+    serialize_using_bincode!();
+}
+
 #[derive(serde::Serialize, serde::Deserialize)]
 pub struct SchedulerPartialInputWrapper(
     pub  SchedulerCircuitInstanceWitness<