From 43219c35820f513150a0150d4daf68697c11afc4 Mon Sep 17 00:00:00 2001
From: Tad Hardesty <tad@dcspark.io>
Date: Tue, 15 Apr 2025 14:04:28 -0700
Subject: [PATCH] Let exit(0) end WASM execution but continue tracing and
 proving

---
 src/wasm_ctx.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/wasm_ctx.rs b/src/wasm_ctx.rs
index ebd1435..7060ecb 100644
--- a/src/wasm_ctx.rs
+++ b/src/wasm_ctx.rs
@@ -227,7 +227,12 @@ pub trait ZKWASMCtx {
     let mut func_results = prepare_func_results(&ty);
 
     // Call the function to invoke.
-    func.call_with_trace(&mut store, &func_args, &mut func_results, tracer.clone())?;
+    match func.call_with_trace(&mut store, &func_args, &mut func_results, tracer.clone()) {
+      Ok(()) => {}
+      // Permit exit(0) to end execution early and continue proving.
+      Err(wasmi::Error::Trap(t)) if t.i32_exit_status() == Some(0) => {}
+      Err(e) => return Err(e.into()),
+    }
     tracing::debug!("wasm func res: {:#?}", func_results);
 
     // Extract the execution trace produced from WASM execution.