Finalize tail call support (#2246)

Adds tail call support to fuzzer and makes small changes to handle return calls in multiple utilities and passes. Makes larger changes to DAE and inlining passes to properly handle tail calls.

Author: tlively

scripts/fuzz_opt.py

@@ -34,7 +34,6 @@
 # simd: known issues with d8
 # atomics, bulk memory: doesn't work in wasm2js
 # truncsat: https://github.com/WebAssembly/binaryen/issues/2198
-# tail-call: WIP
 CONSTANT_FEATURE_OPTS = ['--all-features']
 
 # possible feature options that are sometimes passed to the tools.
@@ -298,7 +297,7 @@ def run(self, wasm):
     return out
 
   def can_run_on_feature_opts(self, feature_opts):
-    return all([x in feature_opts for x in ['--disable-exception-handling', '--disable-simd', '--disable-threads', '--disable-bulk-memory', '--disable-nontrapping-float-to-int']])
+    return all([x in feature_opts for x in ['--disable-exception-handling', '--disable-simd', '--disable-threads', '--disable-bulk-memory', '--disable-nontrapping-float-to-int', '--disable-tail-call']])
 
 
 class Asyncify(TestCaseHandler):
@@ -343,7 +342,7 @@ def do_asyncify(wasm):
     compare(before, after_asyncify, 'Asyncify (before/after_asyncify)')
 
   def can_run_on_feature_opts(self, feature_opts):
-    return all([x in feature_opts for x in ['--disable-exception-handling', '--disable-simd']])
+    return all([x in feature_opts for x in ['--disable-exception-handling', '--disable-simd', '--disable-tail-call']])
 
 
 # The global list of all test case handlers

src/ir/ExpressionAnalyzer.cpp

@@ -132,9 +132,13 @@ template<typename T> void visitImmediates(Expression* curr, T& visitor) {
       }
       visitor.visitScopeName(curr->default_);
     }
-    void visitCall(Call* curr) { visitor.visitNonScopeName(curr->target); }
+    void visitCall(Call* curr) {
+      visitor.visitNonScopeName(curr->target);
+      visitor.visitInt(curr->isReturn);
+    }
     void visitCallIndirect(CallIndirect* curr) {
       visitor.visitNonScopeName(curr->fullType);
+      visitor.visitInt(curr->isReturn);
     }
     void visitLocalGet(LocalGet* curr) { visitor.visitIndex(curr->index); }
     void visitLocalSet(LocalSet* curr) { visitor.visitIndex(curr->index); }

src/ir/ExpressionManipulator.cpp

@@ -71,15 +71,16 @@ flexibleCopy(Expression* original, Module& wasm, CustomCopier custom) {
                                 copy(curr->value));
     }
     Expression* visitCall(Call* curr) {
-      auto* ret = builder.makeCall(curr->target, {}, curr->type);
+      auto* ret =
+        builder.makeCall(curr->target, {}, curr->type, curr->isReturn);
       for (Index i = 0; i < curr->operands.size(); i++) {
         ret->operands.push_back(copy(curr->operands[i]));
       }
       return ret;
     }
     Expression* visitCallIndirect(CallIndirect* curr) {
       auto* ret = builder.makeCallIndirect(
-        curr->fullType, copy(curr->target), {}, curr->type);
+        curr->fullType, copy(curr->target), {}, curr->type, curr->isReturn);
       for (Index i = 0; i < curr->operands.size(); i++) {
         ret->operands.push_back(copy(curr->operands[i]));
       }

src/ir/effects.h

@@ -223,14 +223,22 @@ struct EffectAnalyzer
 
   void visitCall(Call* curr) {
     calls = true;
+    if (curr->isReturn) {
+      branches = true;
+    }
     if (debugInfo) {
       // debugInfo call imports must be preserved very strongly, do not
       // move code around them
       // FIXME: we could check if the call is to an import
       branches = true;
     }
   }
-  void visitCallIndirect(CallIndirect* curr) { calls = true; }
+  void visitCallIndirect(CallIndirect* curr) {
+    calls = true;
+    if (curr->isReturn) {
+      branches = true;
+    }
+  }
   void visitLocalGet(LocalGet* curr) { localsRead.insert(curr->index); }
   void visitLocalSet(LocalSet* curr) { localsWritten.insert(curr->index); }
   void visitGlobalGet(GlobalGet* curr) { globalsRead.insert(curr->name); }

src/passes/Asyncify.cpp

@@ -349,6 +349,9 @@ class ModuleAnalyzer {
         }
         struct Walker : PostWalker<Walker> {
           void visitCall(Call* curr) {
+            if (curr->isReturn) {
+              Fatal() << "tail calls not yet supported in aysncify";
+            }
             auto* target = module->getFunction(curr->target);
             if (target->imported() && target->module == ASYNCIFY) {
               // Redirect the imports to the functions we'll add later.
@@ -375,6 +378,9 @@ class ModuleAnalyzer {
             info->callsTo.insert(target);
           }
           void visitCallIndirect(CallIndirect* curr) {
+            if (curr->isReturn) {
+              Fatal() << "tail calls not yet supported in aysncify";
+            }
             if (canIndirectChangeState) {
               info->canChangeState = true;
             }

src/passes/DeadArgumentElimination.cpp

@@ -57,6 +57,16 @@ struct DAEFunctionInfo {
   // Map of all calls that are dropped, to their drops' locations (so that
   // if we can optimize out the drop, we can replace the drop there).
   std::unordered_map<Call*, Expression**> droppedCalls;
+  // Whether this function contains any tail calls (including indirect tail
+  // calls) and the set of functions this function tail calls. Tail-callers and
+  // tail-callees cannot have their dropped returns removed because of the
+  // constraint that tail-callees must have the same return type as
+  // tail-callers. Indirectly tail called functions are already not optimized
+  // because being in a table inhibits DAE. TODO: Allow the removal of dropped
+  // returns from tail-callers if their tail-callees can have their returns
+  // removed as well.
+  bool hasTailCalls = false;
+  std::unordered_set<Name> tailCallees;
   // Whether the function can be called from places that
   // affect what we can do. For now, any call we don't
   // see inhibits our optimizations, but TODO: an export
@@ -117,6 +127,16 @@ struct DAEScanner
     if (!getModule()->getFunction(curr->target)->imported()) {
       info->calls[curr->target].push_back(curr);
     }
+    if (curr->isReturn) {
+      info->hasTailCalls = true;
+      info->tailCallees.insert(curr->target);
+    }
+  }
+
+  void visitCallIndirect(CallIndirect* curr) {
+    if (curr->isReturn) {
+      info->hasTailCalls = true;
+    }
   }
 
   void visitDrop(Drop* curr) {
@@ -239,6 +259,7 @@ struct DAE : public Pass {
     DAEScanner(&infoMap).run(runner, module);
     // Combine all the info.
     std::unordered_map<Name, std::vector<Call*>> allCalls;
+    std::unordered_set<Name> tailCallees;
     for (auto& pair : infoMap) {
       auto& info = pair.second;
       for (auto& pair : info.calls) {
@@ -247,6 +268,9 @@ struct DAE : public Pass {
         auto& allCallsToName = allCalls[name];
         allCallsToName.insert(allCallsToName.end(), calls.begin(), calls.end());
       }
+      for (auto& callee : info.tailCallees) {
+        tailCallees.insert(callee);
+      }
       for (auto& pair : info.droppedCalls) {
         allDroppedCalls[pair.first] = pair.second;
       }
@@ -314,14 +338,11 @@ struct DAE : public Pass {
           // Great, it's not used. Check if none of the calls has a param with
           // side effects, as that would prevent us removing them (flattening
           // should have been done earlier).
-          bool canRemove = true;
-          for (auto* call : calls) {
-            auto* operand = call->operands[i];
-            if (EffectAnalyzer(runner->options, operand).hasSideEffects()) {
-              canRemove = false;
-              break;
-            }
-          }
+          bool canRemove =
+            std::none_of(calls.begin(), calls.end(), [&](Call* call) {
+              auto* operand = call->operands[i];
+              return EffectAnalyzer(runner->options, operand).hasSideEffects();
+            });
           if (canRemove) {
             // Wonderful, nothing stands in our way! Do it.
             // TODO: parallelize this?
@@ -348,18 +369,21 @@ struct DAE : public Pass {
         if (infoMap[name].hasUnseenCalls) {
           continue;
         }
+        if (infoMap[name].hasTailCalls) {
+          continue;
+        }
+        if (tailCallees.find(name) != tailCallees.end()) {
+          continue;
+        }
         auto iter = allCalls.find(name);
         if (iter == allCalls.end()) {
           continue;
         }
         auto& calls = iter->second;
-        bool allDropped = true;
-        for (auto* call : calls) {
-          if (!allDroppedCalls.count(call)) {
-            allDropped = false;
-            break;
-          }
-        }
+        bool allDropped =
+          std::all_of(calls.begin(), calls.end(), [&](Call* call) {
+            return allDroppedCalls.count(call);
+          });
         if (!allDropped) {
           continue;
         }

src/passes/DeadCodeElimination.cpp

@@ -365,7 +365,12 @@ struct DeadCodeElimination
     return curr;
   }
 
-  void visitCall(Call* curr) { handleCall(curr); }
+  void visitCall(Call* curr) {
+    handleCall(curr);
+    if (curr->isReturn) {
+      reachable = false;
+    }
+  }
 
   void visitCallIndirect(CallIndirect* curr) {
     if (handleCall(curr) != curr) {
@@ -380,6 +385,9 @@ struct DeadCodeElimination
       block->finalize(curr->type);
       replaceCurrent(block);
     }
+    if (curr->isReturn) {
+      reachable = false;
+    }
   }
 
   // Append the reachable operands of the current node to a block, and replace

src/passes/Directize.cpp

@@ -65,7 +65,8 @@ struct FunctionDirectizer : public WalkerPass<PostWalker<FunctionDirectizer>> {
       }
       // Everything looks good!
       replaceCurrent(
-        Builder(*getModule()).makeCall(name, curr->operands, curr->type));
+        Builder(*getModule())
+          .makeCall(name, curr->operands, curr->type, curr->isReturn));
     }
   }
 

src/passes/I64ToI32Lowering.cpp

@@ -262,9 +262,14 @@ struct I64ToI32Lowering : public WalkerPass<PostWalker<I64ToI32Lowering>> {
     return call;
   }
   void visitCall(Call* curr) {
+    if (curr->isReturn &&
+        getModule()->getFunction(curr->target)->result == i64) {
+      Fatal()
+        << "i64 to i32 lowering of return_call values not yet implemented";
+    }
     auto* fixedCall = visitGenericCall<Call>(
       curr, [&](std::vector<Expression*>& args, Type ty) {
-        return builder->makeCall(curr->target, args, ty);
+        return builder->makeCall(curr->target, args, ty, curr->isReturn);
       });
     // If this was to an import, we need to call the legal version. This assumes
     // that legalize-js-interface has been run before.
@@ -275,10 +280,15 @@ struct I64ToI32Lowering : public WalkerPass<PostWalker<I64ToI32Lowering>> {
   }
 
   void visitCallIndirect(CallIndirect* curr) {
+    if (curr->isReturn &&
+        getModule()->getFunctionType(curr->fullType)->result == i64) {
+      Fatal()
+        << "i64 to i32 lowering of return_call values not yet implemented";
+    }
     visitGenericCall<CallIndirect>(
       curr, [&](std::vector<Expression*>& args, Type ty) {
         return builder->makeCallIndirect(
-          curr->fullType, curr->target, args, ty);
+          curr->fullType, curr->target, args, ty, curr->isReturn);
       });
   }
 

src/passes/Inlining.cpp

@@ -146,7 +146,17 @@ struct Planner : public WalkerPass<PostWalker<Planner>> {
     // plan to inline if we know this is valid to inline, and if the call is
     // actually performed - if it is dead code, it's pointless to inline.
     // we also cannot inline ourselves.
-    if (state->worthInlining.count(curr->target) && curr->type != unreachable &&
+    bool isUnreachable;
+    if (curr->isReturn) {
+      // Tail calls are only actually unreachable if an argument is
+      isUnreachable =
+        std::any_of(curr->operands.begin(),
+                    curr->operands.end(),
+                    [](Expression* op) { return op->type == unreachable; });
+    } else {
+      isUnreachable = curr->type == unreachable;
+    }
+    if (state->worthInlining.count(curr->target) && !isUnreachable &&
         curr->target != getFunction()->name) {
       // nest the call in a block. that way the location of the pointer to the
       // call will not change even if we inline multiple times into the same
@@ -164,32 +174,69 @@ struct Planner : public WalkerPass<PostWalker<Planner>> {
   InliningState* state;
 };
 
+struct Updater : public PostWalker<Updater> {
+  Module* module;
+  std::map<Index, Index> localMapping;
+  Name returnName;
+  Builder* builder;
+  void visitReturn(Return* curr) {
+    replaceCurrent(builder->makeBreak(returnName, curr->value));
+  }
+  // Return calls in inlined functions should only break out of the scope of
+  // the inlined code, not the entire function they are being inlined into. To
+  // achieve this, make the call a non-return call and add a break. This does
+  // not cause unbounded stack growth because inlining and return calling both
+  // avoid creating a new stack frame.
+  template<typename T> void handleReturnCall(T* curr, Type targetType) {
+    curr->isReturn = false;
+    curr->type = targetType;
+    if (isConcreteType(targetType)) {
+      replaceCurrent(builder->makeBreak(returnName, curr));
+    } else {
+      replaceCurrent(builder->blockify(curr, builder->makeBreak(returnName)));
+    }
+  }
+  void visitCall(Call* curr) {
+    if (curr->isReturn) {
+      handleReturnCall(curr, module->getFunction(curr->target)->result);
+    }
+  }
+  void visitCallIndirect(CallIndirect* curr) {
+    if (curr->isReturn) {
+      handleReturnCall(curr, module->getFunctionType(curr->fullType)->result);
+    }
+  }
+  void visitLocalGet(LocalGet* curr) {
+    curr->index = localMapping[curr->index];
+  }
+  void visitLocalSet(LocalSet* curr) {
+    curr->index = localMapping[curr->index];
+  }
+};
+
 // Core inlining logic. Modifies the outside function (adding locals as
 // needed), and returns the inlined code.
 static Expression*
 doInlining(Module* module, Function* into, InliningAction& action) {
   Function* from = action.contents;
   auto* call = (*action.callSite)->cast<Call>();
+  // Works for return_call, too
+  Type retType = module->getFunction(call->target)->result;
   Builder builder(*module);
-  auto* block = Builder(*module).makeBlock();
+  auto* block = builder.makeBlock();
   block->name = Name(std::string("__inlined_func$") + from->name.str);
-  *action.callSite = block;
-  // Prepare to update the inlined code's locals and other things.
-  struct Updater : public PostWalker<Updater> {
-    std::map<Index, Index> localMapping;
-    Name returnName;
-    Builder* builder;
-
-    void visitReturn(Return* curr) {
-      replaceCurrent(builder->makeBreak(returnName, curr->value));
+  if (call->isReturn) {
+    if (isConcreteType(retType)) {
+      *action.callSite = builder.makeReturn(block);
+    } else {
+      *action.callSite = builder.makeSequence(block, builder.makeReturn());
     }
-    void visitLocalGet(LocalGet* curr) {
-      curr->index = localMapping[curr->index];
-    }
-    void visitLocalSet(LocalSet* curr) {
-      curr->index = localMapping[curr->index];
-    }
-  } updater;
+  } else {
+    *action.callSite = block;
+  }
+  // Prepare to update the inlined code's locals and other things.
+  Updater updater;
+  updater.module = module;
   updater.returnName = block->name;
   updater.builder = &builder;
   // Set up a locals mapping
@@ -215,12 +262,12 @@ doInlining(Module* module, Function* into, InliningAction& action) {
   }
   updater.walk(contents);
   block->list.push_back(contents);
-  block->type = call->type;
+  block->type = retType;
   // If the function returned a value, we just set the block containing the
   // inlined code to have that type. or, if the function was void and
   // contained void, that is fine too. a bad case is a void function in which
   // we have unreachable code, so we would be replacing a void call with an
-  // unreachable; we need to handle
+  // unreachable.
   if (contents->type == unreachable && block->type == none) {
     // Make the block reachable by adding a break to it
     block->list.push_back(builder.makeBreak(block->name));