introduce visitCallLike to handle callRef, callIndirect

Author: caiiiycuk

src/passes/Asyncify.cpp

@@ -1341,6 +1341,70 @@ struct AsyncifyAssertInNonInstrumented : public Pass {
   Module* module;
 };
 
+struct AsyncifyUnwindWalker
+  : WalkerPass<ExpressionStackWalker<AsyncifyUnwindWalker>> {
+  Function* function;
+  Module* module;
+
+  // Adds a check for Call that is inside a Catch block (we do not handle
+  // unwinding there).
+  template<typename T> void replaceCallWithCheck(T* call) {
+    auto builder = std::make_unique<Builder>(*module);
+    auto check = builder->makeIf(
+      builder->makeBinary(NeInt32,
+                          builder->makeGlobalGet(ASYNCIFY_STATE, Type::i32),
+                          builder->makeConst(int32_t(State::Normal))),
+      builder->makeUnreachable());
+    if (call->type.isConcrete()) {
+      auto temp = builder->addVar(function, call->type);
+      replaceCurrent(builder->makeBlock(
+        {
+          builder->makeLocalSet(temp, call),
+          check,
+          builder->makeLocalGet(temp, call->type),
+        },
+        call->type));
+    } else {
+      replaceCurrent(builder->makeBlock(
+        {
+          call,
+          check,
+        },
+        call->type));
+    }
+  }
+
+  template<typename T> void visitCallLike(T* curr) {
+    assert(!expressionStack.empty());
+    // A return_call (curr->isReturn) can be ignored here: It returns first,
+    // leaving the Catch, before calling.
+    if (curr->isReturn) {
+      return;
+    }
+    // Go up the stack and see if we are in a Catch.
+    Index i = expressionStack.size() - 1;
+    while (i > 0) {
+      auto* expr = expressionStack[i];
+      if (Try* aTry = expr->template dynCast<Try>()) {
+        // check if curr is inside body of aTry (which is safe),
+        // otherwise do replace a call
+        assert(i + 1 < expressionStack.size());
+        if (expressionStack[i + 1] != aTry->body) {
+          replaceCallWithCheck(curr);
+        }
+        break;
+      }
+      i--;
+    }
+  }
+
+  void visitCall(Call* curr) { visitCallLike(curr); }
+
+  void visitCallRef(CallRef* curr) { visitCallLike(curr); }
+
+  void visitCallIndirect(CallIndirect* curr) { visitCallLike(curr); }
+};
+
 struct AsyncifyAssertUnwindCorrectness : Pass {
   bool isFunctionParallel() override { return true; }
 
@@ -1357,63 +1421,7 @@ struct AsyncifyAssertUnwindCorrectness : Pass {
   }
 
   void runOnFunction(Module* module_, Function* function) override {
-    struct UnwindWalker : WalkerPass<ExpressionStackWalker<UnwindWalker>> {
-      Function* function;
-      Module* module;
-
-      // Adds a check for Call that is inside a Catch block (we do not handle unwinding there).
-      void replaceCallWithCheck(Call* call) {
-        auto builder = std::make_unique<Builder>(*module);
-        auto check = builder->makeIf(
-          builder->makeBinary(NeInt32,
-                              builder->makeGlobalGet(ASYNCIFY_STATE, Type::i32),
-                              builder->makeConst(int32_t(State::Normal))),
-          builder->makeUnreachable());
-        if (call->type.isConcrete()) {
-          auto temp = builder->addVar(function, call->type);
-          replaceCurrent(builder->makeBlock(
-            {
-              builder->makeLocalSet(temp, call),
-              check,
-              builder->makeLocalGet(temp, call->type),
-            },
-            call->type));
-        } else {
-          replaceCurrent(builder->makeBlock(
-            {
-              call,
-              check,
-            },
-            call->type));
-        }
-      }
-
-      void visitCall(Call* curr) {
-        assert(!expressionStack.empty());
-        // A return_call (curr->isReturn) can be ignored here: It returns first,
-        // leaving the Catch, before calling.
-        if (curr->isReturn) {
-          return;
-        }
-        // Go up the stack and see if we are in a Catch.
-        Index i = expressionStack.size() - 1;
-        while (i > 0) {
-          auto* expr = expressionStack[i];
-          if (Try* aTry = expr->template dynCast<Try>()) {
-            // check if curr is inside body of aTry (which is safe),
-            // otherwise do replace a call
-            assert(i + 1 < expressionStack.size());
-            if (expressionStack[i + 1] != aTry->body) {
-              replaceCallWithCheck(curr);
-            }
-            break;
-          }
-          i--;
-        }
-      };
-    };
-
-    UnwindWalker walker;
+    AsyncifyUnwindWalker walker;
     walker.function = function;
     walker.module = module_;
     walker.walk(function->body);