fix chunk artifact output (#2013)

Extracted from #2012

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",

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"] }

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);

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();

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(())
             },
         )