2 files changed, 26 insertions, 6 deletions

diff --git a/src/ir/linear-execution.h b/src/ir/linear-execution.h
index c88d0e444..cee039b37 100644
--- a/src/ir/linear-execution.h
+++ b/src/ir/linear-execution.h
@@ -17,9 +17,9 @@
 #ifndef wasm_ir_linear_execution_h
 #define wasm_ir_linear_execution_h
 
-#include <ir/properties.h>
-#include <wasm-traversal.h>
-#include <wasm.h>
+#include "ir/properties.h"
+#include "wasm-traversal.h"
+#include "wasm.h"
 
 namespace wasm {
 
@@ -47,6 +47,17 @@ struct LinearExecutionWalker : public PostWalker<SubType, VisitorType> {
   static void scan(SubType* self, Expression** currp) {
     Expression* curr = *currp;
 
+    auto handleCall = [&](bool isReturn) {
+      // Control is nonlinear if we return, or if EH is enabled or may be.
+      if (isReturn || !self->getModule() ||
+          self->getModule()->features.hasExceptionHandling()) {
+        self->pushTask(SubType::doNoteNonLinear, currp);
+      }
+
+      // Scan the children normally.
+      PostWalker<SubType, VisitorType>::scan(self, currp);
+    };
+
     switch (curr->_id) {
       case Expression::Id::InvalidId:
         WASM_UNREACHABLE("bad id");
@@ -97,6 +108,14 @@ struct LinearExecutionWalker : public PostWalker<SubType, VisitorType> {
         self->maybePushTask(SubType::scan, &curr->cast<Return>()->value);
         break;
       }
+      case Expression::Id::CallId: {
+        handleCall(curr->cast<Call>()->isReturn);
+        return;
+      }
+      case Expression::Id::CallRefId: {
+        handleCall(curr->cast<CallRef>()->isReturn);
+        return;
+      }
       case Expression::Id::TryId: {
         self->pushTask(SubType::doVisitTry, currp);
         self->pushTask(SubType::doNoteNonLinear, currp);
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index a042ba573..ff2141fb4 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -997,9 +997,9 @@ struct SimplifyLocals
     // will inhibit us creating an if return value.
     struct EquivalentOptimizer
       : public LinearExecutionWalker<EquivalentOptimizer> {
+
       std::vector<Index>* numLocalGets;
       bool removeEquivalentSets;
-      Module* module;
       PassOptions passOptions;
 
       bool anotherCycle = false;
@@ -1016,6 +1016,8 @@ struct SimplifyLocals
       }
 
       void visitLocalSet(LocalSet* curr) {
+        auto* module = this->getModule();
+
         // Remove trivial copies, even through a tee
         auto* value =
           Properties::getFallthrough(curr->value, passOptions, *module);
@@ -1123,11 +1125,10 @@ struct SimplifyLocals
     };
 
     EquivalentOptimizer eqOpter;
-    eqOpter.module = this->getModule();
     eqOpter.passOptions = this->getPassOptions();
     eqOpter.numLocalGets = &getCounter.num;
     eqOpter.removeEquivalentSets = allowStructure;
-    eqOpter.walkFunction(func);
+    eqOpter.walkFunctionInModule(func, this->getModule());
     if (eqOpter.refinalize) {
       ReFinalize().walkFunctionInModule(func, this->getModule());
     }