Make table load strategy configurable from CLI

common/src/prover_state/cli.rs

@@ -57,6 +57,8 @@ macro_rules! gen_prover_state_config {
         pub struct CliProverStateConfig {
             #[clap(long, help_heading = HEADING, default_value_t = CircuitPersistence::Disk)]
             pub persistence: CircuitPersistence,
+            #[clap(long, help_heading = HEADING, default_value_t = TableLoadStrategy::OnDemand)]
+            pub load_strategy: TableLoadStrategy,
 
             $(
                 #[clap(
@@ -102,13 +104,16 @@ impl CliProverStateConfig {
         config
     }
 
-    pub fn into_prover_state_manager(
-        self,
-        persistence: super::CircuitPersistence,
-    ) -> ProverStateManager {
+    pub fn into_prover_state_manager(self) -> ProverStateManager {
         ProverStateManager {
-            persistence,
+            persistence: self.persistence.with_load_strategy(self.load_strategy),
             circuit_config: self.into_circuit_config(),
         }
     }
 }
+
+impl From<CliProverStateConfig> for ProverStateManager {
+    fn from(config: CliProverStateConfig) -> Self {
+        config.into_prover_state_manager()
+    }
+}

common/src/prover_state/mod.rs

@@ -11,8 +11,9 @@
 //!   [`evm_arithmetization::fixed_recursive_verifier::AllRecursiveCircuits`].
 //! - Global prover state management via the [`P_STATE`] static and the
 //!   [`set_prover_state_from_config`] function.
-use std::sync::OnceLock;
+use std::{fmt::Display, sync::OnceLock};
 
+use clap::ValueEnum;
 use evm_arithmetization::{proof::AllProof, prover::prove, AllStark, StarkConfig};
 use plonky2::{
     field::goldilocks_field::GoldilocksField, plonk::config::PoseidonGoldilocksConfig,
@@ -75,7 +76,7 @@ pub fn p_manager() -> &'static ProverStateManager {
 }
 
 /// Specifies how to load the table circuits.
-#[derive(Debug, Clone, Copy, Default)]
+#[derive(Debug, Clone, Copy, Default, ValueEnum)]
 pub enum TableLoadStrategy {
     #[default]
     /// Load the circuit tables as needed for shrinking STARK proofs.
@@ -88,6 +89,15 @@ pub enum TableLoadStrategy {
     Monolithic,
 }
 
+impl Display for TableLoadStrategy {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            TableLoadStrategy::OnDemand => write!(f, "on-demand"),
+            TableLoadStrategy::Monolithic => write!(f, "monolithic"),
+        }
+    }
+}
+
 /// Specifies whether to persist the processed circuits.
 #[derive(Debug, Clone, Copy)]
 pub enum CircuitPersistence {
@@ -114,6 +124,16 @@ pub struct ProverStateManager {
 }
 
 impl ProverStateManager {
+    pub fn with_load_strategy(self, load_strategy: TableLoadStrategy) -> Self {
+        match self.persistence {
+            CircuitPersistence::None => self,
+            CircuitPersistence::Disk(_) => Self {
+                circuit_config: self.circuit_config,
+                persistence: CircuitPersistence::Disk(load_strategy),
+            },
+        }
+    }
+
     /// Load the table circuits necessary to shrink the STARK proof.
     ///
     /// [`AllProof`] provides the necessary degree bits for each circuit via the

leader/src/main.rs

@@ -36,13 +36,11 @@ async fn main() -> Result<()> {
     if let paladin::config::Runtime::InMemory = args.paladin.runtime {
         // If running in emulation mode, we'll need to initialize the prover
         // state here.
-        let persistence = args.prover_state_config.persistence;
         args.prover_state_config
-            .into_prover_state_manager(
-                // Use the monolithic load strategy for the prover state when running in
-                // emulation mode.
-                persistence.with_load_strategy(TableLoadStrategy::Monolithic),
-            )
+            .into_prover_state_manager()
+            // Use the monolithic load strategy for the prover state when running in
+            // emulation mode.
+            .with_load_strategy(TableLoadStrategy::Monolithic)
             .initialize()?;
     }
 

verifier/src/main.rs

@@ -18,7 +18,7 @@ fn main() -> Result<()> {
 
     let verifer = args
         .prover_state_config
-        .into_prover_state_manager(Default::default())
+        .into_prover_state_manager()
         .verifier()?;
 
     verifer.verify(&input)?;

worker/src/main.rs

@@ -1,6 +1,6 @@
 use anyhow::Result;
 use clap::Parser;
-use common::prover_state::{cli::CliProverStateConfig, TableLoadStrategy};
+use common::prover_state::cli::CliProverStateConfig;
 use dotenvy::dotenv;
 use ops::register;
 use paladin::runtime::WorkerRuntime;
@@ -21,9 +21,8 @@ async fn main() -> Result<()> {
     init::tracing();
     let args = Cli::parse();
 
-    let persistence = args.prover_state_config.persistence;
     args.prover_state_config
-        .into_prover_state_manager(persistence.with_load_strategy(TableLoadStrategy::OnDemand))
+        .into_prover_state_manager()
         .initialize()?;
 
     let runtime = WorkerRuntime::from_config(&args.paladin, register()).await?;