Support control flow inputs in IRBuilder

Since multivalue was standardized, WebAssembly has supported not only
multiple results but also an arbitrary number of inputs on control flow
structures, but until now Binaryen did not support control flow input.
Binaryen IR still has no way to represent control flow input, so lower
it away using scratch locals in IRBuilder. Since both the text and
binary parsers use IRBuilder, this gives us full support for parsing
control flow inputs.

The lowering scheme is mostly simple. A local.set writing the control
flow inputs to a scratch local is inserted immediately before the
control flow structure begins and a local.get retrieving those inputs is
inserted inside the control flow structure before the rest of its body.
The only complications come from ifs, in which the inputs must be
retrieved at the beginning of both arms, and from loops, where branches
to the beginning of the loop must be transformed so their values are
written to the scratch local along the way.

Resolves #6407.

Author: tlively

src/parser/contexts.h

@@ -1447,11 +1447,7 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
 
   Result<HeapType> getBlockTypeFromTypeUse(Index pos, HeapType type) {
     assert(type.isSignature());
-    if (type.getSignature().params != Type::none) {
-      return in.err(pos, "block parameters not yet supported");
-    }
-    // TODO: Once we support block parameters, return an error here if any of
-    // them are named.
+    // TODO: Error if block parameters are named
     return type;
   }
 
@@ -1822,20 +1818,23 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
                      HeapType type) {
     // TODO: validate labels?
     // TODO: Move error on input types to here?
-    return withLoc(pos,
-                   irBuilder.makeBlock(label ? *label : Name{},
-                                       type.getSignature().results));
+    if (!type.isSignature()) {
+      return in.err(pos, "expected function type");
+    }
+    return withLoc(
+      pos, irBuilder.makeBlock(label ? *label : Name{}, type.getSignature()));
   }
 
   Result<> makeIf(Index pos,
                   const std::vector<Annotation>& annotations,
                   std::optional<Name> label,
                   HeapType type) {
     // TODO: validate labels?
-    // TODO: Move error on input types to here?
+    if (!type.isSignature()) {
+      return in.err(pos, "expected function type");
+    }
     return withLoc(
-      pos,
-      irBuilder.makeIf(label ? *label : Name{}, type.getSignature().results));
+      pos, irBuilder.makeIf(label ? *label : Name{}, type.getSignature()));
   }
 
   Result<> visitElse() { return withLoc(irBuilder.visitElse()); }
@@ -1845,21 +1844,23 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
                     std::optional<Name> label,
                     HeapType type) {
     // TODO: validate labels?
-    // TODO: Move error on input types to here?
+    if (!type.isSignature()) {
+      return in.err(pos, "expected function type");
+    }
     return withLoc(
-      pos,
-      irBuilder.makeLoop(label ? *label : Name{}, type.getSignature().results));
+      pos, irBuilder.makeLoop(label ? *label : Name{}, type.getSignature()));
   }
 
   Result<> makeTry(Index pos,
                    const std::vector<Annotation>& annotations,
                    std::optional<Name> label,
                    HeapType type) {
     // TODO: validate labels?
-    // TODO: Move error on input types to here?
+    if (!type.isSignature()) {
+      return in.err(pos, "expected function type");
+    }
     return withLoc(
-      pos,
-      irBuilder.makeTry(label ? *label : Name{}, type.getSignature().results));
+      pos, irBuilder.makeTry(label ? *label : Name{}, type.getSignature()));
   }
 
   Result<> makeTryTable(Index pos,
@@ -1875,12 +1876,10 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
       labels.push_back(info.label);
       isRefs.push_back(info.isRef);
     }
-    return withLoc(pos,
-                   irBuilder.makeTryTable(label ? *label : Name{},
-                                          type.getSignature().results,
-                                          tags,
-                                          labels,
-                                          isRefs));
+    return withLoc(
+      pos,
+      irBuilder.makeTryTable(
+        label ? *label : Name{}, type.getSignature(), tags, labels, isRefs));
   }
 
   Result<> visitCatch(Index pos, Name tag) {

src/wasm-binary.h

@@ -1467,10 +1467,12 @@ class WasmBinaryReader {
 
   bool getBasicType(int32_t code, Type& out);
   bool getBasicHeapType(int64_t code, HeapType& out);
+  // Get the signature of control flow structure.
+  Signature getBlockType();
   // Read a value and get a type for it.
   Type getType();
   // Get a type given the initial S32LEB has already been read, and is provided.
-  Type getType(int initial);
+  Type getType(int code);
   HeapType getHeapType();
   HeapType getIndexedHeapType();
 

src/wasm-ir-builder.h

@@ -80,15 +80,18 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // the corresponding `makeXYZ` function below instead of `visitXYZStart`, but
   // either way must call `visitEnd` and friends at the appropriate times.
   Result<> visitFunctionStart(Function* func);
-  Result<> visitBlockStart(Block* block);
-  Result<> visitIfStart(If* iff, Name label = {});
+  Result<> visitBlockStart(Block* block, Type inputType = Type::none);
+  Result<> visitIfStart(If* iff, Name label = {}, Type inputType = Type::none);
   Result<> visitElse();
-  Result<> visitLoopStart(Loop* iff);
-  Result<> visitTryStart(Try* tryy, Name label = {});
+  Result<> visitLoopStart(Loop* iff, Type inputType = Type::none);
+  Result<>
+  visitTryStart(Try* tryy, Name label = {}, Type inputType = Type::none);
   Result<> visitCatch(Name tag);
   Result<> visitCatchAll();
   Result<> visitDelegate(Index label);
-  Result<> visitTryTableStart(TryTable* trytable, Name label = {});
+  Result<> visitTryTableStart(TryTable* trytable,
+                              Name label = {},
+                              Type inputType = Type::none);
   Result<> visitEnd();
 
   // Used to visit break nodes when traversing a single block without its
@@ -113,9 +116,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // nodes. This is generally safer than calling `visit` because the function
   // signatures ensure that there are no missing fields.
   Result<> makeNop();
-  Result<> makeBlock(Name label, Type type);
-  Result<> makeIf(Name label, Type type);
-  Result<> makeLoop(Name label, Type type);
+  Result<> makeBlock(Name label, Signature sig);
+  Result<> makeIf(Name label, Signature sig);
+  Result<> makeLoop(Name label, Signature sig);
   Result<> makeBreak(Index label, bool isConditional);
   Result<> makeSwitch(const std::vector<Index>& labels, Index defaultLabel);
   // Unlike Builder::makeCall, this assumes the function already exists.
@@ -180,9 +183,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<> makeTableFill(Name table);
   Result<> makeTableCopy(Name destTable, Name srcTable);
   Result<> makeTableInit(Name elem, Name table);
-  Result<> makeTry(Name label, Type type);
+  Result<> makeTry(Name label, Signature sig);
   Result<> makeTryTable(Name label,
-                        Type type,
+                        Signature sig,
                         const std::vector<Name>& tags,
                         const std::vector<Index>& labels,
                         const std::vector<bool>& isRefs);
@@ -323,13 +326,21 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
 
     // The branch label name for this scope. Always fresh, never shadowed.
     Name label;
+
     // For Try/Catch/CatchAll scopes, we need to separately track a label used
     // for branches, since the normal label is only used for delegates.
     Name branchLabel;
 
     bool labelUsed = 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.
+    Type inputType;
+    Index inputLocal = -1;
+
+    // The stack of instructions being built in this scope.
     std::vector<Expression*> exprStack;
+
     // Whether we have seen an unreachable instruction and are in
     // stack-polymorphic unreachable mode.
     bool unreachable = false;
@@ -338,29 +349,39 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     size_t startPos = 0;
 
     ScopeCtx() : scope(NoScope{}) {}
-    ScopeCtx(Scope scope) : scope(scope) {}
-    ScopeCtx(Scope scope, Name label, bool labelUsed)
-      : scope(scope), label(label), labelUsed(labelUsed) {}
+    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});
+      return ScopeCtx(FuncScope{func}, Type::none);
     }
-    static ScopeCtx makeBlock(Block* block) {
-      return ScopeCtx(BlockScope{block});
+    static ScopeCtx makeBlock(Block* block, Type inputType) {
+      return ScopeCtx(BlockScope{block}, inputType);
     }
-    static ScopeCtx makeIf(If* iff, Name originalLabel = {}) {
-      return ScopeCtx(IfScope{iff, originalLabel});
+    static ScopeCtx makeIf(If* iff, Name originalLabel, Type inputType) {
+      return ScopeCtx(IfScope{iff, originalLabel}, inputType);
     }
-    static ScopeCtx
-    makeElse(If* iff, Name originalLabel, Name label, bool labelUsed) {
-      return ScopeCtx(ElseScope{iff, originalLabel}, label, labelUsed);
+    static ScopeCtx makeElse(If* iff,
+                             Name originalLabel,
+                             Name label,
+                             bool labelUsed,
+                             Type inputType,
+                             Index inputLocal) {
+      return ScopeCtx(
+        ElseScope{iff, originalLabel}, label, labelUsed, inputType, inputLocal);
     }
-    static ScopeCtx makeLoop(Loop* loop) { return ScopeCtx(LoopScope{loop}); }
-    static ScopeCtx makeTry(Try* tryy, Name originalLabel = {}) {
-      return ScopeCtx(TryScope{tryy, originalLabel});
+    static ScopeCtx makeLoop(Loop* loop, Type inputType) {
+      return ScopeCtx(LoopScope{loop}, inputType);
+    }
+    static ScopeCtx makeTry(Try* tryy, Name originalLabel, Type inputType) {
+      return ScopeCtx(TryScope{tryy, originalLabel}, inputType);
     }
     static ScopeCtx makeCatch(Try* tryy,
                               Name originalLabel,
@@ -378,8 +399,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
       return ScopeCtx(
         CatchAllScope{tryy, originalLabel}, label, labelUsed, branchLabel);
     }
-    static ScopeCtx makeTryTable(TryTable* trytable, Name originalLabel = {}) {
-      return ScopeCtx(TryTableScope{trytable, originalLabel});
+    static ScopeCtx
+    makeTryTable(TryTable* trytable, Name originalLabel, Type inputType) {
+      return ScopeCtx(TryTableScope{trytable, originalLabel}, inputType);
     }
 
     bool isNone() { return std::get_if<NoScope>(&scope); }
@@ -518,6 +540,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
       }
       WASM_UNREACHABLE("unexpected scope kind");
     }
+    bool isDelimiter() { return getElse() || getCatch() || getCatchAll(); }
   };
 
   // The stack of block contexts currently being parsed.
@@ -541,7 +564,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Index blockHint = 0;
   Index labelHint = 0;
 
-  void pushScope(ScopeCtx scope) {
+  Result<> pushScope(ScopeCtx&& scope) {
     if (auto label = scope.getOriginalLabel()) {
       // Assign a fresh label to the scope, if necessary.
       if (!scope.label) {
@@ -554,7 +577,21 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
       scope.startPos = lastBinaryPos;
       lastBinaryPos = *binaryPos;
     }
-    scopeStack.push_back(scope);
+    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{};
   }
 
   ScopeCtx& getScope() {
@@ -610,6 +647,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<Type> getLabelType(Index label);
   Result<Type> getLabelType(Name labelName);
 
+  void fixLoopWithInput(Loop* loop, Type inputType, Index scratch);
+
   void dump();
 };
 

src/wasm/wasm-binary.cpp

@@ -2121,30 +2121,30 @@ bool WasmBinaryReader::getBasicHeapType(int64_t code, HeapType& out) {
   }
 }
 
-Type WasmBinaryReader::getType(int initial) {
-  // Single value types are negative; signature indices are non-negative
-  if (initial >= 0) {
-    // TODO: Handle block input types properly.
-    auto sig = getSignatureByTypeIndex(initial);
-    if (sig.params != Type::none) {
-      throwError("control flow inputs are not supported yet");
-    }
-    return sig.results;
+Signature WasmBinaryReader::getBlockType() {
+  // Single value types are negative; signature indices are non-negative.
+  auto code = getS32LEB();
+  if (code >= 0) {
+    return getSignatureByTypeIndex(code);
+  }
+  if (code == BinaryConsts::EncodedType::Empty) {
+    return Signature();
   }
+  return Signature(Type::none, getType(code));
+}
+
+Type WasmBinaryReader::getType(int code) {
   Type type;
-  if (getBasicType(initial, type)) {
+  if (getBasicType(code, type)) {
     return type;
   }
-  switch (initial) {
-    // None only used for block signatures. TODO: Separate out?
-    case BinaryConsts::EncodedType::Empty:
-      return Type::none;
+  switch (code) {
     case BinaryConsts::EncodedType::nullable:
       return Type(getHeapType(), Nullable);
     case BinaryConsts::EncodedType::nonnullable:
       return Type(getHeapType(), NonNullable);
     default:
-      throwError("invalid wasm type: " + std::to_string(initial));
+      throwError("invalid wasm type: " + std::to_string(code));
   }
   WASM_UNREACHABLE("unexpected type");
 }
@@ -2885,11 +2885,11 @@ Result<> WasmBinaryReader::readInst() {
   uint8_t code = getInt8();
   switch (code) {
     case BinaryConsts::Block:
-      return builder.makeBlock(Name(), getType());
+      return builder.makeBlock(Name(), getBlockType());
     case BinaryConsts::If:
-      return builder.makeIf(Name(), getType());
+      return builder.makeIf(Name(), getBlockType());
     case BinaryConsts::Loop:
-      return builder.makeLoop(Name(), getType());
+      return builder.makeLoop(Name(), getBlockType());
     case BinaryConsts::Br:
       return builder.makeBreak(getU32LEB(), false);
     case BinaryConsts::BrIf:
@@ -2974,9 +2974,9 @@ Result<> WasmBinaryReader::readInst() {
     case BinaryConsts::TableSet:
       return builder.makeTableSet(getTableName(getU32LEB()));
     case BinaryConsts::Try:
-      return builder.makeTry(Name(), getType());
+      return builder.makeTry(Name(), getBlockType());
     case BinaryConsts::TryTable: {
-      auto type = getType();
+      auto type = getBlockType();
       std::vector<Name> tags;
       std::vector<Index> labels;
       std::vector<bool> isRefs;