From 9f6f0cfa9e4571e7863a21af5311ffe486f81c5c Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Fri, 14 Apr 2023 23:54:44 -0700
Subject: [PATCH] Interpreter: configurable trace hook.

---
 src/interp.rs    | 36 +++++++++++++++++++++---------------
 src/ir/func.rs   | 10 ++++++++++
 src/ir/module.rs | 22 ++++++++++++++++++++++
 3 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/src/interp.rs b/src/interp.rs
index 0765bae..628f93a 100644
--- a/src/interp.rs
+++ b/src/interp.rs
@@ -12,12 +12,12 @@ mod wasi;
 const WASM_PAGE: usize = 0x1_0000; // 64KiB
 const MAX_PAGES: usize = 2048; // 2048 * 64KiB = 128MiB
 
-#[derive(Debug, Clone)]
 pub struct InterpContext {
     pub memories: PerEntity<Memory, InterpMemory>,
     pub tables: PerEntity<Table, InterpTable>,
     pub globals: PerEntity<Global, ConstVal>,
     pub fuel: u64,
+    pub trace_handler: Option<Box<dyn Fn(usize, Vec<ConstVal>) -> bool + Send>>,
 }
 
 type MultiVal = SmallVec<[ConstVal; 2]>;
@@ -28,6 +28,7 @@ pub enum InterpResult {
     Exit,
     Trap,
     OutOfFuel,
+    TraceHandlerQuit,
 }
 
 impl InterpResult {
@@ -84,6 +85,7 @@ impl InterpContext {
             tables,
             globals,
             fuel: u64::MAX,
+            trace_handler: None,
         })
     }
 
@@ -195,20 +197,24 @@ impl InterpContext {
                         smallvec![result]
                     }
                     &ValueDef::Trace(id, ref args) => {
-                        let args = args
-                            .iter()
-                            .map(|&arg| {
-                                let arg = body.resolve_alias(arg);
-                                let multivalue = frame
-                                    .values
-                                    .get(&arg)
-                                    .ok_or_else(|| format!("Unset SSA value: {}", arg))
-                                    .unwrap();
-                                assert_eq!(multivalue.len(), 1);
-                                multivalue[0]
-                            })
-                            .collect::<Vec<_>>();
-                        eprintln!("TRACE: {}: {:?}", id, &args[..]);
+                        if let Some(handler) = self.trace_handler.as_ref() {
+                            let args = args
+                                .iter()
+                                .map(|&arg| {
+                                    let arg = body.resolve_alias(arg);
+                                    let multivalue = frame
+                                        .values
+                                        .get(&arg)
+                                        .ok_or_else(|| format!("Unset SSA value: {}", arg))
+                                        .unwrap();
+                                    assert_eq!(multivalue.len(), 1);
+                                    multivalue[0]
+                                })
+                                .collect::<Vec<ConstVal>>();
+                            if !handler(id, args) {
+                                return InterpResult::TraceHandlerQuit;
+                            }
+                        }
                         smallvec![]
                     }
                     &ValueDef::None | &ValueDef::Placeholder(..) | &ValueDef::BlockParam(..) => {
diff --git a/src/ir/func.rs b/src/ir/func.rs
index c107573..acecf9e 100644
--- a/src/ir/func.rs
+++ b/src/ir/func.rs
@@ -97,6 +97,16 @@ impl<'a> FuncDecl<'a> {
             FuncDecl::None => panic!("No name for FuncDecl::None"),
         }
     }
+
+    pub fn without_orig_bytes(self) -> FuncDecl<'static> {
+        match self {
+            FuncDecl::Body(sig, name, body) => FuncDecl::Body(sig, name, body),
+            FuncDecl::Import(sig, name) => FuncDecl::Import(sig, name),
+            FuncDecl::Compiled(sig, name, func) => FuncDecl::Compiled(sig, name, func),
+            FuncDecl::None => FuncDecl::None,
+            FuncDecl::Lazy(..) => panic!("Trying to strip lifetime from lazy decl"),
+        }
+    }
 }
 
 #[derive(Clone, Debug, Default)]
diff --git a/src/ir/module.rs b/src/ir/module.rs
index 61286dc..3e1b014 100644
--- a/src/ir/module.rs
+++ b/src/ir/module.rs
@@ -145,6 +145,28 @@ impl<'a> Module<'a> {
             debug_map: DebugMap::default(),
         }
     }
+
+    pub fn without_orig_bytes(self) -> Module<'static> {
+        Module {
+            orig_bytes: &[],
+            funcs: EntityVec::from(
+                self.funcs
+                    .into_vec()
+                    .into_iter()
+                    .map(|decl| decl.without_orig_bytes())
+                    .collect::<Vec<_>>(),
+            ),
+            signatures: self.signatures,
+            globals: self.globals,
+            tables: self.tables,
+            imports: self.imports,
+            exports: self.exports,
+            memories: self.memories,
+            start_func: self.start_func,
+            debug: self.debug,
+            debug_map: self.debug_map,
+        }
+    }
 }
 
 impl<'a> Module<'a> {