1 files changed, 20 insertions, 1 deletions
diff --git a/src/passes/LogExecution.cpp b/src/passes/LogExecution.cpp
index 5978ecc72..63ada5819 100644
--- a/src/passes/LogExecution.cpp
+++ b/src/passes/LogExecution.cpp
@@ -59,7 +59,26 @@ struct LogExecution : public WalkerPass<PostWalker<LogExecution>> {
     // Add the import
     auto import =
       Builder::makeFunction(LOGGER, Signature(Type::i32, Type::none), {});
-    import->module = ENV;
+
+    // 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) {
+      for (auto& func : curr->functions) {
+        if (func->imported()) {
+          import->module = func->module;
+          break;
+        }
+      }
+    }
+
     import->base = LOGGER;
     curr->addFunction(std::move(import));
   }