[StackIR] Support control flow input parameters

src/ir/module-utils.cpp

@@ -24,6 +24,7 @@
 #include "support/topological_sort.h"
 #include "support/utilities.h"
 #include "wasm-builder.h"
+#include "wasm-stack.h"
 
 namespace wasm::ModuleUtils {
 
@@ -398,9 +399,10 @@ struct TypeInfos {
     }
   }
   void noteControlFlow(Signature sig) {
-    // TODO: support control flow input parameters.
-    assert(sig.params.size() == 0);
-    if (sig.results.isTuple()) {
+    if (sig.results == Type::unreachable) {
+      sig.results = Type::none;
+    }
+    if (sig.params != Type::none || sig.results.isTuple()) {
       // We have to use a function type.
       ++controlFlowSignatures[sig];
     } else if (sig.results != Type::none) {
@@ -417,9 +419,13 @@ using ReferencedFuncs = std::unordered_map<Name, std::atomic<bool>>;
 struct CodeScanner : PostWalker<CodeScanner> {
   TypeInfos& info;
   ReferencedFuncs& referencedFuncs;
+  bool hasStackIR;
 
-  CodeScanner(Module& wasm, TypeInfos& info, ReferencedFuncs& referencedFuncs)
-    : info(info), referencedFuncs(referencedFuncs) {
+  CodeScanner(Module& wasm,
+              TypeInfos& info,
+              ReferencedFuncs& referencedFuncs,
+              bool hasStackIR = false)
+    : info(info), referencedFuncs(referencedFuncs), hasStackIR(hasStackIR) {
     setModule(&wasm);
   }
 
@@ -475,19 +481,29 @@ struct CodeScanner : PostWalker<CodeScanner> {
     info.note(curr->type);
   }
   void visitBlock(Block* curr) {
-    info.noteControlFlow(Signature(Type::none, curr->type));
+    if (!hasStackIR) {
+      info.noteControlFlow(Signature(Type::none, curr->type));
+    }
   }
   void visitIf(If* curr) {
-    info.noteControlFlow(Signature(Type::none, curr->type));
+    if (!hasStackIR) {
+      info.noteControlFlow(Signature(Type::none, curr->type));
+    }
   }
   void visitLoop(Loop* curr) {
-    info.noteControlFlow(Signature(Type::none, curr->type));
+    if (!hasStackIR) {
+      info.noteControlFlow(Signature(Type::none, curr->type));
+    }
   }
   void visitTry(Try* curr) {
-    info.noteControlFlow(Signature(Type::none, curr->type));
+    if (!hasStackIR) {
+      info.noteControlFlow(Signature(Type::none, curr->type));
+    }
   }
   void visitTryTable(TryTable* curr) {
-    info.noteControlFlow(Signature(Type::none, curr->type));
+    if (!hasStackIR) {
+      info.noteControlFlow(Signature(Type::none, curr->type));
+    }
   }
   void visitRefFunc(RefFunc* curr) { referencedFuncs.at(curr->func) = true; }
 };
@@ -503,7 +519,8 @@ InsertOrderedMap<HeapType, HeapTypeInfo>
 collectHeapTypeInfo(Module& wasm,
                     WorldMode worldMode,
                     TypeInclusion inclusion,
-                    VisibilityHandling visibility) {
+                    VisibilityHandling visibility,
+                    const ModuleStackIR* stackIR) {
   // Collect module-level info.
   TypeInfos info;
   ReferencedFuncs referencedFuncs;
@@ -531,11 +548,27 @@ collectHeapTypeInfo(Module& wasm,
       for (auto type : func->vars) {
         info.note(type);
       }
+      auto* funcStackIR = stackIR ? stackIR->getStackIROrNull(func) : nullptr;
+      if (funcStackIR) {
+        for (auto* inst : *funcStackIR) {
+          if (!inst) {
+            continue;
+          }
+          if (inst->op == StackInst::BlockBegin ||
+              inst->op == StackInst::IfBegin ||
+              inst->op == StackInst::LoopBegin ||
+              inst->op == StackInst::TryBegin ||
+              inst->op == StackInst::TryTableBegin) {
+            info.noteControlFlow(Signature(inst->params, inst->origin->type));
+          }
+        }
+      }
       // Don't just use `func->imported()` here because we also might be
       // printing an error message on a partially parsed module whose declared
       // function bodies have not all been parsed yet.
       if (func->body) {
-        CodeScanner(wasm, info, referencedFuncs).walk(func->body);
+        CodeScanner(wasm, info, referencedFuncs, funcStackIR != nullptr)
+          .walk(func->body);
       }
     });
 
@@ -975,9 +1008,13 @@ std::vector<HeapType> getPrivateHeapTypes(Module& wasm, WorldMode worldMode) {
   return types;
 }
 
-IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm) {
-  auto counts =
-    collectHeapTypeInfo(wasm, WorldMode::Open, TypeInclusion::BinaryTypes);
+IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm,
+                                              const ModuleStackIR* stackIR) {
+  auto counts = collectHeapTypeInfo(wasm,
+                                    WorldMode::Open,
+                                    TypeInclusion::BinaryTypes,
+                                    VisibilityHandling::NoVisibility,
+                                    stackIR);
 
   // Collect the rec groups.
   std::unordered_map<RecGroup, size_t> groupIndices;

src/ir/module-utils.h

@@ -22,6 +22,10 @@
 #include "support/unique_deferring_queue.h"
 #include "wasm.h"
 
+namespace wasm {
+class ModuleStackIR;
+} // namespace wasm
+
 namespace wasm::ModuleUtils {
 
 // Copies a function into a module. If newName is provided it is used as the
@@ -475,7 +479,8 @@ InsertOrderedMap<HeapType, HeapTypeInfo> collectHeapTypeInfo(
   Module& wasm,
   WorldMode worldMode,
   TypeInclusion inclusion = TypeInclusion::AllTypes,
-  VisibilityHandling visibility = VisibilityHandling::NoVisibility);
+  VisibilityHandling visibility = VisibilityHandling::NoVisibility,
+  const ModuleStackIR* stackIR = nullptr);
 
 // Helper function for collecting all the non-basic heap types used in the
 // module, i.e. the types that would appear in the type section.
@@ -503,7 +508,9 @@ struct IndexedHeapTypes {
 // Similar to `collectHeapTypes`, but provides fast lookup of the index for each
 // type as well. Also orders the types to be valid and sorts the types by
 // frequency of use to minimize code size.
-IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm);
+IndexedHeapTypes
+getOptimizedIndexedHeapTypes(Module& wasm,
+                             const ModuleStackIR* stackIR = nullptr);
 
 } // namespace wasm::ModuleUtils
 

src/passes/Print.cpp

@@ -252,18 +252,21 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
   }
 
   std::ostream& printBlockType(Signature sig) {
-    assert(sig.params == Type::none);
-    if (sig.results == Type::none) {
-      return o;
-    }
-    if (sig.results.isTuple()) {
+    if (sig.params != Type::none || sig.results.isTuple()) {
       if (auto it = signatureTypes.find(sig); it != signatureTypes.end()) {
-        o << "(type ";
+        o << " (type ";
         printHeapTypeName(it->second);
-        o << ") ";
+        o << ')';
       }
     }
-    printResultType(sig.results);
+    if (sig.params != Type::none) {
+      o << ' ';
+      printParamType(sig.params);
+    }
+    if (sig.results != Type::none) {
+      o << ' ';
+      printResultType(sig.results);
+    }
     return o;
   }
 
@@ -362,6 +365,33 @@ struct PrintExpressionContents
       currFunction(parent.currFunction), o(parent.o),
       features(wasm ? wasm->features : FeatureSet::All), full(isFullForced()) {}
 
+  using OverriddenVisitor<PrintExpressionContents>::visit;
+
+  void visit(StackInst* inst) {
+    switch (inst->op) {
+      case StackInst::Basic:
+        visit(inst->origin);
+        break;
+      case StackInst::BlockBegin:
+        visitBlock(inst->origin->cast<Block>(), inst->params);
+        break;
+      case StackInst::IfBegin:
+        visitIf(inst->origin->cast<If>(), inst->params);
+        break;
+      case StackInst::LoopBegin:
+        visitLoop(inst->origin->cast<Loop>(), inst->params);
+        break;
+      case StackInst::TryBegin:
+        visitTry(inst->origin->cast<Try>(), inst->params);
+        break;
+      case StackInst::TryTableBegin:
+        visitTryTable(inst->origin->cast<TryTable>(), inst->params);
+        break;
+      default:
+        WASM_UNREACHABLE("unexpected op");
+    }
+  }
+
   std::ostream& printType(Type type) { return parent.printType(type); }
 
   std::ostream& printHeapTypeName(HeapType type) {
@@ -376,8 +406,11 @@ struct PrintExpressionContents
     return parent.printParamType(type);
   }
 
-  std::ostream& printBlockType(Signature sig) {
-    return parent.printBlockType(sig);
+  void printBlockType(Type results, Type params = Type::none) {
+    if (!results.isConcrete()) {
+      results = Type::none;
+    }
+    parent.printBlockType(Signature(params, results));
   }
 
   void printMemoryPostfix(uint8_t bytes, Type type) {
@@ -407,18 +440,15 @@ struct PrintExpressionContents
     return o;
   }
 
-  void visitBlock(Block* curr) {
+  void visitBlock(Block* curr, Type params = Type::none) {
     printMedium(o, "block");
     if (curr->name.is()) {
       o << ' ';
       curr->name.print(o);
     }
-    if (curr->type.isConcrete()) {
-      o << ' ';
-      printBlockType(Signature(Type::none, curr->type));
-    }
+    printBlockType(curr->type, params);
   }
-  void visitIf(If* curr) {
+  void visitIf(If* curr, Type params = Type::none) {
     printMedium(o, "if");
     // Ifs are unreachable if their condition is unreachable, but in that case
     // the arms might have some concrete type we have to account for to produce
@@ -427,21 +457,15 @@ struct PrintExpressionContents
     if (curr->condition->type == Type::unreachable && curr->ifFalse) {
       type = Type::getLeastUpperBound(curr->ifTrue->type, curr->ifFalse->type);
     }
-    if (type.isConcrete()) {
-      o << ' ';
-      printBlockType(Signature(Type::none, type));
-    }
+    printBlockType(type, params);
   }
-  void visitLoop(Loop* curr) {
+  void visitLoop(Loop* curr, Type params = Type::none) {
     printMedium(o, "loop");
     if (curr->name.is()) {
       o << ' ';
       curr->name.print(o);
     }
-    if (curr->type.isConcrete()) {
-      o << ' ';
-      printBlockType(Signature(Type::none, curr->type));
-    }
+    printBlockType(curr->type, params);
   }
   void visitBreak(Break* curr) {
     if (curr->condition) {
@@ -2136,23 +2160,17 @@ struct PrintExpressionContents
     printMedium(o, "elem.drop ");
     curr->segment.print(o);
   }
-  void visitTry(Try* curr) {
+  void visitTry(Try* curr, Type params = Type::none) {
     printMedium(o, "try");
     if (curr->name.is()) {
       o << ' ';
       curr->name.print(o);
     }
-    if (curr->type.isConcrete()) {
-      o << ' ';
-      printBlockType(Signature(Type::none, curr->type));
-    }
+    printBlockType(curr->type, params);
   }
-  void visitTryTable(TryTable* curr) {
+  void visitTryTable(TryTable* curr, Type params = Type::none) {
     printMedium(o, "try_table");
-    if (curr->type.isConcrete()) {
-      o << ' ';
-      printBlockType(Signature(Type::none, curr->type));
-    }
+    printBlockType(curr->type, params);
     for (Index i = 0; i < curr->catchTags.size(); i++) {
       o << " (";
       if (curr->catchTags[i]) {
@@ -2713,7 +2731,9 @@ struct PrintExpressionContents
 void PrintSExpression::setModule(Module* module) {
   currModule = module;
   if (module) {
-    heapTypes = ModuleUtils::getOptimizedIndexedHeapTypes(*module).types;
+    heapTypes = ModuleUtils::getOptimizedIndexedHeapTypes(
+                  *module, moduleStackIR ? &*moduleStackIR : nullptr)
+                  .types;
     for (auto type : heapTypes) {
       if (type.isSignature()) {
         signatureTypes.insert({type.getSignature(), type});
@@ -3842,14 +3862,16 @@ printStackInst(StackInst* inst, std::ostream& o, Function* func) {
     case StackInst::BlockBegin:
     case StackInst::IfBegin:
     case StackInst::LoopBegin:
-    case StackInst::TryBegin: {
-      PrintExpressionContents(printer).visit(inst->origin);
+    case StackInst::TryBegin:
+    case StackInst::TryTableBegin: {
+      PrintExpressionContents(printer).visit(inst);
       break;
     }
     case StackInst::BlockEnd:
     case StackInst::IfEnd:
     case StackInst::LoopEnd:
-    case StackInst::TryEnd: {
+    case StackInst::TryEnd:
+    case StackInst::TryTableEnd: {
       printMedium(o, "end");
       o << " ;; type: ";
       printer.printType(inst->type);
@@ -3914,7 +3936,7 @@ static std::ostream& printStackIR(StackIR* ir, PrintSExpression& printer) {
       case StackInst::TryTableBegin: {
         controlFlowDepth++;
         doIndent();
-        PrintExpressionContents(printer).visit(inst->origin);
+        PrintExpressionContents(printer).visit(inst);
         indent++;
         break;
       }

src/wasm-binary.h

@@ -1309,6 +1309,8 @@ enum FeaturePrefix { FeatureUsed = '+', FeatureDisallowed = '-' };
 // (local index in IR, tuple index) => binary local index
 using MappedLocals = std::unordered_map<std::pair<Index, Index>, size_t>;
 
+class ModuleStackIR;
+
 // Writes out wasm to the binary format
 
 class WasmBinaryWriter {
@@ -1383,10 +1385,8 @@ class WasmBinaryWriter {
 public:
   WasmBinaryWriter(Module* input,
                    BufferWithRandomAccess& o,
-                   const PassOptions& options)
-    : wasm(input), o(o), options(options), indexes(*input) {
-    prepare();
-  }
+                   const PassOptions& options);
+  ~WasmBinaryWriter();
 
   // locations in the output binary for the various parts of the module
   struct TableOfContents {
@@ -1552,6 +1552,7 @@ class WasmBinaryWriter {
   Function::DebugLocation lastDebugLocation;
 
   std::unique_ptr<ImportInfo> importInfo;
+  std::unique_ptr<ModuleStackIR> moduleStackIR;
 
   // General debugging info: track locations as we write.
   BinaryLocations binaryLocations;