Implement on-demand recursive circuit table loading (#21)

* Implement on-demand recursive circuit table loading * Update common/src/prover_state/mod.rs Co-authored-by: Robin Salen <[email protected]> * Update common/src/prover_state/mod.rs Co-authored-by: Robin Salen <[email protected]> * Update common/src/prover_state/persistence.rs Co-authored-by: Robin Salen <[email protected]> * address PR feedback * Make table load strategy configurable from CLI --------- Co-authored-by: Robin Salen <[email protected]>
0xPolygonZero · Feb 23, 2024 · 2b37fbd · 2b37fbd
1 parent 4c66002
commit 2b37fbd
Show file tree

Hide file tree

Showing 10 changed files with 602 additions and 217 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/common/Cargo.toml b/common/Cargo.toml
@@ -15,3 +15,5 @@ proof_gen = { workspace = true }
 plonky2 = { workspace = true }
 evm_arithmetization = { workspace = true }
 clap = { workspace = true }
+anyhow = { workspace = true }
+trace_decoder = { workspace = true }
diff --git a/common/src/prover_state/circuit.rs b/common/src/prover_state/circuit.rs
@@ -11,7 +11,9 @@ use proof_gen::types::AllRecursiveCircuits;
 use crate::parsing::{parse_range, RangeParseError};
 
 /// Number of tables defined in plonky2.
-const NUM_TABLES: usize = 7;
+///
+/// TODO: This should be made public in the evm_arithmetization crate.
+pub(crate) const NUM_TABLES: usize = 7;
 
 /// New type wrapper for [`Range`] that implements [`FromStr`] and [`Display`].
 ///
@@ -111,6 +113,19 @@ impl Circuit {
             Circuit::Memory => "memory",
         }
     }
+
+    /// Get the circuit name as a short str literal.
+    pub const fn as_short_str(&self) -> &'static str {
+        match self {
+            Circuit::Arithmetic => "a",
+            Circuit::BytePacking => "bp",
+            Circuit::Cpu => "c",
+            Circuit::Keccak => "k",
+            Circuit::KeccakSponge => "ks",
+            Circuit::Logic => "l",
+            Circuit::Memory => "m",
+        }
+    }
 }
 
 impl From<usize> for Circuit {
@@ -128,11 +143,27 @@ impl From<usize> for Circuit {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct CircuitConfig {
     circuits: [Range<usize>; NUM_TABLES],
 }
 
+impl std::ops::Index<usize> for CircuitConfig {
+    type Output = Range<usize>;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.circuits[index]
+    }
+}
+
+impl std::ops::Index<Circuit> for CircuitConfig {
+    type Output = Range<usize>;
+
+    fn index(&self, index: Circuit) -> &Self::Output {
+        &self.circuits[index as usize]
+    }
+}
+
 impl Default for CircuitConfig {
     fn default() -> Self {
         Self {
@@ -176,16 +207,8 @@ impl CircuitConfig {
     /// Get a unique string representation of the config.
     pub fn get_configuration_digest(&self) -> String {
         self.enumerate()
-            .map(|(circuit, range)| match circuit {
-                Circuit::Arithmetic => format!("a_{}-{}", range.start, range.end),
-                Circuit::BytePacking => format!("b_p_{}-{}", range.start, range.end),
-                Circuit::Cpu => format!("c_{}-{}", range.start, range.end),
-                Circuit::Keccak => format!("k_{}-{}", range.start, range.end),
-                Circuit::KeccakSponge => {
-                    format!("k_s_{}-{}", range.start, range.end)
-                }
-                Circuit::Logic => format!("l_{}-{}", range.start, range.end),
-                Circuit::Memory => format!("m_{}-{}", range.start, range.end),
+            .map(|(circuit, range)| {
+                format!("{}_{}-{}", circuit.as_short_str(), range.start, range.end)
             })
             .fold(String::new(), |mut acc, s| {
                 if !acc.is_empty() {

diff --git a/common/src/prover_state/cli.rs b/common/src/prover_state/cli.rs
@@ -1,10 +1,12 @@
 //! CLI arguments for constructing a [`CircuitConfig`], which can be used to
 //! construct table circuits.
-use clap::Args;
+use std::fmt::Display;
+
+use clap::{Args, ValueEnum};
 
 use super::{
     circuit::{Circuit, CircuitConfig, CircuitSize},
-    CircuitPersistence, ProverStateConfig,
+    ProverStateManager, TableLoadStrategy,
 };
 
 /// The help heading for the circuit arguments.
@@ -21,13 +23,42 @@ fn circuit_arg_desc(circuit_name: &str) -> String {
     format!("The min/max size for the {circuit_name} table circuit.")
 }
 
+/// Specifies whether to persist the processed circuits.
+#[derive(Debug, Clone, Copy, ValueEnum)]
+pub enum CircuitPersistence {
+    /// Do not persist the processed circuits.
+    None,
+    /// Persist the processed circuits to disk.
+    Disk,
+}
+
+impl CircuitPersistence {
+    pub fn with_load_strategy(self, load_strategy: TableLoadStrategy) -> super::CircuitPersistence {
+        match self {
+            CircuitPersistence::None => super::CircuitPersistence::None,
+            CircuitPersistence::Disk => super::CircuitPersistence::Disk(load_strategy),
+        }
+    }
+}
+
+impl Display for CircuitPersistence {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CircuitPersistence::None => write!(f, "none"),
+            CircuitPersistence::Disk => write!(f, "disk"),
+        }
+    }
+}
+
 /// Macro for generating the [`CliCircuitConfig`] struct.
 macro_rules! gen_prover_state_config {
     ($($name:ident: $circuit:expr),*) => {
         #[derive(Args, Debug)]
         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(
@@ -73,16 +104,16 @@ impl CliProverStateConfig {
         config
     }
 
-    pub fn into_prover_state_config(self) -> ProverStateConfig {
-        ProverStateConfig {
-            persistence: self.persistence,
+    pub fn into_prover_state_manager(self) -> ProverStateManager {
+        ProverStateManager {
+            persistence: self.persistence.with_load_strategy(self.load_strategy),
             circuit_config: self.into_circuit_config(),
         }
     }
 }
 
-impl From<CliProverStateConfig> for ProverStateConfig {
-    fn from(item: CliProverStateConfig) -> Self {
-        item.into_prover_state_config()
+impl From<CliProverStateConfig> for ProverStateManager {
+    fn from(config: CliProverStateConfig) -> Self {
+        config.into_prover_state_manager()
     }
 }