1 files changed, 62 insertions, 8 deletions

diff --git a/src/passes/Asyncify.cpp b/src/passes/Asyncify.cpp
index 2708ae4e0..d4a788adb 100644
--- a/src/passes/Asyncify.cpp
+++ b/src/passes/Asyncify.cpp
@@ -291,6 +291,7 @@
 //      of their original range.
 //
 
+#include "cfg/liveness-traversal.h"
 #include "ir/effects.h"
 #include "ir/find_all.h"
 #include "ir/literal-utils.h"
@@ -1149,8 +1150,9 @@ struct AsyncifyLocals : public WalkerPass<PostWalker<AsyncifyLocals>> {
     if (!analyzer->needsInstrumentation(func)) {
       return;
     }
-    // Note the locals we want to preserve, before we add any more helpers.
-    numPreservableLocals = func->getNumLocals();
+    // Find the locals that we actually need to load and save: any local that is
+    // alive at a relevant call site must be handled, but others can be ignored.
+    findRelevantLiveLocals(func);
     // The new function body has a prelude to load locals if rewinding,
     // then the actual main body, which does all its unwindings by breaking
     // to the unwind block, which then handles pushing the call index, as
@@ -1198,16 +1200,61 @@ private:
   std::unique_ptr<AsyncifyBuilder> builder;
 
   Index rewindIndex;
-  Index numPreservableLocals;
   std::map<Type, Index> fakeCallLocals;
+  std::set<Index> relevantLiveLocals;
+
+  void findRelevantLiveLocals(Function* func) {
+    struct RelevantLiveLocalsWalker
+      : public LivenessWalker<RelevantLiveLocalsWalker,
+                              Visitor<RelevantLiveLocalsWalker>> {
+      // Basic blocks that have a possible unwind/rewind in them.
+      std::set<BasicBlock*> relevantBasicBlocks;
+
+      void visitCall(Call* curr) {
+        if (!currBasicBlock) {
+          return;
+        }
+        // Note blocks where we might unwind/rewind, all of which have a
+        // possible call to ASYNCIFY_CHECK_CALL_INDEX emitted right before the
+        // actual call.
+        // Note that each relevant original call was turned into a sequence of
+        // instructions, one of which is an if and then a call to this special
+        // intrinsic. We rely on the fact that if a local was live at the
+        // original call, it also would be in all that sequence of instructions,
+        // and in particular at the call we look for here (which is right before
+        // the call, and so anything that has its final use at the call is still
+        // live here).
+        if (curr->target == ASYNCIFY_CHECK_CALL_INDEX) {
+          relevantBasicBlocks.insert(currBasicBlock);
+        }
+      }
+    };
+
+    RelevantLiveLocalsWalker walker;
+    walker.walkFunctionInModule(func, getModule());
+    // The relevant live locals are ones that are alive at an unwind/rewind
+    // location. TODO look more precisely inside basic blocks, as one might stop
+    // being live in the middle
+    for (auto* block : walker.liveBlocks) {
+      if (walker.relevantBasicBlocks.count(block)) {
+        for (auto local : block->contents.start) {
+          relevantLiveLocals.insert(local);
+        }
+      }
+    }
+  }
 
   Expression* makeLocalLoading() {
-    if (numPreservableLocals == 0) {
+    if (relevantLiveLocals.empty()) {
       return builder->makeNop();
     }
     auto* func = getFunction();
+    auto numLocals = func->getNumLocals();
     Index total = 0;
-    for (Index i = 0; i < numPreservableLocals; i++) {
+    for (Index i = 0; i < numLocals; i++) {
+      if (!relevantLiveLocals.count(i)) {
+        continue;
+      }
       total += func->getLocalType(i).getByteSize();
     }
     auto* block = builder->makeBlock();
@@ -1216,7 +1263,10 @@ private:
     block->list.push_back(
       builder->makeLocalSet(tempIndex, builder->makeGetStackPos()));
     Index offset = 0;
-    for (Index i = 0; i < numPreservableLocals; i++) {
+    for (Index i = 0; i < numLocals; i++) {
+      if (!relevantLiveLocals.count(i)) {
+        continue;
+      }
       const auto& types = func->getLocalType(i).expand();
       SmallVector<Expression*, 1> loads;
       for (Index j = 0; j < types.size(); j++) {
@@ -1248,16 +1298,20 @@ private:
   }
 
   Expression* makeLocalSaving() {
-    if (numPreservableLocals == 0) {
+    if (relevantLiveLocals.empty()) {
       return builder->makeNop();
     }
     auto* func = getFunction();
+    auto numLocals = func->getNumLocals();
     auto* block = builder->makeBlock();
     auto tempIndex = builder->addVar(func, Type::i32);
     block->list.push_back(
       builder->makeLocalSet(tempIndex, builder->makeGetStackPos()));
     Index offset = 0;
-    for (Index i = 0; i < numPreservableLocals; i++) {
+    for (Index i = 0; i < numLocals; i++) {
+      if (!relevantLiveLocals.count(i)) {
+        continue;
+      }
       auto localType = func->getLocalType(i);
       const auto& types = localType.expand();
       for (Index j = 0; j < types.size(); j++) {