From 699b74ac5b032113270a2f419ef192bdb7fc0857 Mon Sep 17 00:00:00 2001
From: Leo <leo@powdrlabs.com>
Date: Fri, 1 Nov 2024 14:55:04 +0100
Subject: [PATCH] fix chunk artifact output (#2013)

Extracted from https://github.com/powdr-labs/powdr/pull/2012
---
 examples/keccak/Cargo.toml       |  2 +-
 examples/keccak/guest/Cargo.toml |  4 ++--
 examples/keccak/src/main.rs      |  1 +
 powdr/src/lib.rs                 |  5 ++++-
 riscv/src/continuations.rs       | 14 ++++++++++++--
 5 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/examples/keccak/Cargo.toml b/examples/keccak/Cargo.toml
index 632029264..cf936b028 100644
--- a/examples/keccak/Cargo.toml
+++ b/examples/keccak/Cargo.toml
@@ -8,7 +8,7 @@ default = []
 simd = ["powdr/plonky3-simd"]
 
 [dependencies]
-powdr = { git = "https://github.com/powdr-labs/powdr", features = ["plonky3"] }
+powdr = { git = "https://github.com/powdr-labs/powdr", tag = "v0.1.1", features = ["plonky3"] }
 hex = "0.4"
 serde = { version = "1.0", default-features = false, features = [
   "alloc",
diff --git a/examples/keccak/guest/Cargo.toml b/examples/keccak/guest/Cargo.toml
index 50f884a9b..20faea95f 100644
--- a/examples/keccak/guest/Cargo.toml
+++ b/examples/keccak/guest/Cargo.toml
@@ -4,8 +4,8 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-powdr-riscv-syscalls = { git = "https://github.com/powdr-labs/powdr" }
-powdr-riscv-runtime = { git = "https://github.com/powdr-labs/powdr", features = [
+powdr-riscv-syscalls = { git = "https://github.com/powdr-labs/powdr", tag = "v0.1.1" }
+powdr-riscv-runtime = { git = "https://github.com/powdr-labs/powdr", tag = "v0.1.1", features = [
   "std",
 ] }
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
diff --git a/examples/keccak/src/main.rs b/examples/keccak/src/main.rs
index 06f0d73d4..ae3e16cbc 100644
--- a/examples/keccak/src/main.rs
+++ b/examples/keccak/src/main.rs
@@ -19,6 +19,7 @@ fn main() {
     let mut session = Session::builder()
         .guest_path("./guest")
         .out_path("powdr-target")
+        .chunk_size_log2(18)
         .build()
         .write(1, &challenge)
         .write(2, &preimg);
diff --git a/powdr/src/lib.rs b/powdr/src/lib.rs
index 3c8111b2e..72d2ecaa2 100644
--- a/powdr/src/lib.rs
+++ b/powdr/src/lib.rs
@@ -281,7 +281,10 @@ pub fn prove(pipeline: &mut Pipeline<GoldilocksField>) {
 
     pipeline.rollback_from_witness();
 
-    println!("Running witness and proof generation for all chunks...");
+    println!(
+        "Running witness and proof generation for {} chunks...",
+        bootloader_inputs.bootloader_inputs.len()
+    );
     let start = Instant::now();
     riscv::continuations::rust_continuations(pipeline, generate_proof, bootloader_inputs).unwrap();
     let duration = start.elapsed();
diff --git a/riscv/src/continuations.rs b/riscv/src/continuations.rs
index cae0b2f75..c3ecbc0f7 100644
--- a/riscv/src/continuations.rs
+++ b/riscv/src/continuations.rs
@@ -92,9 +92,11 @@ where
         .map(
             |(i, (bootloader_inputs, start_of_shutdown_routine))| -> Result<(), E> {
                 log::info!("\nRunning chunk {} / {}...", i + 1, num_chunks);
-                if let Some(parent_dir) = pipeline.output_dir() {
-                    let force_overwrite = pipeline.is_force_overwrite();
 
+                let parent_dir = pipeline.output_dir().clone();
+                let force_overwrite = pipeline.is_force_overwrite();
+
+                if let Some(parent_dir) = parent_dir.clone() {
                     let chunk_dir = parent_dir.join(format!("chunk_{i}"));
                     create_dir_all(&chunk_dir).unwrap();
 
@@ -106,6 +108,9 @@ where
                     }
                     hard_link(parent_dir.join("constants.bin"), link_to_consts).unwrap();
 
+                    // The output directory is set here to output witness and proof artifacts
+                    // inside the chunk directory.
+                    // TODO This is hacky and should be improved.
                     pipeline.set_output(chunk_dir, force_overwrite)
                 }
 
@@ -141,6 +146,11 @@ where
                     ),
                 ]);
                 pipeline_callback(pipeline)?;
+
+                if let Some(original_dir) = parent_dir {
+                    pipeline.set_output(original_dir, force_overwrite);
+                }
+
                 Ok(())
             },
         )