Support try-catch blocks in Asyncify.

src/ir/eh-utils.cpp

@@ -32,7 +32,7 @@ namespace EHUtils {
 // no way to get the given expression's address. But that's fine because pop's
 // pointer is only necessary (in handleBlockNestedPops) to fix it up when it is
 // nested, and if 'catchBody' itself is a pop, we don't need to fix it up.
-static Expression*
+Expression*
 getFirstPop(Expression* catchBody, bool& isPopNested, Expression**& popPtr) {
   Expression* firstChild = catchBody;
   isPopNested = false;

src/ir/eh-utils.h

@@ -24,6 +24,18 @@ namespace wasm {
 
 namespace EHUtils {
 
+// This returns three values, some of them as output parameters:
+// - Return value: 'pop' expression (Expression*), when there is one in
+//   first-descendant line. If there's no such pop, it returns null.
+// - isPopNested: Whether the discovered 'pop' is nested within a block
+// - popPtr: 'pop' expression's pointer (Expression**), when there is one found
+//
+// When 'catchBody' itself is a 'pop', 'pop''s pointer is null, because there is
+// no way to get the given expression's address. But that's fine because pop's
+// pointer is only necessary (in handleBlockNestedPops) to fix it up when it is
+// nested, and if 'catchBody' itself is a pop, we don't need to fix it up.
+Expression* getFirstPop(Expression* catchBody, bool& isPopNested, Expression**& popPtr);
+
 // Returns true if a 'pop' instruction exists in a valid location, which means
 // right after a 'catch' instruction in binary writing order.
 // - This assumes there should be at least a single pop. So given a catch body

src/passes/Asyncify.cpp

@@ -312,6 +312,7 @@
 #include "ir/memory-utils.h"
 #include "ir/module-utils.h"
 #include "ir/utils.h"
+#include "ir/eh-utils.h"
 #include "pass.h"
 #include "support/file.h"
 #include "support/string.h"
@@ -1078,6 +1079,71 @@ struct AsyncifyFlow : public Pass {
         results.pop_back();
         results.push_back(loop);
         continue;
+      } else if (auto* tryy = curr->dynCast<Try>()) {
+        if (item.phase == Work::Scan) {
+          work.push_back(Work{curr, Work::Finish});
+          auto& catchBodies = tryy->catchBodies;
+          auto& catchTags = tryy->catchTags;
+          for (size_t i = catchBodies.size(); i > 0; i--) {
+            // TODO: Can optimize if !analyzer->canChangeState(child, func)
+            auto* child = catchBodies[i - 1];
+            Type type = Type::none;
+            if (i - 1 < catchTags.size()) {
+              type = module->getTag(catchTags[i - 1])->sig.params;
+              if (type != Type::none) {
+                bool isPopNested = false;
+                Expression** popPtr = nullptr;
+                auto* pop = EHUtils::getFirstPop(child, isPopNested, popPtr);
+                assert(pop != nullptr && popPtr != nullptr);
+                auto argsTemp = builder->addVar(func, type);
+                *popPtr = builder->makeLocalGet(argsTemp, type);
+                catchBodies[i - 1] = builder->makeLocalSet(argsTemp, pop);
+                replaceRethrow(child, tryy->name, catchTags[i - 1], type, argsTemp);
+              } else {
+                catchBodies[i - 1] = builder->makeNop();
+                replaceRethrowNoArgs(child, tryy->name, catchTags[i - 1]);
+              }
+            } else {
+              catchBodies[i - 1] = builder->makeNop();
+              assertNoRethrow(child);
+            }
+            work.push_back(Work{child, Work::Scan});
+          }
+          work.push_back(Work{tryy->body, Work::Scan});
+          continue;
+        }
+        tryy->body = results.back();
+        results.pop_back();
+        if (tryy->isDelegate()) {
+          results.push_back(tryy);
+          continue;
+        }
+        auto tagTemp = builder->addVar(func, Type::i32);
+        auto* successTag = builder->makeConstantExpression(Literal((int32_t)0));
+        auto* pre = makeMaybeSkip(builder->makeLocalSet(tagTemp, successTag));
+        std::vector<Expression*> parts = {pre, tryy};
+        auto& catchBodies = tryy->catchBodies;
+        for (size_t i = 0; i < catchBodies.size(); i++) {
+          auto* currentTag =
+            builder->makeConstantExpression(Literal((int32_t)(i + 1)));
+          auto* newCatchBody = results.back();
+          results.pop_back();
+          catchBodies[i] = builder->makeBlock({
+            catchBodies[i],
+            builder->makeLocalSet(tagTemp, currentTag),
+          });
+          // Create catch equilvalent as an if
+          auto* newCatch = builder->makeIf(
+            builder->makeBinary(
+              OrInt32,
+              builder->makeBinary(
+                EqInt32, builder->makeLocalGet(tagTemp, Type::i32), currentTag),
+              builder->makeStateCheck(State::Rewinding)),
+            newCatchBody);
+          parts.push_back(newCatch);
+        }
+        results.push_back(builder->makeBlock(parts));
+        continue;
       } else if (doesCall(curr)) {
         results.push_back(makeCallSupport(curr));
         continue;
@@ -1160,6 +1226,85 @@ struct AsyncifyFlow : public Pass {
     return builder->makeCall(ASYNCIFY_GET_CALL_INDEX, {}, Type::none);
   }
 
+  void replaceRethrow(Expression*& expr, Name name, Name tag, Type type, Index args) {
+    struct Walker : ControlFlowWalker<Walker> {
+      void visitRethrow(Rethrow* rethrow) {
+        if (rethrow->target != name) {
+          return;
+        }
+        for (Expression* parent : controlFlowStack) {
+          if (auto* parent_try = parent->dynCast<Try>()) {
+            if (parent_try->name == name) {
+              return;
+            }
+          }
+        }
+        // TODO: This is probably not correct. Need to splay the tuple value gotten from args into a vector of args?
+        // Seems unfortunate, to splay and then repack. Opportunity for a throw instruction taking a tuple in binaryen?
+        replaceCurrent(builder->makeThrow(tag, {builder->makeLocalGet(args, type)}));
+      }
+      AsyncifyBuilder* builder;
+      Name name;
+      Name tag;
+      Type type;
+      Index args;
+    };
+    Walker walker;
+    walker.builder = builder.get();
+    walker.name = name;
+    walker.tag = tag;
+    walker.type = type;
+    walker.args = args;
+    walker.walk(expr);
+  }
+
+  void replaceRethrowNoArgs(Expression*& expr, Name name, Name tag) {
+    struct Walker : ControlFlowWalker<Walker> {
+      void visitRethrow(Rethrow* rethrow) {
+        if (rethrow->target != name) {
+          return;
+        }
+        for (Expression* parent : controlFlowStack) {
+          if (auto* parent_try = parent->dynCast<Try>()) {
+            if (parent_try->name == name) {
+              return;
+            }
+          }
+        }
+        replaceCurrent(builder->makeThrow(tag, {}));
+      }
+      AsyncifyBuilder* builder;
+      Name name;
+      Name tag;
+    };
+    Walker walker;
+    walker.builder = builder.get();
+    walker.name = name;
+    walker.tag = tag;
+    walker.walk(expr);
+  }
+
+  void assertNoRethrow(Expression* expr) {
+    struct Walker : PostWalker<Walker> {
+      void visitRethrow(Rethrow* rethrow) {
+        if (rethrow->target != name) {
+          return;
+        }
+        for (Expression* parent : controlFlowStack) {
+          if (auto* parent_try = parent->dynCast<Try>()) {
+            if (parent_try->name == name) {
+              return;
+            }
+          }
+        }
+        WASM_UNREACHABLE("Cannot Asyncify rethrow in catch_all block");
+      }
+    };
+    Walker walker;
+    walker.name = name;
+    walker.walk(expr);
+  }
+
   // Given a function that is not instrumented - because we proved it doesn't
   // need it, or depending on the only-list / remove-list - add assertions that
   // verify that property at runtime.

src/wasm.h

@@ -1427,7 +1427,7 @@ class Try : public SpecificExpression<Expression::TryId> {
 public:
   Try(MixedArena& allocator) : catchTags(allocator), catchBodies(allocator) {}
 
-  Name name; // label that can only be targeted by 'delegate's
+  Name name; // label that can only be targeted by 'delegate's or 'rethrow's
   Expression* body;
   ArenaVector<Name> catchTags;
   ExpressionList catchBodies;