From 0c23394e9d73252000c3161fb33344ca7bbf247c Mon Sep 17 00:00:00 2001
From: YAMAMOTO Takashi <yamamoto@midokura.com>
Date: Fri, 31 May 2024 13:04:11 +0900
Subject: [PATCH] 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).
---
 src/passes/LogExecution.cpp            | 42 ++++++++++++++++++--------
 test/lit/passes/log-execution_arg.wast | 24 +++++++++++++++
 2 files changed, 54 insertions(+), 12 deletions(-)
 create mode 100644 test/lit/passes/log-execution_arg.wast

diff --git a/src/passes/LogExecution.cpp b/src/passes/LogExecution.cpp
index 1b90842bc3d..f1d48012f85 100644
--- 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
new file mode 100644
index 00000000000..e987655d6c5
--- /dev/null
+++ 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)
+  )
+)
+