From 85d641287813a4c95fee2d09efcb6b26d444e953 Mon Sep 17 00:00:00 2001
From: Zhang Junyu <junyu92@gmail.com>
Date: Fri, 28 Jun 2024 09:32:10 +0000
Subject: [PATCH] feat: support for skipping first n slices proving

---
 crates/cli/src/app_builder.rs | 25 +++++++++++++++++++++++--
 crates/cli/src/command.rs     |  2 ++
 crates/cli/src/config.rs      | 11 ++++++++++-
 crates/cli/src/main.rs        |  1 +
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/crates/cli/src/app_builder.rs b/crates/cli/src/app_builder.rs
index bc21fc0d7..4ce9c0160 100644
--- a/crates/cli/src/app_builder.rs
+++ b/crates/cli/src/app_builder.rs
@@ -134,6 +134,20 @@ impl ArgBuilder<bool> for MockTestArg {
     }
 }
 
+struct SkipArg;
+impl ArgBuilder<usize> for SkipArg {
+    fn builder() -> Arg<'static> {
+        arg!(--skip [SKIP_SIZE] "Skip first SKIP_SIZE slice(s) proving")
+            .default_value("0")
+            .value_parser(value_parser!(usize))
+            .multiple_values(false)
+    }
+
+    fn parse(matches: &ArgMatches) -> usize {
+        matches.get_one("skip").cloned().unwrap_or_default()
+    }
+}
+
 fn setup_command() -> Command<'static> {
     let command = Command::new("setup")
         .about("Setup a new zkWasm circuit for provided Wasm image")
@@ -180,7 +194,7 @@ fn dry_run_command() -> Command<'static> {
 }
 
 fn prove_command() -> Command<'static> {
-    Command::new("prove")
+    let command = Command::new("prove")
         .about("Execute the Wasm image and generate a proof")
         .arg(WasmImageArg::builder())
         .arg(PublicInputsArg::builder())
@@ -189,7 +203,13 @@ fn prove_command() -> Command<'static> {
         .arg(ContextOutputArg::builder())
         .arg(OutputDirArg::builder())
         .arg(MockTestArg::builder())
-        .arg(FileBackendArg::builder())
+        .arg(FileBackendArg::builder());
+
+    if cfg!(feature = "continuation") {
+        command.arg(SkipArg::builder())
+    } else {
+        command
+    }
 }
 
 fn verify_command() -> Command<'static> {
@@ -259,6 +279,7 @@ impl From<&ArgMatches> for ProveArg {
             running_arg: val.into(),
             mock_test: MockTestArg::parse(val),
             file_backend: FileBackendArg::parse(val),
+            skip: SkipArg::parse(val),
         }
     }
 }
diff --git a/crates/cli/src/command.rs b/crates/cli/src/command.rs
index b8cecaeaf..bc17a3b9f 100644
--- a/crates/cli/src/command.rs
+++ b/crates/cli/src/command.rs
@@ -261,6 +261,8 @@ pub(crate) struct ProveArg {
     pub(crate) running_arg: RunningArg,
     pub(crate) mock_test: bool,
     pub(crate) file_backend: bool,
+    // skip first n slice(s) proving.
+    pub(crate) skip: usize,
 }
 
 /// Verify the proof.
diff --git a/crates/cli/src/config.rs b/crates/cli/src/config.rs
index b9ecd5da0..e5c0ac988 100644
--- a/crates/cli/src/config.rs
+++ b/crates/cli/src/config.rs
@@ -245,6 +245,7 @@ impl Config {
         context_output_filename: Option<String>,
         mock_test: bool,
         table_backend: TraceBackend,
+        skip: usize,
     ) -> anyhow::Result<()> {
         let mut cached_proving_key = None;
 
@@ -321,7 +322,15 @@ impl Config {
 
         let progress_bar = ProgressBar::new(tables.execution_tables.etable.len() as u64);
 
-        let mut slices = Slices::new(self.k, tables)?.enumerate().peekable();
+        if skip != 0 {
+            progress_bar.inc(skip as u64);
+            println!("skip first {} slice(s)", skip);
+        }
+
+        let mut slices = Slices::new(self.k, tables)?
+            .enumerate()
+            .skip(skip)
+            .peekable();
         while let Some((index, circuit)) = slices.next() {
             let circuit = circuit?;
 
diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
index ecc5442d4..d0fae44e2 100644
--- a/crates/cli/src/main.rs
+++ b/crates/cli/src/main.rs
@@ -155,6 +155,7 @@ fn main() -> Result<()> {
                 arg.running_arg.context_output,
                 arg.mock_test,
                 trace_backend,
+                arg.skip,
             )?;
         }
         Subcommands::Verify(arg) => {