feat: Adding ability to generate 4844 setup key and refactor

prover/prover_fri/tests/basic_test.rs

@@ -57,9 +57,11 @@ async fn prover_and_assert_base_layer(
         CircuitWrapper::Recursive(_) => anyhow::bail!("Expected base layer circuit"),
     };
     let keystore = Keystore::default();
-    let circuit_setup_data =
-        generate_setup_data_common(&keystore, true, circuit.numeric_circuit_type())
-            .context("generate_cpu_base_layers_setup_data()")?;
+    let circuit_setup_data = generate_setup_data_common(
+        &keystore,
+        ProverServiceDataKey::new_basic(circuit.numeric_circuit_type()),
+    )
+    .context("generate_cpu_base_layers_setup_data()")?;
     let setup_data = Arc::new(circuit_setup_data.into());
     let setup_key = ProverServiceDataKey::new(circuit_id, aggregation_round);
     let prover_job = ProverJob::new(block_number, expected_proof_id, circuit_wrapper, setup_key);

prover/prover_fri_types/src/lib.rs

@@ -10,7 +10,9 @@ use circuit_definitions::{
             ZkSyncRecursionLayerProof, ZkSyncRecursionLayerStorageType, ZkSyncRecursiveLayerCircuit,
         },
     },
-    zkevm_circuits::scheduler::block_header::BlockAuxilaryOutputWitness,
+    zkevm_circuits::scheduler::{
+        aux::BaseLayerCircuitType, block_header::BlockAuxilaryOutputWitness,
+    },
     ZkSyncDefaultRoundFunction,
 };
 use zksync_object_store::{serialize_using_bincode, Bucket, StoredObject};
@@ -138,6 +140,30 @@ impl ProverServiceDataKey {
             round: get_round_for_recursive_circuit_type(circuit_id),
         }
     }
+    pub fn new_basic(circuit_id: u8) -> Self {
+        Self {
+            circuit_id,
+            round: AggregationRound::BasicCircuits,
+        }
+    }
+
+    pub fn all_boojum() -> Vec<ProverServiceDataKey> {
+        let mut results = vec![];
+        for numeric_circuit in
+            BaseLayerCircuitType::VM as u8..=BaseLayerCircuitType::L1MessagesHasher as u8
+        {
+            results.push(ProverServiceDataKey::new_basic(numeric_circuit))
+        }
+        for numeric_circuit in ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8
+            ..=ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8
+        {
+            results.push(ProverServiceDataKey::new_recursive(numeric_circuit))
+        }
+        results.push(ProverServiceDataKey::eip4844());
+        // Don't include snark, as it uses the old proving system.
+
+        results
+    }
 
     /// Data key for snark wrapper.
     pub fn snark() -> Self {
@@ -163,6 +189,10 @@ impl ProverServiceDataKey {
             && self.round == AggregationRound::BasicCircuits
     }
 
+    pub fn is_base_layer(&self) -> bool {
+        self.round == AggregationRound::BasicCircuits && !self.is_eip4844()
+    }
+
     pub fn name(&self) -> String {
         if self.is_eip4844() {
             return "eip4844".to_string();

prover/vk_setup_data_generator_server_fri/src/main.rs

@@ -18,7 +18,10 @@ use zkevm_test_harness::{
         WrapperConfig,
     },
 };
-use zksync_prover_fri_types::circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType;
+use zksync_prover_fri_types::{
+    circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType,
+    ProverServiceDataKey,
+};
 use zksync_vk_setup_data_server_fri::{
     keystore::Keystore,
     setup_data_generator::{CPUSetupDataGenerator, GPUSetupDataGenerator, SetupDataGenerator},
@@ -85,6 +88,8 @@ enum CircuitSelector {
     Recursive,
     /// Select circuits from basic group.
     Basic,
+    /// EIP 4844 circuit
+    Eip4844,
 }
 
 #[derive(Debug, Parser)]
@@ -163,39 +168,31 @@ fn generate_setup_keys(
     generator: &dyn SetupDataGenerator,
     options: &GeneratorOptions,
 ) -> anyhow::Result<()> {
-    match options.circuits_type {
+    let circuit_type = match options.circuits_type {
         CircuitSelector::All => {
             let digests = generator.generate_all(options.dry_run)?;
             tracing::info!("Setup keys md5(s):");
-            print_stats(digests)
-        }
-        CircuitSelector::Recursive => {
-            let digest = generator
-                .generate_and_write_setup_data(
-                    false,
-                    options
-                        .numeric_circuit
-                        .expect("--numeric-circuit must be provided"),
-                    options.dry_run,
-                )
-                .context("generate_setup_data()")?;
-            tracing::info!("digest: {:?}", digest);
-            Ok(())
-        }
-        CircuitSelector::Basic => {
-            let digest = generator
-                .generate_and_write_setup_data(
-                    true,
-                    options
-                        .numeric_circuit
-                        .expect("--numeric-circuit must be provided"),
-                    options.dry_run,
-                )
-                .context("generate_setup_data()")?;
-            tracing::info!("digest: {:?}", digest);
-            Ok(())
+            print_stats(digests)?;
+            return Ok(());
         }
-    }
+        CircuitSelector::Recursive => ProverServiceDataKey::new_recursive(
+            options
+                .numeric_circuit
+                .expect("--numeric-circuit must be provided"),
+        ),
+        CircuitSelector::Basic => ProverServiceDataKey::new_basic(
+            options
+                .numeric_circuit
+                .expect("--numeric-circuit must be provided"),
+        ),
+        CircuitSelector::Eip4844 => ProverServiceDataKey::eip4844(),
+    };
+
+    let digest = generator
+        .generate_and_write_setup_data(circuit_type, options.dry_run)
+        .context("generate_setup_data()")?;
+    tracing::info!("digest: {:?}", digest);
+    Ok(())
 }
 
 fn main() -> anyhow::Result<()> {

prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs

@@ -4,16 +4,13 @@
 use std::collections::HashMap;
 
 use anyhow::Context as _;
-use circuit_definitions::{
-    circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType,
-    zkevm_circuits::scheduler::aux::BaseLayerCircuitType,
-};
 use zkevm_test_harness::{
-    compute_setups::{generate_circuit_setup_data, CircuitSetupData},
+    compute_setups::{
+        generate_circuit_setup_data, generate_circuit_setup_data_4844, CircuitSetupData,
+    },
     data_source::SetupDataSource,
 };
 use zksync_prover_fri_types::ProverServiceDataKey;
-use zksync_types::basic_fri_types::AggregationRound;
 #[cfg(feature = "gpu")]
 use {
     crate::GpuProverSetupData, shivini::cs::setup::GpuSetup, shivini::ProverContext,
@@ -26,104 +23,94 @@ use crate::{keystore::Keystore, GoldilocksProverSetupData};
 /// It also does a final sanity check to make sure that verification keys didn't change.
 pub fn generate_setup_data_common(
     keystore: &Keystore,
-    is_base_layer: bool,
-    circuit_type: u8,
+    circuit: ProverServiceDataKey,
 ) -> anyhow::Result<CircuitSetupData> {
     let mut data_source = keystore.load_keys_to_data_source()?;
-    let circuit_setup_data =
-        generate_circuit_setup_data(is_base_layer, circuit_type, &mut data_source).unwrap();
+    let circuit_setup_data = if circuit.is_eip4844() {
+        generate_circuit_setup_data_4844().unwrap()
+    } else {
+        generate_circuit_setup_data(
+            circuit.is_base_layer(),
+            circuit.circuit_id,
+            &mut data_source,
+        )
+        .unwrap()
+    };
 
-    let (finalization, vk) = if is_base_layer {
+    let (finalization, vk) = if circuit.is_eip4844() {
+        (None, data_source.get_eip4844_vk().unwrap())
+    } else if circuit.is_base_layer() {
         (
-            keystore.load_finalization_hints(ProverServiceDataKey::new(
-                circuit_type,
-                AggregationRound::BasicCircuits,
-            ))?,
+            Some(keystore.load_finalization_hints(circuit.clone())?),
             data_source
-                .get_base_layer_vk(circuit_type)
+                .get_base_layer_vk(circuit.circuit_id)
                 .unwrap()
                 .into_inner(),
         )
     } else {
         (
-            keystore.load_finalization_hints(ProverServiceDataKey::new_recursive(circuit_type))?,
+            Some(keystore.load_finalization_hints(circuit.clone())?),
             data_source
-                .get_recursion_layer_vk(circuit_type)
+                .get_recursion_layer_vk(circuit.circuit_id)
                 .unwrap()
                 .into_inner(),
         )
     };
 
     // Sanity check to make sure that generated setup data is matching.
-    if finalization != circuit_setup_data.finalization_hint {
-        anyhow::bail!("finalization hint mismatch for circuit: {circuit_type}");
+    if let Some(finalization) = finalization {
+        if finalization != circuit_setup_data.finalization_hint {
+            anyhow::bail!(
+                "finalization hint mismatch for circuit: {:?}",
+                circuit.name()
+            );
+        }
     }
     if vk != circuit_setup_data.vk {
-        anyhow::bail!("vk mismatch for circuit: {circuit_type}");
+        anyhow::bail!("vk mismatch for circuit: {:?}", circuit.name());
     }
     Ok(circuit_setup_data)
 }
 
 /// Trait to cover GPU and CPU setup data generation.
 pub trait SetupDataGenerator {
     /// Generates the setup keys.
-    fn generate_setup_data(
-        &self,
-        is_base_layer: bool,
-        numeric_circuit: u8,
-    ) -> anyhow::Result<Vec<u8>>;
+    fn generate_setup_data(&self, circuit: ProverServiceDataKey) -> anyhow::Result<Vec<u8>>;
 
     fn keystore(&self) -> &Keystore;
 
     /// Generates and stores the setup keys.
     /// Returns the md5 check sum of the stored file.
     fn generate_and_write_setup_data(
         &self,
-        is_base_layer: bool,
-        numeric_circuit: u8,
+        circuit: ProverServiceDataKey,
         dry_run: bool,
     ) -> anyhow::Result<String> {
-        let serialized = self.generate_setup_data(is_base_layer, numeric_circuit)?;
+        let serialized = self.generate_setup_data(circuit.clone())?;
         let digest = md5::compute(&serialized);
 
         if !dry_run {
             self.keystore()
-                .save_setup_data_for_circuit_type(
-                    if is_base_layer {
-                        ProverServiceDataKey::new(numeric_circuit, AggregationRound::BasicCircuits)
-                    } else {
-                        ProverServiceDataKey::new_recursive(numeric_circuit)
-                    },
-                    &serialized,
-                )
+                .save_setup_data_for_circuit_type(circuit, &serialized)
                 .context("save_setup_data()")?;
         } else {
             tracing::warn!("Dry run - not writing the key");
         }
         Ok(format!("{:?}", digest))
     }
 
+    /// Generate all setup keys for boojum circuits.
     fn generate_all(&self, dry_run: bool) -> anyhow::Result<HashMap<String, String>> {
-        let mut result = HashMap::new();
-
-        for numeric_circuit in
-            BaseLayerCircuitType::VM as u8..=BaseLayerCircuitType::L1MessagesHasher as u8
-        {
-            let digest = self
-                .generate_and_write_setup_data(true, numeric_circuit, dry_run)
-                .context(format!("base layer, circuit {:?}", numeric_circuit))?;
-            result.insert(format!("base_{}", numeric_circuit), digest);
-        }
-
-        for numeric_circuit in ZkSyncRecursionLayerStorageType::SchedulerCircuit as u8
-            ..=ZkSyncRecursionLayerStorageType::LeafLayerCircuitForL1MessagesHasher as u8
-        {
-            let digest = self
-                .generate_and_write_setup_data(false, numeric_circuit, dry_run)
-                .context(format!("recursive layer, circuit {:?}", numeric_circuit))?;
-            result.insert(format!("recursive_{}", numeric_circuit), digest);
-        }
-        Ok(result)
+        Ok(ProverServiceDataKey::all_boojum()
+            .iter()
+            .map(|circuit| {
+                let digest = self
+                    .generate_and_write_setup_data(circuit.clone(), dry_run)
+                    .context(circuit.name())
+                    .unwrap();
+                (circuit.name(), digest)
+            })
+            .collect())
     }
 }
 
@@ -132,13 +119,8 @@ pub struct CPUSetupDataGenerator {
 }
 
 impl SetupDataGenerator for CPUSetupDataGenerator {
-    fn generate_setup_data(
-        &self,
-        is_base_layer: bool,
-        numeric_circuit: u8,
-    ) -> anyhow::Result<Vec<u8>> {
-        let circuit_setup_data =
-            generate_setup_data_common(&self.keystore, is_base_layer, numeric_circuit)?;
+    fn generate_setup_data(&self, circuit: ProverServiceDataKey) -> anyhow::Result<Vec<u8>> {
+        let circuit_setup_data = generate_setup_data_common(&self.keystore, circuit)?;
 
         let prover_setup_data: GoldilocksProverSetupData = circuit_setup_data.into();
 
@@ -155,22 +137,17 @@ pub struct GPUSetupDataGenerator {
 }
 
 impl SetupDataGenerator for GPUSetupDataGenerator {
-    fn generate_setup_data(
-        &self,
-        is_base_layer: bool,
-        numeric_circuit: u8,
-    ) -> anyhow::Result<Vec<u8>> {
+    fn generate_setup_data(&self, circuit: ProverServiceDataKey) -> anyhow::Result<Vec<u8>> {
         #[cfg(not(feature = "gpu"))]
         {
-            let _ = (is_base_layer, numeric_circuit);
+            let _ = circuit;
             anyhow::bail!("Must compile with --gpu feature to use this option.");
         }
         #[cfg(feature = "gpu")]
         {
             let _context =
                 ProverContext::create().context("failed initializing gpu prover context")?;
-            let circuit_setup_data =
-                generate_setup_data_common(&self.keystore, is_base_layer, numeric_circuit)?;
+            let circuit_setup_data = generate_setup_data_common(&self.keystore, circuit)?;
 
             let worker = Worker::new();
             let gpu_setup_data = GpuSetup::from_setup_and_hints(