LogExecution: Optionally take a module name for the logger function (#…

…6629) --log-execution=NAME will use NAME as the module for the logger function import, rather than infer it. If the name is not provided (--log-execution as before this PR) then we will try to automatically decide which to use ("env", unless we see another module name is used, which can be the case in optimized modules).
WebAssembly · May 31, 2024 · 0c23394 · 0c23394
1 parent 18bb1f5
commit 0c23394
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 12 deletions.
diff --git a/src/passes/LogExecution.cpp b/src/passes/LogExecution.cpp
@@ -39,9 +39,18 @@ namespace wasm {
 Name LOGGER("log_execution");
 
 struct LogExecution : public WalkerPass<PostWalker<LogExecution>> {
+  // The module name the logger function is imported from.
+  IString loggerModule;
+
   // Adds calls to new imports.
   bool addsEffects() override { return true; }
 
+  void run(Module* module) override {
+    auto& options = getPassOptions();
+    loggerModule = options.getArgumentOrDefault("log-execution", "");
+    super::run(module);
+  }
+
   void visitLoop(Loop* curr) { curr->body = makeLogCall(curr->body); }
 
   void visitReturn(Return* curr) { replaceCurrent(makeLogCall(curr)); }
@@ -63,23 +72,32 @@ struct LogExecution : public WalkerPass<PostWalker<LogExecution>> {
     auto import =
       Builder::makeFunction(LOGGER, Signature(Type::i32, Type::none), {});
 
-    // Import the log function from import "env" if the module
-    // imports other functions from that name.
-    for (auto& func : curr->functions) {
-      if (func->imported() && func->module == ENV) {
-        import->module = func->module;
-        break;
-      }
-    }
-
-    // If not, then pick the import name of the first function we find.
-    if (!import->module) {
+    if (loggerModule != "") {
+      import->module = loggerModule;
+    } else {
+      // Import the log function from import "env" if the module
+      // imports other functions from that name.
       for (auto& func : curr->functions) {
-        if (func->imported()) {
+        if (func->imported() && func->module == ENV) {
           import->module = func->module;
           break;
         }
       }
+
+      // If not, then pick the import name of the first function we find.
+      if (!import->module) {
+        for (auto& func : curr->functions) {
+          if (func->imported()) {
+            import->module = func->module;
+            break;
+          }
+        }
+      }
+
+      // If no function was found, use ENV.
+      if (!import->module) {
+        import->module = ENV;
+      }
     }
 
     import->base = LOGGER;

diff --git a/test/lit/passes/log-execution_arg.wast b/test/lit/passes/log-execution_arg.wast
@@ -0,0 +1,24 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; Test the option to provide the module name as an argument.
+;; RUN: foreach %s %t wasm-opt --log-execution=foo -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (type $1 (func (param i32)))
+
+  ;; CHECK:      (import "env" "func" (func $import))
+  (import "env" "func" (func $import))
+  ;; CHECK:      (import "foo" "log_execution" (func $log_execution (param i32)))
+
+  ;; CHECK:      (func $nopp
+  ;; CHECK-NEXT:  (call $log_execution
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $nopp
+    (nop)
+  )
+)
+