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).

2 files changed, 54 insertions, 12 deletions

diff --git a/src/passes/LogExecution.cpp b/src/passes/LogExecution.cpp
index 1b90842bc..f1d48012f 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 000000000..e987655d6
--- /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)
+  )
+)
+