Support control flow inputs in IRBuilder

CHANGELOG.md

@@ -17,6 +17,7 @@ Current Trunk
 
  - BinaryenSelect no longer takes a type parameter.
  - AutoDrop APIs have been removed.
+ - Binaryen now supports parsing control flow structures with parameter types.
 
 v120
 ----

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;