From 52d5a67a1b8d8cd9e2741aca2ece88ee494583bf Mon Sep 17 00:00:00 2001
From: "xingqiang.yuan" <xingqiang.yuan@okg.com>
Date: Thu, 16 Nov 2023 19:10:14 +0800
Subject: [PATCH] update

---
 crates/zkwasm/src/loader/mod.rs                | 6 +-----
 crates/zkwasm/src/runtime/wasmi_interpreter.rs | 6 +++---
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/crates/zkwasm/src/loader/mod.rs b/crates/zkwasm/src/loader/mod.rs
index 0e6b9cc13..dad7d688e 100644
--- a/crates/zkwasm/src/loader/mod.rs
+++ b/crates/zkwasm/src/loader/mod.rs
@@ -181,7 +181,7 @@ impl<E: MultiMillerLoop> ZkWasmLoader<E> {
         compiled_module.dry_run(&mut env)
     }
 
-    pub fn dry_run_trace_count(&self, arg: ExecutionArg) -> Result<Option<RuntimeValue>> {
+    pub fn dry_run_trace_count(&self, arg: ExecutionArg) -> Result<(Option<RuntimeValue>, usize)> {
         let (mut env, _) = HostEnv::new_with_full_foreign_plugins(
             arg.public_inputs,
             arg.private_inputs,
@@ -356,7 +356,3 @@ mod tests {
         }
     }
 }
-
-pub fn get_trace_count() -> usize {
-    return wasmi::tracer::phantom::get_trace_count();
-}
\ No newline at end of file
diff --git a/crates/zkwasm/src/runtime/wasmi_interpreter.rs b/crates/zkwasm/src/runtime/wasmi_interpreter.rs
index bdfc528ad..775cf5011 100644
--- a/crates/zkwasm/src/runtime/wasmi_interpreter.rs
+++ b/crates/zkwasm/src/runtime/wasmi_interpreter.rs
@@ -36,7 +36,7 @@ impl WasmRuntimeIO {
 pub trait Execution<R> {
     fn dry_run<E: Externals>(self, externals: &mut E) -> Result<Option<R>>;
 
-    fn dry_run_trace_count<E: Externals>(self, externals: &mut E) -> Result<Option<RuntimeValue>>;
+    fn dry_run_trace_count<E: Externals>(self, externals: &mut E) -> Result<(Option<RuntimeValue>, usize)>;
 
     fn run<E: Externals>(
         self,
@@ -56,14 +56,14 @@ impl Execution<RuntimeValue>
         Ok(result)
     }
 
-    fn dry_run_trace_count<E: Externals>(self, externals: &mut E) -> Result<Option<RuntimeValue>> {
+    fn dry_run_trace_count<E: Externals>(self, externals: &mut E) -> Result<(Option<RuntimeValue>, usize)> {
         let instance = self.instance.run_start(externals).unwrap();
 
         let result = instance.invoke_export_trace_count(
             &self.entry, &[], externals, self.tracer.clone()
         )?;
 
-        Ok(result)
+        Ok((result, wasmi::tracer::phantom::get_trace_count()))
     }
 
     fn run<E: Externals>(