[IRBuilder] Tunnel block inputs into nested scopes

src/passes/Outlining.cpp

@@ -34,12 +34,6 @@
 #define ODBG(statement)
 #endif
 
-// Check a Result or MaybeResult for error and call Fatal() if the error exists.
-#define ASSERT_OK(val)                                                         \
-  if (auto _val = (val); auto err = _val.getErr()) {                           \
-    Fatal() << err->msg;                                                       \
-  }
-
 namespace wasm {
 
 // This custom hasher conforms to std::hash<Key>. Its purpose is to provide

src/support/result.h

@@ -17,6 +17,7 @@
 #ifndef wasm_support_result_h
 #define wasm_support_result_h
 
+#include <cassert>
 #include <string>
 #include <variant>
 
@@ -39,6 +40,13 @@ struct Err {
     return (typename decltype(_val)::ErrorType)(*err);                         \
   }
 
+// Assert that a Result or MaybeResult is not an error.
+#define ASSERT_OK(val)                                                         \
+  do {                                                                         \
+    [[maybe_unused]] auto _val = (val);                                        \
+    assert(!_val.getErr());                                                    \
+  } while (false)
+
 // Represent a result of type T or an error message.
 template<typename T = Ok, typename E = Err> struct [[nodiscard]] Result {
   using ErrorType = E;

src/wasm-ir-builder.h

@@ -433,10 +433,10 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     // function's `explicitLabelNames`.
     bool labelExplicit = false;
 
-    // If the control flow scope has an input type, we need to lower it using a
-    // scratch local because we cannot represent control flow input in the IR.
+    // If the control flow scope has an input type, we need to lower it using
+    // scratch locals because we cannot represent control flow input in the IR.
     Type inputType;
-    Index inputLocal = -1;
+    std::vector<Index> inputLocals;
 
     // If there are br_on_*, try_table, or resume branches that target this
     // scope and carry additional values, we need to use a scratch local to
@@ -461,13 +461,6 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     ScopeCtx() : scope(NoScope{}) {}
     ScopeCtx(Scope scope, Type inputType)
       : scope(scope), inputType(inputType) {}
-    ScopeCtx(
-      Scope scope, Name label, bool labelUsed, Type inputType, Index inputLocal)
-      : scope(scope), label(label), labelUsed(labelUsed), inputType(inputType),
-        inputLocal(inputLocal) {}
-    ScopeCtx(Scope scope, Name label, bool labelUsed, Name branchLabel)
-      : scope(scope), label(label), branchLabel(branchLabel),
-        labelUsed(labelUsed) {}
 
     static ScopeCtx makeFunc(Function* func) {
       return ScopeCtx(FuncScope{func}, Type::none);
@@ -514,7 +507,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
       unreachable = false;
       if (!keepInput) {
         inputType = Type::none;
-        inputLocal = -1;
+        inputLocals.clear();
       }
     }
     bool isNone() { return std::get_if<NoScope>(&scope); }
@@ -715,35 +708,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Index blockHint = 0;
   Index labelHint = 0;
 
-  Result<> pushScope(ScopeCtx&& scope) {
-    if (auto label = scope.getOriginalLabel()) {
-      // Assign a fresh label to the scope, if necessary.
-      if (!scope.label) {
-        scope.label = makeFresh(label);
-      }
-      // Record the original label to handle references to it correctly.
-      labelDepths[label].push_back(scopeStack.size() + 1);
-    }
-    if (binaryPos) {
-      scope.startPos = lastBinaryPos;
-      lastBinaryPos = *binaryPos;
-    }
-    bool hasInput = scope.inputType != Type::none;
-    Index inputLocal = scope.inputLocal;
-    if (hasInput && !scope.isDelimiter()) {
-      if (inputLocal == Index(-1)) {
-        auto scratch = addScratchLocal(scope.inputType);
-        CHECK_ERR(scratch);
-        inputLocal = scope.inputLocal = *scratch;
-      }
-      CHECK_ERR(makeLocalSet(inputLocal));
-    }
-    scopeStack.emplace_back(std::move(scope));
-    if (hasInput) {
-      CHECK_ERR(makeLocalGet(inputLocal));
-    }
-    return Ok{};
-  }
+  Result<> pushScope(ScopeCtx&& scope);
 
   ScopeCtx& getScope() {
     if (scopeStack.empty()) {
@@ -806,7 +771,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<std::pair<Index, Name>> getExtraOutputLocalAndLabel(Index label,
                                                              size_t extraArity);
   Expression* fixExtraOutput(ScopeCtx& scope, Name label, Expression* expr);
-  void fixLoopWithInput(Loop* loop, Type inputType, Index scratch);
+  void fixLoopWithInput(Loop* loop,
+                        Type inputType,
+                        const std::vector<Index>& scratch);
 
   void applyAnnotations(Expression* expr, const CodeAnnotation& annotation);
 

src/wasm/wasm-ir-builder.cpp

@@ -778,6 +778,17 @@ struct IRBuilder::ChildPopper
     ConstraintCollector{builder, children}.visitStackSwitch(curr, ct);
     return popConstrainedChildren(children);
   }
+
+  Result<std::vector<Expression*>> popInputs(Type inputType) {
+    std::vector<Expression*> inputs(inputType.size());
+    std::vector<Child> children;
+    children.reserve(inputType.size());
+    for (size_t i = 0; i < inputType.size(); ++i) {
+      children.push_back({&inputs[i], {inputType[i]}});
+    }
+    CHECK_ERR(popConstrainedChildren(children));
+    return inputs;
+  }
 };
 
 Result<> IRBuilder::visit(Expression* curr) {
@@ -877,6 +888,79 @@ IRBuilder::visitTryTableStart(TryTable* trytable, Name label, Type inputType) {
   return pushScope(ScopeCtx::makeTryTable(trytable, label, inputType));
 }
 
+Result<> IRBuilder::pushScope(ScopeCtx&& scope) {
+  if (auto label = scope.getOriginalLabel()) {
+    // Assign a fresh label to the scope, if necessary.
+    if (!scope.label) {
+      scope.label = makeFresh(label);
+    }
+    // Record the original label to handle references to it correctly.
+    labelDepths[label].push_back(scopeStack.size() + 1);
+  }
+  if (binaryPos) {
+    scope.startPos = lastBinaryPos;
+    lastBinaryPos = *binaryPos;
+  }
+  bool hasInput = scope.inputType != Type::none;
+  std::vector<Expression*> inputs;
+  if (!scope.isDelimiter()) {
+    bool isUnnamedBlock = scope.getBlock() && !scope.getBlock()->name;
+    if (hasInput) {
+      auto popped = ChildPopper{*this}.popInputs(scope.inputType);
+      CHECK_ERR(popped);
+      inputs = std::move(*popped);
+    } else if (isUnnamedBlock && getScope().exprStack.size() == 1 &&
+               getScope().exprStack[0]->is<Pop>()) {
+      // Allow pops to move into unnamed blocks even without input types so we
+      // can round-trip WAT without the printer needing to synthesize block
+      // input types. (Note that the block may receive a name later if we
+      // encounter a break to it; in that case, the fixups that run later will
+      // have work to do.)
+      inputs.push_back(getScope().exprStack.back());
+      getScope().exprStack.pop_back();
+    }
+    const auto& parentLocals = getScope().inputLocals;
+    // The new scope may only receive a suffix of the outer scope's parameters,
+    // so match inputs starting from the remaining size of the outer stack
+    // rather than from the beginning.
+    size_t baseIndex = getScope().exprStack.size();
+    for (size_t i = 0; i < inputs.size(); ++i) {
+      // Move an enclosing scope's input expression (a scratch local.get or a
+      // catch pop) directly into this scope rather than allocating a new
+      // scratch local.
+      auto* expr = inputs[i];
+      size_t parentIndex = baseIndex + i;
+      bool isInputGet = expr->is<LocalGet>() &&
+                        expr->type == scope.inputType[i] &&
+                        parentIndex < parentLocals.size();
+      bool isPop = isUnnamedBlock && inputs.size() == 1 && expr->is<Pop>();
+      if (isInputGet) {
+        assert(expr->cast<LocalGet>()->index == parentLocals[parentIndex]);
+        scope.inputLocals.push_back(parentLocals[parentIndex]);
+      } else if (!isPop) {
+        auto scratch = addScratchLocal(scope.inputType[i]);
+        CHECK_ERR(scratch);
+        scope.inputLocals.push_back(*scratch);
+        pushSynthetic(builder.makeLocalSet(*scratch, expr));
+        inputs[i] = builder.makeLocalGet(*scratch, scope.inputType[i]);
+      }
+    }
+  } else if (hasInput) {
+    // Delimiters (such as `else`) reuse the scratch locals that were already
+    // populated when the enclosing scope started, pushing fresh `local.get`s
+    // onto the new scope's stack.
+    for (size_t i = 0; i < scope.inputLocals.size(); ++i) {
+      inputs.push_back(
+        builder.makeLocalGet(scope.inputLocals[i], scope.inputType[i]));
+    }
+  }
+  scopeStack.emplace_back(std::move(scope));
+  for (auto* input : inputs) {
+    pushSynthetic(input);
+  }
+  return Ok{};
+}
+
 Result<Expression*> IRBuilder::finishScope(Block* block) {
 #if IR_BUILDER_DEBUG
   if (auto* loc = std::get_if<Function::DebugLocation>(&debugLoc)) {
@@ -1212,8 +1296,8 @@ Result<> IRBuilder::visitEnd() {
     if (scope.inputType != Type::none && scope.labelUsed) {
       // Branches to this loop carry values, but Binaryen IR does not support
       // that. Fix this by trampolining the branches through new code that sets
-      // the branch value to the appropriate scratch local.
-      fixLoopWithInput(loop, scope.inputType, scope.inputLocal);
+      // the branch value to the appropriate scratch locals.
+      fixLoopWithInput(loop, scope.inputType, scope.inputLocals);
     }
     loop->finalize(loop->type);
     push(loop);
@@ -1222,8 +1306,19 @@ Result<> IRBuilder::visitEnd() {
     if (scope.inputType != Type::none) {
       // Normally an if without an else must have type none, but if there is an
       // input parameter, the empty else arm must propagate its value.
-      // Synthesize an else arm that loads the value from the scratch local.
-      iff->ifFalse = builder.makeLocalGet(scope.inputLocal, scope.inputType);
+      // Synthesize an else arm that loads the value from the scratch locals.
+      if (scope.inputType.isSingle()) {
+        iff->ifFalse =
+          builder.makeLocalGet(scope.inputLocals[0], scope.inputType);
+      } else {
+        std::vector<Expression*> gets;
+        gets.reserve(scope.inputLocals.size());
+        for (size_t i = 0; i < scope.inputLocals.size(); ++i) {
+          gets.push_back(
+            builder.makeLocalGet(scope.inputLocals[i], scope.inputType[i]));
+        }
+        iff->ifFalse = builder.makeTupleMake(gets);
+      }
     } else {
       iff->ifFalse = nullptr;
     }
@@ -1368,17 +1463,18 @@ IRBuilder::fixExtraOutput(ScopeCtx& scope, Name label, Expression* curr) {
   return curr;
 }
 
-// Branches to this loop need to be trampolined through code that sets the value
-// carried by the branch to the appropriate scratch local before branching to
-// the loop. Transform this:
+// Branches to this loop need to be trampolined through code that sets the
+// values carried by the branch to the appropriate scratch locals before
+// branching to the loop. Transform this:
 //
 //   (loop $l (param t1) (result t2) ...)
 //
-// to this:
+// to this (or, when t1 is a tuple, using a tuple scratch local to extract each
+// element into its corresponding input scratch local):
 //
 //  (loop $l0 (result t2)
 //    (block $l1 (result t2)
-//      (local.set $scratch ;; set the branch values to the scratch local
+//      (local.set $scratch ;; set the branch value to the scratch local
 //        (block $l (result t1)
 //          (br $l1 ;; exit the loop with the fallthrough value, if any.
 //            ...   ;; contains branches to $l
@@ -1388,7 +1484,9 @@ IRBuilder::fixExtraOutput(ScopeCtx& scope, Name label, Expression* curr) {
 //      (br $l0) ;; continue the loop
 //    )
 //  )
-void IRBuilder::fixLoopWithInput(Loop* loop, Type inputType, Index scratch) {
+void IRBuilder::fixLoopWithInput(Loop* loop,
+                                 Type inputType,
+                                 const std::vector<Index>& scratch) {
   auto l = loop->name;
   auto l0 = makeFresh(l, 0);
   auto l1 = makeFresh(l, 1);
@@ -1399,10 +1497,22 @@ void IRBuilder::fixLoopWithInput(Loop* loop, Type inputType, Index scratch) {
           loop->body, l, builder.makeBreak(l1), inputType)
       : builder.makeBlock(l, {builder.makeBreak(l1, loop->body)}, inputType);
 
-  Block* outer = builder.makeBlock(
-    l1,
-    {builder.makeLocalSet(scratch, inner), builder.makeBreak(l0)},
-    loop->type);
+  std::vector<Expression*> list;
+  if (inputType.isSingle()) {
+    list.push_back(builder.makeLocalSet(scratch[0], inner));
+  } else {
+    auto tupleScratch = addScratchLocal(inputType);
+    ASSERT_OK(tupleScratch);
+    list.push_back(builder.makeLocalSet(*tupleScratch, inner));
+    for (Index i = 0; i < scratch.size(); ++i) {
+      list.push_back(builder.makeLocalSet(
+        scratch[i],
+        builder.makeTupleExtract(builder.makeLocalGet(*tupleScratch, inputType),
+                                 i)));
+    }
+  }
+  list.push_back(builder.makeBreak(l0));
+  Block* outer = builder.makeBlock(l1, list, loop->type);
 
   loop->body = outer;
   loop->name = l0;
@@ -1934,38 +2044,22 @@ Result<> IRBuilder::makeUnreachable() {
 Result<> IRBuilder::makePop(Type type) {
   // We don't actually want to create a new Pop expression here because we
   // already create them automatically when starting a legacy catch block that
-  // needs one. Just verify that the Pop we are being asked to make is the same
-  // type as the Pop we have already made, and move it into the current scope if
-  // we have entered unnamed blocks since the catch began. (Note that the block
-  // may receive a name later, if we encounter a break to it, if we process
-  // that after the pop. In that case, the fixups that run later will have work
-  // to do.)
-  ScopeCtx* catchScope = nullptr;
-  for (int i = int(scopeStack.size()) - 1; i >= 0; --i) {
-    auto& scope = scopeStack[i];
-    if (scope.getCatch()) {
-      catchScope = &scope;
-      break;
-    }
-    auto* block = scope.getBlock();
-    if (!block || block->name || !scope.exprStack.empty()) {
-      break;
+  // needs one (and move them into nested unnamed blocks in pushScope). Just
+  // verify that the Pop we are being asked to make is the same type as the Pop
+  // we have already made.
+  auto& scope = getScope();
+  if (scope.exprStack.size() != 1 || !scope.exprStack[0]->is<Pop>()) {
+    if (auto* tryy = scope.getCatch();
+        tryy && wasm.getTag(tryy->catchTags.back())->params() == Type::none) {
+      return Err{"cannot pop from empty tag"};
     }
-  }
-  if (!catchScope || catchScope->exprStack.size() != 1 ||
-      !catchScope->exprStack[0]->is<Pop>()) {
     return Err{
       "pop instructions may only appear at the beginning of catch blocks"};
   }
-  auto* pop = catchScope->exprStack[0];
-  auto expectedType = pop->type;
-  if (!Type::isSubType(expectedType, type)) {
+  auto expectedType = scope.exprStack[0]->type;
+  if (type != expectedType) {
     return Err{std::string("Expected pop of type ") + expectedType.toString()};
   }
-  if (catchScope != &getScope()) {
-    catchScope->exprStack.pop_back();
-    pushSynthetic(pop);
-  }
   return Ok{};
 }