Replace shell usage

0xPolygonZero · Dec 1, 2024 · 4d48829 · 4d48829
1 parent 6937c5a
commit 4d48829
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 219 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -1,4 +1,4 @@
----   # Workflow with multiple jobs to test different parts of the project
+--- # Workflow with multiple jobs to test different parts of the project
 
 name: Continuous Integration
 
@@ -66,7 +66,7 @@ jobs:
         with:
           cache-on-failure: true
 
-      - name: build  # build separately so test logs are actually nice
+      - name: build # build separately so test logs are actually nice
         run: cargo build --tests --manifest-path trace_decoder/Cargo.toml
 
       - name: test
@@ -171,7 +171,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Run the script
-        run: ./scripts/prove_stdio.sh artifacts/witness_b19807080.json
+        run: cargo xtask prove-stdio verify artifacts/witness_b19807080.json
 
   simple_proof_witness_only:
     name: Execute bash script to generate the proof witness for a small block.
@@ -182,7 +182,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Run the script
-        run: ./scripts/prove_stdio.sh artifacts/witness_b19807080.json test_only
+        run: cargo xtask prove-stdio test artifacts/witness_b19807080.json
 
   multi_blocks_proof_regular:
     name: Execute bash script to generate and verify a proof for multiple blocks using parallel proving.
@@ -193,4 +193,4 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Run the script
-        run: ./scripts/prove_stdio.sh artifacts/witness_b3_b6.json "" use_test_config
+        run: cargo xtask prove-stdio test artifacts/witness_b3_b6.json --use-test-config
diff --git a/scripts/prove_stdio.rs b/scripts/prove_stdio.rs
@@ -27,7 +27,7 @@ pub struct ProveStdioArgs {
     block_batch_size: u32,
     /// The directory to output the proof files. If it does not exist, it will
     /// recursively be created.
-    #[arg(short = 'o', long, value_hint = ValueHint::DirPath, default_value = ".")]
+    #[arg(short = 'o', long, value_hint = ValueHint::DirPath, default_value = "./proofs")]
     output_dir: PathBuf,
 }
 
@@ -45,51 +45,6 @@ pub fn prove_via_stdio(args: ProveStdioArgs) -> anyhow::Result<()> {
 
     match args.mode {
         RunMode::Test => {
-            envs.extend([
-                ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
-                ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
-                ("CPU_CIRCUIT_SIZE", "9..20"),
-                ("KECCAK_CIRCUIT_SIZE", "7..18"),
-                ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"),
-                ("LOGIC_CIRCUIT_SIZE", "5..17"),
-                ("MEMORY_CIRCUIT_SIZE", "17..22"),
-                ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"),
-                ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"),
-                // TODO(Robin): update Poseidon ranges here and below once Kernel ASM supports
-                ("POSEIDON_CIRCUIT_SIZE", "4..8"),
-            ]);
-            let witness_filename = args
-                .input_witness_file
-                .to_str()
-                .ok_or(anyhow::anyhow!("Invalid witness file path"))?;
-            if witness_filename.contains("witness_b19807080") {
-                envs.extend([
-                    ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
-                    ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
-                    ("CPU_CIRCUIT_SIZE", "9..20"),
-                    ("KECCAK_CIRCUIT_SIZE", "7..18"),
-                    ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"),
-                    ("LOGIC_CIRCUIT_SIZE", "5..17"),
-                    ("MEMORY_CIRCUIT_SIZE", "17..22"),
-                    ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"),
-                    ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"),
-                    ("POSEIDON_CIRCUIT_SIZE", "4..8"),
-                ]);
-            } else if witness_filename.contains("witness_b3_b6") {
-                envs.extend([
-                    ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
-                    ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
-                    ("CPU_CIRCUIT_SIZE", "10..20"),
-                    ("KECCAK_CIRCUIT_SIZE", "4..13"),
-                    ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..9"),
-                    ("LOGIC_CIRCUIT_SIZE", "4..14"),
-                    ("MEMORY_CIRCUIT_SIZE", "17..22"),
-                    ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..18"),
-                    ("MEMORY_AFTER_CIRCUIT_SIZE", "7..8"),
-                    ("POSEIDON_CIRCUIT_SIZE", "4..8"),
-                ]);
-            }
-
             let mut cmd = prove_command(args, envs)?;
             let status = cmd.spawn()?.wait()?;
             ensure!(status.success(), "command failed with {}", status);
@@ -105,7 +60,9 @@ pub fn prove_via_stdio(args: ProveStdioArgs) -> anyhow::Result<()> {
             ensure!(status.success(), "command failed with {}", status);
 
             // Construct the command to run.
+            add_verify_envs(&args, &mut envs)?;
             let mut cmd = prove_command(args, envs)?;
+
             // Time the proving.
             let start = std::time::Instant::now();
             let status = cmd.spawn()?.wait()?;
@@ -117,6 +74,55 @@ pub fn prove_via_stdio(args: ProveStdioArgs) -> anyhow::Result<()> {
     }
 }
 
+fn add_verify_envs(args: &ProveStdioArgs, envs: &mut Vec<(&str, &str)>) -> anyhow::Result<()> {
+    let witness_filename = args
+        .input_witness_file
+        .to_str()
+        .ok_or(anyhow::anyhow!("Invalid witness file path"))?;
+    if witness_filename.contains("witness_b19807080") {
+        envs.extend([
+            ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
+            ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
+            ("CPU_CIRCUIT_SIZE", "9..20"),
+            ("KECCAK_CIRCUIT_SIZE", "7..18"),
+            ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"),
+            ("LOGIC_CIRCUIT_SIZE", "5..17"),
+            ("MEMORY_CIRCUIT_SIZE", "17..22"),
+            ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"),
+            ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"),
+            ("POSEIDON_CIRCUIT_SIZE", "4..8"),
+        ]);
+    } else if witness_filename.contains("witness_b3_b6") {
+        envs.extend([
+            ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
+            ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
+            ("CPU_CIRCUIT_SIZE", "10..20"),
+            ("KECCAK_CIRCUIT_SIZE", "4..13"),
+            ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..9"),
+            ("LOGIC_CIRCUIT_SIZE", "4..14"),
+            ("MEMORY_CIRCUIT_SIZE", "17..22"),
+            ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..18"),
+            ("MEMORY_AFTER_CIRCUIT_SIZE", "7..8"),
+            ("POSEIDON_CIRCUIT_SIZE", "4..8"),
+        ]);
+    } else {
+        envs.extend([
+            ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
+            ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
+            ("CPU_CIRCUIT_SIZE", "9..20"),
+            ("KECCAK_CIRCUIT_SIZE", "7..18"),
+            ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"),
+            ("LOGIC_CIRCUIT_SIZE", "5..17"),
+            ("MEMORY_CIRCUIT_SIZE", "17..22"),
+            ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"),
+            ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"),
+            // TODO(Robin): update Poseidon ranges here and below once Kernel ASM supports
+            ("POSEIDON_CIRCUIT_SIZE", "4..8"),
+        ]);
+    }
+    Ok(())
+}
+
 fn prove_command(args: ProveStdioArgs, envs: Vec<(&str, &str)>) -> anyhow::Result<Command> {
     let witness_file = File::open(&args.input_witness_file)?;
     let mut cmd = Command::new("cargo");
@@ -139,10 +145,10 @@ fn prove_command(args: ProveStdioArgs, envs: Vec<(&str, &str)>) -> anyhow::Resul
         args.output_dir
             .to_str()
             .ok_or(anyhow::anyhow!("Invalid output dir path"))?,
-        "stdio",
     ]);
     if args.use_test_config {
         cmd.arg("--use-test-config");
     }
+    cmd.arg("stdio");
     Ok(cmd)
 }
diff --git a/scripts/prove_stdio.sh b/scripts/prove_stdio.sh