impl artifacts manager for remaining parts

prover/crates/bin/witness_generator/src/artifacts.rs

@@ -32,7 +32,7 @@ pub(crate) trait ArtifactsManager {
     async fn get_artifacts(
         metadata: &Self::InputMetadata,
         object_store: &dyn ObjectStore,
-    ) -> Self::InputArtifacts;
+    ) -> anyhow::Result<Self::InputArtifacts>;
 
     async fn save_artifacts(
         job_id: u32,

prover/crates/bin/witness_generator/src/basic_circuits/artifacts.rs

@@ -7,10 +7,9 @@ use zksync_prover_fri_types::AuxOutputWitnessWrapper;
 use zksync_prover_fri_utils::get_recursive_layer_circuit_id_for_base_layer;
 use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
 
-use crate::basic_circuits::BasicCircuitArtifacts;
 use crate::{
     artifacts::{ArtifactsManager, BlobUrls},
-    basic_circuits::{BasicWitnessGenerator, BasicWitnessGeneratorJob},
+    basic_circuits::{BasicCircuitArtifacts, BasicWitnessGenerator, BasicWitnessGeneratorJob},
     utils::SchedulerPartialInputWrapper,
 };
 
@@ -23,13 +22,13 @@ impl ArtifactsManager for BasicWitnessGenerator {
     async fn get_artifacts(
         metadata: &Self::InputMetadata,
         object_store: &dyn ObjectStore,
-    ) -> Self::InputArtifacts {
+    ) -> anyhow::Result<Self::InputArtifacts> {
         let l1_batch_number = *metadata;
-        let job = object_store.get(l1_batch_number).await.unwrap();
-        BasicWitnessGeneratorJob {
+        let data = object_store.get(l1_batch_number).await.unwrap();
+        Ok(BasicWitnessGeneratorJob {
             block_number: l1_batch_number,
-            job,
-        }
+            data,
+        })
     }
 
     async fn save_artifacts(

prover/crates/bin/witness_generator/src/basic_circuits/job_processor.rs

@@ -41,10 +41,9 @@ impl JobProcessor for BasicWitnessGenerator {
                     block_number
                 );
                 let started_at = Instant::now();
-                let job = Self::get_artifacts(&block_number, &*self.object_store).await;
+                let job = Self::get_artifacts(&block_number, &*self.object_store).await?;
 
-                crate::metrics::WITNESS_GENERATOR_METRICS.blob_fetch_time
-                    [&AggregationRound::BasicCircuits.into()]
+                WITNESS_GENERATOR_METRICS.blob_fetch_time[&AggregationRound::BasicCircuits.into()]
                     .observe(started_at.elapsed());
 
                 Ok(Some((block_number, job)))

prover/crates/bin/witness_generator/src/basic_circuits/mod.rs

@@ -69,7 +69,7 @@ pub struct BasicCircuitArtifacts {
 #[derive(Clone)]
 pub struct BasicWitnessGeneratorJob {
     pub(super) block_number: L1BatchNumber,
-    pub(super) job: WitnessInputData,
+    pub(super) data: WitnessInputData,
 }
 
 #[derive(Debug)]
@@ -115,7 +115,10 @@ impl BasicWitnessGenerator {
         started_at: Instant,
         max_circuits_in_flight: usize,
     ) -> Option<BasicCircuitArtifacts> {
-        let BasicWitnessGeneratorJob { block_number, job } = basic_job;
+        let BasicWitnessGeneratorJob {
+            block_number,
+            data: job,
+        } = basic_job;
 
         tracing::info!(
             "Starting witness generation of type {:?} for block {}",

prover/crates/bin/witness_generator/src/leaf_aggregation/artifacts.rs

@@ -1,6 +1,6 @@
-use async_trait::async_trait;
 use std::time::Instant;
 
+use async_trait::async_trait;
 use zksync_object_store::ObjectStore;
 use zksync_prover_dal::{ConnectionPool, Prover, ProverDal};
 use zksync_prover_fri_types::keys::ClosedFormInputKey;
@@ -23,16 +23,18 @@ impl ArtifactsManager for LeafAggregationWitnessGenerator {
     async fn get_artifacts(
         metadata: &Self::InputMetadata,
         object_store: &dyn ObjectStore,
-    ) -> Self::InputArtifacts {
+    ) -> anyhow::Result<Self::InputArtifacts> {
         let key = ClosedFormInputKey {
             block_number: metadata.block_number,
             circuit_id: metadata.circuit_id,
         };
 
-        object_store
+        let artifacts = object_store
             .get(key)
             .await
-            .unwrap_or_else(|_| panic!("leaf aggregation job artifacts missing: {:?}", key))
+            .unwrap_or_else(|_| panic!("leaf aggregation job artifacts missing: {:?}", key));
+
+        Ok(artifacts)
     }
 
     #[tracing::instrument(

prover/crates/bin/witness_generator/src/leaf_aggregation/mod.rs

@@ -118,7 +118,7 @@ pub async fn prepare_leaf_aggregation_job(
 ) -> anyhow::Result<LeafAggregationWitnessGeneratorJob> {
     let started_at = Instant::now();
     let closed_form_input =
-        LeafAggregationWitnessGenerator::get_artifacts(&metadata, object_store).await;
+        LeafAggregationWitnessGenerator::get_artifacts(&metadata, object_store).await?;
 
     WITNESS_GENERATOR_METRICS.blob_fetch_time[&AggregationRound::LeafAggregation.into()]
         .observe(started_at.elapsed());

prover/crates/bin/witness_generator/src/node_aggregation/artifacts.rs

@@ -1,6 +1,6 @@
-use async_trait::async_trait;
 use std::time::Instant;
 
+use async_trait::async_trait;
 use zksync_object_store::ObjectStore;
 use zksync_prover_dal::{ConnectionPool, Prover, ProverDal};
 use zksync_prover_fri_types::keys::AggregationsKey;
@@ -26,18 +26,20 @@ impl ArtifactsManager for NodeAggregationWitnessGenerator {
     async fn get_artifacts(
         metadata: &Self::InputMetadata,
         object_store: &dyn ObjectStore,
-    ) -> Self::InputArtifacts {
+    ) -> anyhow::Result<Self::InputArtifacts> {
         let key = AggregationsKey {
             block_number: metadata.block_number,
             circuit_id: metadata.circuit_id,
             depth: metadata.depth,
         };
-        object_store.get(key).await.unwrap_or_else(|error| {
+        let artifacts = object_store.get(key).await.unwrap_or_else(|error| {
             panic!(
                 "node aggregation job artifacts getting error. Key: {:?}, error: {:?}",
                 key, error
             )
-        })
+        });
+
+        Ok(artifacts)
     }
 
     #[tracing::instrument(

prover/crates/bin/witness_generator/src/node_aggregation/mod.rs

@@ -225,7 +225,7 @@ pub async fn prepare_job(
     keystore: Keystore,
 ) -> anyhow::Result<NodeAggregationWitnessGeneratorJob> {
     let started_at = Instant::now();
-    let artifacts = NodeAggregationWitnessGenerator::get_artifacts(&metadata, object_store).await;
+    let artifacts = NodeAggregationWitnessGenerator::get_artifacts(&metadata, object_store).await?;
 
     WITNESS_GENERATOR_METRICS.blob_fetch_time[&AggregationRound::NodeAggregation.into()]
         .observe(started_at.elapsed());

prover/crates/bin/witness_generator/src/recursion_tip/artifacts.rs

@@ -1,10 +1,14 @@
-use async_trait::async_trait;
-use std::time::Instant;
+use std::{collections::HashMap, time::Instant};
 
-use circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType;
+use async_trait::async_trait;
+use circuit_definitions::{
+    circuit_definitions::recursion_layer::{ZkSyncRecursionLayerStorageType, ZkSyncRecursionProof},
+    zkevm_circuits::scheduler::aux::BaseLayerCircuitType,
+};
+use zkevm_test_harness::empty_node_proof;
 use zksync_object_store::ObjectStore;
 use zksync_prover_dal::{ConnectionPool, Prover, ProverDal};
-use zksync_prover_fri_types::{keys::FriCircuitKey, CircuitWrapper};
+use zksync_prover_fri_types::{keys::FriCircuitKey, CircuitWrapper, FriProofWrapper};
 use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
 
 use crate::{
@@ -14,15 +18,59 @@ use crate::{
 
 #[async_trait]
 impl ArtifactsManager for RecursionTipWitnessGenerator {
-    type InputMetadata = ();
-    type InputArtifacts = ();
+    type InputMetadata = Vec<(u8, u32)>;
+    type InputArtifacts = Vec<ZkSyncRecursionProof>;
     type OutputArtifacts = RecursionTipArtifacts;
 
+    /// Loads all proofs for a given recursion tip's job ids.
+    /// Note that recursion tip may not have proofs for some specific circuits (because the batch didn't contain them).
+    /// In this scenario, we still need to pass a proof, but it won't be taken into account during proving.
+    /// For this scenario, we use an empty_proof, but any proof would suffice.
     async fn get_artifacts(
-        metadata: &Self::InputMetadata,
+        metadata: &Vec<(u8, u32)>,
         object_store: &dyn ObjectStore,
-    ) -> Self::InputArtifacts {
-        todo!()
+    ) -> anyhow::Result<Vec<ZkSyncRecursionProof>> {
+        let job_mapping: HashMap<u8, u32> = metadata
+            .clone()
+            .into_iter()
+            .map(|(leaf_circuit_id, job_id)| {
+                (
+                    ZkSyncRecursionLayerStorageType::from_leaf_u8_to_basic_u8(leaf_circuit_id),
+                    job_id,
+                )
+            })
+            .collect();
+
+        let empty_proof = empty_node_proof().into_inner();
+
+        let mut proofs = Vec::new();
+        for circuit_id in BaseLayerCircuitType::as_iter_u8() {
+            if job_mapping.contains_key(&circuit_id) {
+                let fri_proof_wrapper = object_store
+                    .get(*job_mapping.get(&circuit_id).unwrap())
+                    .await
+                    .unwrap_or_else(|_| {
+                        panic!(
+                            "Failed to load proof with circuit_id {} for recursion tip",
+                            circuit_id
+                        )
+                    });
+                match fri_proof_wrapper {
+                    FriProofWrapper::Base(_) => {
+                        return Err(anyhow::anyhow!(
+                        "Expected only recursive proofs for recursion tip, got Base for circuit {}",
+                        circuit_id
+                    ));
+                    }
+                    FriProofWrapper::Recursive(recursive_proof) => {
+                        proofs.push(recursive_proof.into_inner());
+                    }
+                }
+            } else {
+                proofs.push(empty_proof.clone());
+            }
+        }
+        Ok(proofs)
     }
 
     async fn save_artifacts(

prover/crates/bin/witness_generator/src/recursion_tip/mod.rs

@@ -1,6 +1,3 @@
-mod artifacts;
-mod job_processor;
-
 use std::{sync::Arc, time::Instant};
 
 use anyhow::Context;
@@ -47,10 +44,12 @@ use zksync_types::{
 };
 
 use crate::{
-    metrics::WITNESS_GENERATOR_METRICS,
-    utils::{load_proofs_for_recursion_tip, ClosedFormInputWrapper},
+    artifacts::ArtifactsManager, metrics::WITNESS_GENERATOR_METRICS, utils::ClosedFormInputWrapper,
 };
 
+mod artifacts;
+mod job_processor;
+
 #[derive(Clone)]
 pub struct RecursionTipWitnessGeneratorJob {
     block_number: L1BatchNumber,
@@ -148,7 +147,8 @@ pub async fn prepare_job(
 ) -> anyhow::Result<RecursionTipWitnessGeneratorJob> {
     let started_at = Instant::now();
     let recursion_tip_proofs =
-        load_proofs_for_recursion_tip(final_node_proof_job_ids, object_store).await?;
+        RecursionTipWitnessGenerator::get_artifacts(&final_node_proof_job_ids, object_store)
+            .await?;
     WITNESS_GENERATOR_METRICS.blob_fetch_time[&AggregationRound::RecursionTip.into()]
         .observe(started_at.elapsed());
 

prover/crates/bin/witness_generator/src/scheduler/artifacts.rs

@@ -1,10 +1,10 @@
-use async_trait::async_trait;
 use std::time::Instant;
 
+use async_trait::async_trait;
 use circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType;
 use zksync_object_store::ObjectStore;
 use zksync_prover_dal::{ConnectionPool, Prover, ProverDal};
-use zksync_prover_fri_types::{keys::FriCircuitKey, CircuitWrapper};
+use zksync_prover_fri_types::{keys::FriCircuitKey, CircuitWrapper, FriProofWrapper};
 use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
 
 use crate::{
@@ -14,15 +14,17 @@ use crate::{
 
 #[async_trait]
 impl ArtifactsManager for SchedulerWitnessGenerator {
-    type InputMetadata = ();
-    type InputArtifacts = ();
+    type InputMetadata = u32;
+    type InputArtifacts = FriProofWrapper;
     type OutputArtifacts = SchedulerArtifacts;
 
     async fn get_artifacts(
         metadata: &Self::InputMetadata,
         object_store: &dyn ObjectStore,
-    ) -> Self::InputArtifacts {
-        todo!()
+    ) -> anyhow::Result<Self::InputArtifacts> {
+        let artifacts = object_store.get(*metadata).await?;
+
+        Ok(artifacts)
     }
 
     async fn save_artifacts(

prover/crates/bin/witness_generator/src/scheduler/mod.rs

@@ -1,6 +1,3 @@
-mod artifacts;
-mod job_processor;
-
 use std::{convert::TryInto, sync::Arc, time::Instant};
 
 use anyhow::Context as _;
@@ -30,7 +27,13 @@ use zksync_types::{
     basic_fri_types::AggregationRound, protocol_version::ProtocolSemanticVersion, L1BatchNumber,
 };
 
-use crate::{metrics::WITNESS_GENERATOR_METRICS, utils::SchedulerPartialInputWrapper};
+use crate::{
+    artifacts::ArtifactsManager, metrics::WITNESS_GENERATOR_METRICS,
+    utils::SchedulerPartialInputWrapper,
+};
+
+mod artifacts;
+mod job_processor;
 
 #[derive(Clone)]
 pub struct SchedulerArtifacts {
@@ -132,7 +135,8 @@ pub async fn prepare_job(
     keystore: Keystore,
 ) -> anyhow::Result<SchedulerWitnessGeneratorJob> {
     let started_at = Instant::now();
-    let wrapper = object_store.get(recursion_tip_job_id).await?;
+    let wrapper =
+        SchedulerWitnessGenerator::get_artifacts(&recursion_tip_job_id, object_store).await?;
     let recursion_tip_proof = match wrapper {
         FriProofWrapper::Base(_) => Err(anyhow::anyhow!(
             "Expected only recursive proofs for scheduler l1 batch {l1_batch_number}, got Base"