Merge pull request #2882 from o1-labs/martin/cannon-cli-improvements

More robust o1vm cli

o1vm/Cargo.toml

@@ -12,10 +12,6 @@ license = "Apache-2.0"
 [lib]
 path = "src/lib.rs"
 
-[[bin]]
-name = "test_optimism_preimage_read"
-path = "src/test_preimage_read.rs"
-
 [[bin]]
 name = "legacy_o1vm"
 path = "src/legacy/main.rs"
@@ -35,7 +31,7 @@ ark-ec.workspace = true
 ark-ff.workspace = true
 ark-poly.workspace = true
 base64.workspace = true
-clap.workspace = true
+clap = { workspace = true, features = ["derive"] }
 command-fds.workspace = true
 elf.workspace = true
 env_logger.workspace = true

o1vm/run-vm.sh

@@ -30,7 +30,7 @@ fi
 cargo run --bin ${BINARY_FLAVOR} \
     --all-features \
     --release \
-    -p o1vm -- \
+    -p o1vm cannon run \
     --pprof.cpu \
     --info-at "${INFO_AT:-%10000000}" \
     --snapshot-state-at "${SNAPSHOT_STATE_AT:-%10000000}" \

o1vm/src/cannon.rs

@@ -147,33 +147,36 @@ pub enum StepFrequency {
     Range(u64, Option<u64>),
 }
 
-// Simple parser for Cannon's "frequency format"
-// A frequency input is either
-// - never/always
-// - =<n> (only at step n)
-// - %<n> (every steps multiple of n)
-// - n..[m] (from n on, until m excluded if specified, until the end otherwise)
-pub fn step_frequency_parser(s: &str) -> std::result::Result<StepFrequency, String> {
-    use StepFrequency::*;
-
-    let mod_re = Regex::new(r"^%([0-9]+)").unwrap();
-    let eq_re = Regex::new(r"^=([0-9]+)").unwrap();
-    let ival_re = Regex::new(r"^([0-9]+)..([0-9]+)?").unwrap();
-
-    match s {
-        "never" => Ok(Never),
-        "always" => Ok(Always),
-        s => {
-            if let Some(m) = mod_re.captures(s) {
-                Ok(Every(m[1].parse::<u64>().unwrap()))
-            } else if let Some(m) = eq_re.captures(s) {
-                Ok(Exactly(m[1].parse::<u64>().unwrap()))
-            } else if let Some(m) = ival_re.captures(s) {
-                let lo = m[1].parse::<u64>().unwrap();
-                let hi_opt = m.get(2).map(|x| x.as_str().parse::<u64>().unwrap());
-                Ok(Range(lo, hi_opt))
-            } else {
-                Err(format!("Unknown frequency format {}", s))
+impl FromStr for StepFrequency {
+    type Err = String;
+    // Simple parser for Cannon's "frequency format"
+    // A frequency input is either
+    // - never/always
+    // - =<n> (only at step n)
+    // - %<n> (every steps multiple of n)
+    // - n..[m] (from n on, until m excluded if specified, until the end otherwise)
+    fn from_str(s: &str) -> std::result::Result<StepFrequency, String> {
+        use StepFrequency::*;
+
+        let mod_re = Regex::new(r"^%([0-9]+)").unwrap();
+        let eq_re = Regex::new(r"^=([0-9]+)").unwrap();
+        let ival_re = Regex::new(r"^([0-9]+)..([0-9]+)?").unwrap();
+
+        match s {
+            "never" => Ok(Never),
+            "always" => Ok(Always),
+            s => {
+                if let Some(m) = mod_re.captures(s) {
+                    Ok(Every(m[1].parse::<u64>().unwrap()))
+                } else if let Some(m) = eq_re.captures(s) {
+                    Ok(Exactly(m[1].parse::<u64>().unwrap()))
+                } else if let Some(m) = ival_re.captures(s) {
+                    let lo = m[1].parse::<u64>().unwrap();
+                    let hi_opt = m.get(2).map(|x| x.as_str().parse::<u64>().unwrap());
+                    Ok(Range(lo, hi_opt))
+                } else {
+                    Err(format!("Unknown frequency format {}", s))
+                }
             }
         }
     }
@@ -296,13 +299,13 @@ mod tests {
     #[test]
     fn sp_parser() {
         use StepFrequency::*;
-        assert_eq!(step_frequency_parser("never"), Ok(Never));
-        assert_eq!(step_frequency_parser("always"), Ok(Always));
-        assert_eq!(step_frequency_parser("=123"), Ok(Exactly(123)));
-        assert_eq!(step_frequency_parser("%123"), Ok(Every(123)));
-        assert_eq!(step_frequency_parser("1..3"), Ok(Range(1, Some(3))));
-        assert_eq!(step_frequency_parser("1.."), Ok(Range(1, None)));
-        assert!(step_frequency_parser("@123").is_err());
+        assert_eq!(StepFrequency::from_str("never"), Ok(Never));
+        assert_eq!(StepFrequency::from_str("always"), Ok(Always));
+        assert_eq!(StepFrequency::from_str("=123"), Ok(Exactly(123)));
+        assert_eq!(StepFrequency::from_str("%123"), Ok(Every(123)));
+        assert_eq!(StepFrequency::from_str("1..3"), Ok(Range(1, Some(3))));
+        assert_eq!(StepFrequency::from_str("1.."), Ok(Range(1, None)));
+        assert!(StepFrequency::from_str("@123").is_err());
     }
 
     // This sample is a subset taken from a Cannon-generated "meta.json" file

o1vm/src/cli/cannon.rs

@@ -0,0 +1,112 @@
+use crate::cannon::*;
+use clap::{arg, Parser, Subcommand};
+
+#[derive(Parser, Debug, Clone)]
+pub struct MipsVmConfigurationArgs {
+    #[arg(
+        long,
+        value_name = "FILE",
+        default_value = "state.json",
+        help = "initial state file"
+    )]
+    input: String,
+
+    #[arg(
+        long,
+        value_name = "FILE",
+        default_value = "out.json",
+        help = "output state file"
+    )]
+    output: String,
+
+    #[arg(
+        long,
+        value_name = "FILE",
+        default_value = "meta.json",
+        help = "metadata file"
+    )]
+    meta: String,
+
+    #[arg(
+        long = "proof-at",
+        short = 'p',
+        long,
+        value_name = "FREQ",
+        default_value = "never"
+    )]
+    proof_at: StepFrequency,
+
+    #[arg(
+        long = "proof-fmt",
+        value_name = "FORMAT",
+        default_value = "proof-%d.json"
+    )]
+    proof_fmt: String,
+
+    #[arg(
+        long = "snapshot-fmt",
+        value_name = "FORMAT",
+        default_value = "state-%d.json"
+    )]
+    snapshot_fmt: String,
+
+    #[arg(long = "stop-at", value_name = "FREQ", default_value = "never")]
+    stop_at: StepFrequency,
+
+    #[arg(long = "info-at", value_name = "FREQ", default_value = "never")]
+    info_at: StepFrequency,
+
+    #[arg(long = "pprof.cpu", action = clap::ArgAction::SetTrue)]
+    pprof_cpu: bool,
+
+    #[arg(
+        long = "snapshot-state-at",
+        value_name = "FREQ",
+        default_value = "never"
+    )]
+    snapshot_state_at: StepFrequency,
+
+    #[arg(name = "host", value_name = "HOST", help = "host program specification <host program> [host program arguments]", num_args = 1.., last = true)]
+    host: Vec<String>,
+}
+
+impl From<MipsVmConfigurationArgs> for VmConfiguration {
+    fn from(cfg: MipsVmConfigurationArgs) -> Self {
+        VmConfiguration {
+            input_state_file: cfg.input,
+            output_state_file: cfg.output,
+            metadata_file: cfg.meta,
+            proof_at: cfg.proof_at,
+            stop_at: cfg.stop_at,
+            snapshot_state_at: cfg.snapshot_state_at,
+            info_at: cfg.info_at,
+            proof_fmt: cfg.proof_fmt,
+            snapshot_fmt: cfg.snapshot_fmt,
+            pprof_cpu: cfg.pprof_cpu,
+            host: if cfg.host.is_empty() {
+                None
+            } else {
+                Some(HostProgram {
+                    name: cfg.host[0].to_string(),
+                    arguments: cfg.host[1..].to_vec(),
+                })
+            },
+        }
+    }
+}
+
+#[derive(Parser, Debug, Clone)]
+pub struct RunArgs {
+    #[arg(long = "preimage-db-dir", value_name = "PREIMAGE_DB_DIR")]
+    pub preimage_db_dir: Option<String>,
+    // it's important that vm_cfg is last in order to properly parse the host field
+    #[command(flatten)]
+    pub vm_cfg: MipsVmConfigurationArgs,
+}
+
+#[derive(Subcommand, Clone, Debug)]
+pub enum Cannon {
+    Run(RunArgs),
+    #[command(name = "test-optimism-preimage-read")]
+    TestPreimageRead(RunArgs),
+}

o1vm/src/cli/mod.rs

@@ -0,0 +1,14 @@
+use clap::Parser;
+
+pub mod cannon;
+
+#[derive(Parser, Debug, Clone)]
+#[command(
+    name = "o1vm",
+    version = "0.1",
+    about = "o1vm - a generic purpose zero-knowledge virtual machine"
+)]
+pub enum Commands {
+    #[command(subcommand)]
+    Cannon(cannon::Cannon),
+}