feat(prover): Added --recompute-if-missing option to key generator

prover/setup.sh

@@ -23,32 +23,5 @@ rm ../etc/env/base/fri_proof_compressor.toml.backup
 
 zk config compile dev
 
-
-SHOULD_GENERATE_KEYS=false
-
-for i in {1..13}
-do
-    if ! [ -f vk_setup_data_generator_server_fri/data/setup_basic_${i}_data.bin ]; then
-        SHOULD_GENERATE_KEYS=true        
-    fi
-done
-
-if ! [ -f vk_setup_data_generator_server_fri/data/setup_scheduler_data.bin ]; then
-    SHOULD_GENERATE_KEYS=true            
-fi
-
-if ! [ -f vk_setup_data_generator_server_fri/data/setup_node_data.bin ]; then
-    SHOULD_GENERATE_KEYS=true        
-fi
-
-for i in {3..15}
-do
-    if ! [ -f vk_setup_data_generator_server_fri/data/setup_leaf_${i}_data.bin ]; then
-        SHOULD_GENERATE_KEYS=true        
-    fi
-done
-
-
-if [ "$SHOULD_GENERATE_KEYS" = "true" ]; then
-    zk f cargo run $GPU_FLAG --release --bin key_generator -- $GENERATE_SK_COMMAND all
-fi
+# Update setup keys (only if they are not present)
+zk f cargo run $GPU_FLAG --release --bin key_generator -- $GENERATE_SK_COMMAND all --recompute-if-missing

prover/vk_setup_data_generator_server_fri/src/keystore.rs

@@ -1,4 +1,8 @@
-use std::{fs, fs::File, io::Read};
+use std::{
+    fs::{self, File},
+    io::Read,
+    path::Path,
+};
 
 use anyhow::Context as _;
 use circuit_definitions::{
@@ -293,6 +297,10 @@ impl Keystore {
         })
     }
 
+    pub fn is_setup_data_present(&self, key: &ProverServiceDataKey) -> bool {
+        Path::new(&self.get_file_path(key.clone(), ProverServiceDataType::SetupData)).exists()
+    }
+
     pub fn save_setup_data_for_circuit_type(
         &self,
         key: ProverServiceDataKey,

prover/vk_setup_data_generator_server_fri/src/main.rs

@@ -104,6 +104,11 @@ struct GeneratorOptions {
     #[arg(long, default_value = "false")]
     dry_run: bool,
 
+    /// If true, then generate the setup key, only if it is missing.
+    // Warning: this doesn't check the correctness of the existing file.
+    #[arg(long, default_value = "false")]
+    recompute_if_missing: bool,
+
     #[arg(long)]
     path: Option<String>,
 
@@ -170,7 +175,7 @@ fn generate_setup_keys(
 ) -> anyhow::Result<()> {
     let circuit_type = match options.circuits_type {
         CircuitSelector::All => {
-            let digests = generator.generate_all(options.dry_run)?;
+            let digests = generator.generate_all(options.dry_run, options.recompute_if_missing)?;
             tracing::info!("Setup keys md5(s):");
             print_stats(digests)?;
             return Ok(());
@@ -189,7 +194,7 @@ fn generate_setup_keys(
     };
 
     let digest = generator
-        .generate_and_write_setup_data(circuit_type, options.dry_run)
+        .generate_and_write_setup_data(circuit_type, options.dry_run, options.recompute_if_missing)
         .context("generate_setup_data()")?;
     tracing::info!("digest: {:?}", digest);
     Ok(())

prover/vk_setup_data_generator_server_fri/src/setup_data_generator.rs

@@ -85,7 +85,15 @@ pub trait SetupDataGenerator {
         &self,
         circuit: ProverServiceDataKey,
         dry_run: bool,
+        recompute_if_missing: bool,
     ) -> anyhow::Result<String> {
+        if recompute_if_missing && self.keystore().is_setup_data_present(&circuit) {
+            tracing::info!(
+                "Skipping setup computation for {:?} as file is already present.",
+                circuit.name(),
+            );
+            return Ok("Skipped".to_string());
+        }
         let serialized = self.generate_setup_data(circuit.clone())?;
         let digest = md5::compute(&serialized);
 
@@ -100,12 +108,16 @@ pub trait SetupDataGenerator {
     }
 
     /// Generate all setup keys for boojum circuits.
-    fn generate_all(&self, dry_run: bool) -> anyhow::Result<HashMap<String, String>> {
+    fn generate_all(
+        &self,
+        dry_run: bool,
+        recompute_if_missing: bool,
+    ) -> anyhow::Result<HashMap<String, String>> {
         Ok(ProverServiceDataKey::all_boojum()
             .iter()
             .map(|circuit| {
                 let digest = self
-                    .generate_and_write_setup_data(circuit.clone(), dry_run)
+                    .generate_and_write_setup_data(circuit.clone(), dry_run, recompute_if_missing)
                     .context(circuit.name())
                     .unwrap();
                 (circuit.name(), digest)