Add setup keys step to init

zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs

@@ -8,7 +8,7 @@ use url::Url;
 use xshell::Shell;
 use zksync_config::configs::fri_prover::CloudConnectionMode;
 
-use super::init_bellman_cuda::InitBellmanCudaArgs;
+use super::{init_bellman_cuda::InitBellmanCudaArgs, setup_keys::SetupKeysArgs};
 use crate::{
     commands::prover::gcs::get_project_ids,
     consts::{DEFAULT_CREDENTIALS_FILE, DEFAULT_PROOF_STORE_DIR},
@@ -18,25 +18,24 @@ use crate::{
         MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT, MSG_CREATE_GCS_BUCKET_NAME_PROMTP,
         MSG_CREATE_GCS_BUCKET_PROJECT_ID_NO_PROJECTS_PROMPT,
         MSG_CREATE_GCS_BUCKET_PROJECT_ID_PROMPT, MSG_CREATE_GCS_BUCKET_PROMPT,
-        MSG_DOWNLOAD_SETUP_KEY_PROMPT, MSG_GETTING_PROOF_STORE_CONFIG,
+        MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT, MSG_GETTING_PROOF_STORE_CONFIG,
         MSG_GETTING_PUBLIC_STORE_CONFIG, MSG_PROOF_STORE_CONFIG_PROMPT, MSG_PROOF_STORE_DIR_PROMPT,
         MSG_PROOF_STORE_GCS_BUCKET_BASE_URL_ERR, MSG_PROOF_STORE_GCS_BUCKET_BASE_URL_PROMPT,
         MSG_PROOF_STORE_GCS_CREDENTIALS_FILE_PROMPT, MSG_PROVER_DB_NAME_HELP,
-        MSG_PROVER_DB_URL_HELP, MSG_SAVE_TO_PUBLIC_BUCKET_PROMPT, MSG_SETUP_KEY_PATH_PROMPT,
+        MSG_PROVER_DB_URL_HELP, MSG_SAVE_TO_PUBLIC_BUCKET_PROMPT,
+        MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT, MSG_SETUP_KEYS_PROMPT,
         MSG_USE_DEFAULT_DATABASES_HELP,
     },
 };
 
-#[derive(Debug, Clone, Serialize, Deserialize, Parser, Default)]
+#[derive(Debug, Clone, Parser, Default)]
 pub struct ProverInitArgs {
     // Proof store object
     #[clap(long)]
     pub proof_store_dir: Option<String>,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub proof_store_gcs_config: ProofStorageGCSTmp,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub create_gcs_bucket_config: ProofStorageGCSCreateBucketTmp,
 
     // Public store object
@@ -45,20 +44,21 @@ pub struct ProverInitArgs {
     #[clap(long)]
     pub public_store_dir: Option<String>,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub public_store_gcs_config: PublicStorageGCSTmp,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub public_create_gcs_bucket_config: PublicStorageGCSCreateBucketTmp,
 
     // Bellman cuda
     #[clap(flatten)]
-    #[serde(flatten)]
     pub bellman_cuda_config: InitBellmanCudaArgs,
 
     #[clap(flatten)]
-    #[serde(flatten)]
-    pub setup_key_config: SetupKeyConfigTmp,
+    pub setup_compressor_key_config: SetupCompressorKeyConfigTmp,
+
+    #[clap(flatten)]
+    pub setup_keys_args: SetupKeysArgs,
+    #[clap(long, default_missing_value = "true", num_args = 0..=1)]
+    pub setup_keys: Option<bool>,
 
     #[clap(long)]
     pub setup_database: Option<bool>,
@@ -137,7 +137,7 @@ pub struct PublicStorageGCSCreateBucketTmp {
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, Parser, Default)]
-pub struct SetupKeyConfigTmp {
+pub struct SetupCompressorKeyConfigTmp {
     #[clap(long)]
     pub download_key: Option<bool>,
     #[clap(long)]
@@ -172,7 +172,7 @@ pub enum ProofStorageConfig {
 }
 
 #[derive(Debug, Clone)]
-pub struct SetupKeyConfig {
+pub struct SetupCompressorKeyConfig {
     pub download_key: bool,
     pub setup_key_path: String,
 }
@@ -187,7 +187,8 @@ pub struct ProverDatabaseConfig {
 pub struct ProverInitArgsFinal {
     pub proof_store: ProofStorageConfig,
     pub public_store: Option<ProofStorageConfig>,
-    pub setup_key_config: SetupKeyConfig,
+    pub setup_compressor_key_config: SetupCompressorKeyConfig,
+    pub setup_keys: Option<SetupKeysArgs>,
     pub bellman_cuda_config: InitBellmanCudaArgs,
     pub cloud_type: CloudConnectionMode,
     pub database_config: Option<ProverDatabaseConfig>,
@@ -202,15 +203,18 @@ impl ProverInitArgs {
     ) -> anyhow::Result<ProverInitArgsFinal> {
         let proof_store = self.fill_proof_storage_values_with_prompt(shell)?;
         let public_store = self.fill_public_storage_values_with_prompt(shell)?;
-        let setup_key_config = self.fill_setup_key_values_with_prompt(setup_key_path);
+        let setup_compressor_key_config =
+            self.fill_setup_compressor_key_values_with_prompt(setup_key_path);
         let bellman_cuda_config = self.fill_bellman_cuda_values_with_prompt()?;
         let cloud_type = self.get_cloud_type_with_prompt();
         let database_config = self.fill_database_values_with_prompt(chain_config);
+        let setup_keys = self.fill_setup_keys_values_with_prompt();
 
         Ok(ProverInitArgsFinal {
             proof_store,
             public_store,
-            setup_key_config,
+            setup_compressor_key_config,
+            setup_keys,
             bellman_cuda_config,
             cloud_type,
             database_config,
@@ -336,32 +340,49 @@ impl ProverInitArgs {
         }
     }
 
-    fn fill_setup_key_values_with_prompt(&self, setup_key_path: &str) -> SetupKeyConfig {
+    fn fill_setup_compressor_key_values_with_prompt(
+        &self,
+        setup_key_path: &str,
+    ) -> SetupCompressorKeyConfig {
         let download_key = self
             .clone()
-            .setup_key_config
+            .setup_compressor_key_config
             .download_key
             .unwrap_or_else(|| {
-                PromptConfirm::new(MSG_DOWNLOAD_SETUP_KEY_PROMPT)
+                PromptConfirm::new(MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT)
                     .default(true)
                     .ask()
             });
         let setup_key_path = self
             .clone()
-            .setup_key_config
+            .setup_compressor_key_config
             .setup_key_path
             .unwrap_or_else(|| {
-                Prompt::new(MSG_SETUP_KEY_PATH_PROMPT)
+                Prompt::new(MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT)
                     .default(setup_key_path)
                     .ask()
             });
 
-        SetupKeyConfig {
+        SetupCompressorKeyConfig {
             download_key,
             setup_key_path,
         }
     }
 
+    fn fill_setup_keys_values_with_prompt(&self) -> Option<SetupKeysArgs> {
+        let args = self.setup_keys_args.clone();
+
+        if self.setup_keys.unwrap_or_else(|| {
+            PromptConfirm::new(MSG_SETUP_KEYS_PROMPT)
+                .default(true)
+                .ask()
+        }) {
+            Some(args)
+        } else {
+            None
+        }
+    }
+
     fn partial_create_gcs_bucket_config_provided(
         &self,
         bucket_name: Option<String>,

zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs

@@ -2,7 +2,7 @@ use clap::{Parser, ValueEnum};
 use common::PromptSelect;
 use strum::{EnumIter, IntoEnumIterator};
 
-use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_HELP, MSG_SETUP_KEYS_REGION_PROMPT};
+use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT, MSG_SETUP_KEYS_REGION_PROMPT};
 
 #[derive(Debug, Clone, Parser, Default)]
 pub struct SetupKeysArgs {
@@ -33,9 +33,9 @@ pub enum Region {
 
 impl SetupKeysArgs {
     pub fn fill_values_with_prompt(self) -> SetupKeysArgsFinal {
-        let mode = self
-            .mode
-            .unwrap_or_else(|| PromptSelect::new(MSG_SETUP_KEYS_DOWNLOAD_HELP, Mode::iter()).ask());
+        let mode = self.mode.unwrap_or_else(|| {
+            PromptSelect::new(MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT, Mode::iter()).ask()
+        });
 
         if mode == Mode::Download {
             let region = self.region.unwrap_or_else(|| {

zk_toolbox/crates/zk_inception/src/commands/prover/init.rs

@@ -21,6 +21,7 @@ use super::{
     args::init::{ProofStorageConfig, ProverInitArgs},
     gcs::create_gcs_bucket,
     init_bellman_cuda::run as init_bellman_cuda,
+    setup_keys,
     utils::get_link_to_prover,
 };
 use crate::{
@@ -55,14 +56,18 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<(
     let proof_object_store_config = get_object_store_config(shell, Some(args.proof_store))?;
     let public_object_store_config = get_object_store_config(shell, args.public_store)?;
 
-    if args.setup_key_config.download_key {
+    if args.setup_compressor_key_config.download_key {
         download_setup_key(
             shell,
             &general_config,
-            &args.setup_key_config.setup_key_path,
+            &args.setup_compressor_key_config.setup_key_path,
         )?;
     }
 
+    if let Some(args) = args.setup_keys {
+        setup_keys::run(args, &shell).await?;
+    }
+
     let mut prover_config = general_config
         .prover_config
         .expect(MSG_PROVER_CONFIG_NOT_FOUND_ERR);
@@ -81,7 +86,7 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<(
     let mut proof_compressor_config = general_config
         .proof_compressor_config
         .expect(MSG_PROOF_COMPRESSOR_CONFIG_NOT_FOUND_ERR);
-    proof_compressor_config.universal_setup_path = args.setup_key_config.setup_key_path;
+    proof_compressor_config.universal_setup_path = args.setup_compressor_key_config.setup_key_path;
     general_config.proof_compressor_config = Some(proof_compressor_config);
 
     chain_config.save_general_config(&general_config)?;

zk_toolbox/crates/zk_inception/src/messages.rs

@@ -5,7 +5,7 @@ use ethers::{
     utils::format_ether,
 };
 
-pub(super) const MSG_SETUP_KEYS_DOWNLOAD_HELP: &str =
+pub(super) const MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT: &str =
     "Do you want to download the setup keys or generate them?";
 pub(super) const MSG_SETUP_KEYS_REGION_PROMPT: &str =
     "From which region you want setup keys to be downloaded?";
@@ -302,8 +302,9 @@ pub(super) const MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT: &str = "What location do
 pub(super) const MSG_PROOF_COMPRESSOR_CONFIG_NOT_FOUND_ERR: &str =
     "Proof compressor config not found";
 pub(super) const MSG_DOWNLOADING_SETUP_KEY_SPINNER: &str = "Downloading setup key...";
-pub(super) const MSG_DOWNLOAD_SETUP_KEY_PROMPT: &str = "Do you want to download the setup key?";
-pub(super) const MSG_SETUP_KEY_PATH_PROMPT: &str = "Provide the path to the setup key:";
+pub(super) const MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT: &str =
+    "Do you want to download the setup key for compressor?";
+pub(super) const MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT: &str = "Provide the path to the setup key:";
 pub(super) const MSG_GETTING_GCP_PROJECTS_SPINNER: &str = "Getting GCP projects...";
 pub(super) const MSG_GETTING_PROOF_STORE_CONFIG: &str = "Getting proof store configuration...";
 pub(super) const MSG_GETTING_PUBLIC_STORE_CONFIG: &str = "Getting public store configuration...";
@@ -325,6 +326,7 @@ pub(super) const MSG_BELLMAN_CUDA_SELECTION_CLONE: &str = "Clone for me (recomme
 pub(super) const MSG_BELLMAN_CUDA_SELECTION_PATH: &str = "I have the code already";
 pub(super) const MSG_CLOUD_TYPE_PROMPT: &str = "Select the cloud connection mode:";
 pub(super) const MSG_THREADS_PROMPT: &str = "Provide the number of threads:";
+pub(super) const MSG_SETUP_KEYS_PROMPT: &str = "Do you want to setup keys?";
 
 pub(super) fn msg_bucket_created(bucket_name: &str) -> String {
     format!("Bucket created successfully with url: gs://{bucket_name}")