feat: implement tail call optimization

src/passes/CMakeLists.txt

@@ -116,6 +116,7 @@ set(passes_SOURCES
   ReorderGlobals.cpp
   ReorderLocals.cpp
   ReReloop.cpp
+  TailCall.cpp
   TrapMode.cpp
   TypeGeneralizing.cpp
   TypeRefining.cpp

src/passes/TailCall.cpp

@@ -0,0 +1,102 @@
+
+#include "ir/properties.h"
+#include "ir/utils.h"
+#include "pass.h"
+#include "wasm-traversal.h"
+#include "wasm.h"
+#include <stack>
+#include <vector>
+
+namespace wasm {
+
+namespace {
+
+struct Finder : TryDepthWalker<Finder> {
+  explicit Finder(const PassOptions& passOptions)
+    : TryDepthWalker<Finder>(), passOptions(passOptions) {}
+  const PassOptions& passOptions;
+  std::vector<Call*> tailCalls;
+  std::vector<CallIndirect*> tailCallIndirects;
+  void visitFunction(Function* curr) {
+    if (passOptions.shrinkLevel > 0 && passOptions.optimizeLevel == 0) {
+      // When we more force on the binary size, add return_call will increase
+      // the code size.
+      return;
+    }
+    checkTailCall(curr->body);
+  }
+  void visitReturn(Return* curr) {
+    if (tryDepth > 0) {
+      // (return (call ...)) is not equal to (return_call ...) in try block
+      return;
+    }
+    checkTailCall(curr->value);
+  }
+
+private:
+  void checkTailCall(Expression* const curr) {
+    std::stack<Expression*> workList{};
+    workList.push(curr);
+    while (!workList.empty()) {
+      Expression* const target = workList.top();
+      workList.pop();
+      if (auto* call = target->dynCast<Call>()) {
+        if (!call->isReturn && call->type == getFunction()->getResults()) {
+          tailCalls.push_back(call);
+        }
+      } else if (auto* call = target->dynCast<CallIndirect>()) {
+        if (!call->isReturn && call->type == getFunction()->getResults()) {
+          tailCallIndirects.push_back(call);
+        }
+      } else if (auto* ifElse = target->dynCast<If>()) {
+        workList.push(ifElse->ifTrue);
+        workList.push(ifElse->ifFalse);
+      } else if (auto* tryy = target->dynCast<Try>()) {
+        for (Expression* catchBody : tryy->catchBodies) {
+          workList.push(catchBody);
+        }
+      } else if (auto* block = target->dynCast<Block>()) {
+        if (!block->list.empty()) {
+          workList.push(block->list.back());
+        }
+      } else {
+        Expression* const next = Properties::getImmediateFallthrough(
+          target, passOptions, *getModule());
+        if (next != target) {
+          workList.push(next);
+        }
+      }
+    }
+  }
+};
+
+} // namespace
+
+struct TailCallOptimizer : public Pass {
+  bool isFunctionParallel() override { return true; }
+  std::unique_ptr<Pass> create() override {
+    return std::make_unique<TailCallOptimizer>();
+  }
+  void runOnFunction(Module* module, Function* function) override {
+    if (!module->features.hasTailCall()) {
+      return;
+    }
+    Finder finder{getPassOptions()};
+    finder.walkFunctionInModule(function, module);
+    for (Call* call : finder.tailCalls) {
+      if (!call->isReturn) {
+        call->isReturn = true;
+      }
+    }
+    for (CallIndirect* call : finder.tailCallIndirects) {
+      if (!call->isReturn) {
+        call->isReturn = true;
+      }
+    }
+    ReFinalize{}.walkFunctionInModule(function, module);
+  }
+};
+
+Pass* createTailCallPass() { return new TailCallOptimizer(); }
+
+} // namespace wasm

src/passes/pass.cpp

@@ -558,6 +558,9 @@ void PassRegistry::registerPasses() {
   registerPass("strip-target-features",
                "strip the wasm target features section",
                createStripTargetFeaturesPass);
+  registerPass("tail-call-optimization",
+               "transform call to return call",
+               createTailCallPass);
   registerPass("translate-to-new-eh",
                "deprecated; same as translate-to-exnref",
                createTranslateToExnrefPass);

src/passes/passes.h

@@ -178,6 +178,7 @@ Pass* createStripEHPass();
 Pass* createStubUnsupportedJSOpsPass();
 Pass* createSSAifyPass();
 Pass* createSSAifyNoMergePass();
+Pass* createTailCallPass();
 Pass* createTable64LoweringPass();
 Pass* createTranslateToExnrefPass();
 Pass* createTrapModeClamp();

test/lit/help/wasm-metadce.test

@@ -530,6 +530,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --table64-lowering                            alias for memory64-lowering
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --tail-call-optimization                      transform call to return call 
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --trace-calls                                 instrument the build with code
 ;; CHECK-NEXT:                                                 to intercept specific function
 ;; CHECK-NEXT:                                                 calls

test/lit/help/wasm-opt.test

@@ -554,6 +554,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --table64-lowering                            alias for memory64-lowering
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --tail-call-optimization                      transform call to return call 
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --trace-calls                                 instrument the build with code
 ;; CHECK-NEXT:                                                 to intercept specific function
 ;; CHECK-NEXT:                                                 calls

test/lit/help/wasm2js.test

@@ -494,6 +494,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --table64-lowering                            alias for memory64-lowering
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --tail-call-optimization                      transform call to return call 
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --trace-calls                                 instrument the build with code
 ;; CHECK-NEXT:                                                 to intercept specific function
 ;; CHECK-NEXT:                                                 calls

test/lit/tail-call-optimization-eh.wast

@@ -0,0 +1,94 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --tail-call-optimization --enable-tail-call --enable-exception-handling --optimize-level 2 --shrink-level 0 -S -o - | filecheck %s
+
+;; Tests for tail call optimization with exception handling
+
+(module $exception
+  ;; CHECK:      (type $0 (func (result i32)))
+
+  ;; CHECK:      (type $1 (func))
+
+  ;; CHECK:      (tag $empty (type $1))
+  (tag $empty)
+  ;; CHECK:      (func $f (result i32)
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $f (result i32)
+    i32.const 0
+  )
+  ;; CHECK:      (func $in_try (result i32)
+  ;; CHECK-NEXT:  (try (result i32)
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:    (return_call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $in_try (result i32)
+    try (result i32)
+      call $f
+    catch $empty
+      call $f
+    end
+  )
+  ;; CHECK:      (func $out_try (result i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (return_call $f)
+  ;; CHECK-NEXT: )
+  (func $out_try (result i32)
+    try
+    catch $empty
+    end
+    call $f
+  )
+  ;; CHECK:      (func $in_catch (result i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (return_call $f)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $in_catch (result i32)
+    try
+    catch $empty
+      call $f
+      return
+    end
+    i32.const 0
+  )
+  ;; CHECK:      (func $implicit_in_catch (result i32)
+  ;; CHECK-NEXT:  (try (result i32)
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:    (return_call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch_all
+  ;; CHECK-NEXT:    (return_call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $implicit_in_catch (result i32)
+    try (result i32)
+      i32.const 0
+    catch $empty
+      call $f
+    catch_all
+      call $f
+    end
+  )
+)

test/lit/tail-call-optimization-shrink.wast

@@ -0,0 +1,22 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --tail-call-optimization --enable-tail-call --optimize-level 0 --shrink-level 2 -S -o - | filecheck %s
+
+;; Tests for tail call optimization
+
+(module
+  ;; CHECK:      (type $0 (func (result i32)))
+
+  ;; CHECK:      (func $f (result i32)
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $f (result i32)
+    i32.const 0
+  )
+  ;; CHECK:      (func $implicit_return (result i32)
+  ;; CHECK-NEXT:  (call $f)
+  ;; CHECK-NEXT: )
+  (func $implicit_return (result i32)
+    call $f
+  )
+)