[Parser] Parse try_table (#6237)

4 files changed, 190 insertions, 6 deletions

diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 7e9add5a7..8f74502ad 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -297,6 +297,8 @@ template<typename Ctx> struct TypeParserCtx {
 
 struct NullInstrParserCtx {
   using ExprT = Ok;
+  using CatchT = Ok;
+  using CatchListT = Ok;
 
   using FieldIdxT = Ok;
   using FuncIdxT = Ok;
@@ -362,6 +364,17 @@ struct NullInstrParserCtx {
   Result<> visitDelegate(Index, LabelIdxT) { return Ok{}; }
   Result<> visitEnd() { return Ok{}; }
 
+  CatchListT makeCatchList() { return Ok{}; }
+  void appendCatch(CatchListT&, CatchT) {}
+  CatchT makeCatch(TagIdxT, LabelIdxT) { return Ok{}; }
+  CatchT makeCatchRef(TagIdxT, LabelIdxT) { return Ok{}; }
+  CatchT makeCatchAll(LabelIdxT) { return Ok{}; }
+  CatchT makeCatchAllRef(LabelIdxT) { return Ok{}; }
+  template<typename BlockTypeT>
+  Result<> makeTryTable(Index, std::optional<Name>, BlockTypeT, CatchListT) {
+    return Ok{};
+  }
+
   Result<> makeUnreachable(Index) { return Ok{}; }
   Result<> makeNop(Index) { return Ok{}; }
   Result<> makeBinary(Index, BinaryOp) { return Ok{}; }
@@ -1023,6 +1036,10 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
   using ExprT = Expression*;
   using ElemListT = std::vector<Expression*>;
 
+  struct CatchInfo;
+  using CatchT = CatchInfo;
+  using CatchListT = std::vector<CatchInfo>;
+
   using FieldIdxT = Index;
   using FuncIdxT = Name;
   using LocalIdxT = Index;
@@ -1106,6 +1123,21 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
 
   TableTypeT makeTableType(LimitsT, Type) { return Ok{}; }
 
+  struct CatchInfo {
+    Name tag;
+    Index label;
+    bool isRef;
+  };
+
+  std::vector<CatchInfo> makeCatchList() { return {}; }
+  void appendCatch(std::vector<CatchInfo>& list, CatchInfo info) {
+    list.push_back(info);
+  }
+  CatchInfo makeCatch(Name tag, Index label) { return {tag, label, false}; }
+  CatchInfo makeCatchRef(Name tag, Index label) { return {tag, label, true}; }
+  CatchInfo makeCatchAll(Index label) { return {{}, label, false}; }
+  CatchInfo makeCatchAllRef(Index label) { return {{}, label, true}; }
+
   Result<HeapTypeT> getHeapTypeFromIdx(Index idx) {
     if (idx >= types.size()) {
       return in.err("type index out of bounds");
@@ -1374,6 +1406,26 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
       irBuilder.makeTry(label ? *label : Name{}, type.getSignature().results));
   }
 
+  Result<> makeTryTable(Index pos,
+                        std::optional<Name> label,
+                        HeapType type,
+                        const std::vector<CatchInfo>& info) {
+    std::vector<Name> tags;
+    std::vector<Index> labels;
+    std::vector<bool> isRefs;
+    for (auto& info : info) {
+      tags.push_back(info.tag);
+      labels.push_back(info.label);
+      isRefs.push_back(info.isRef);
+    }
+    return withLoc(pos,
+                   irBuilder.makeTryTable(label ? *label : Name{},
+                                          type.getSignature().results,
+                                          tags,
+                                          labels,
+                                          isRefs));
+  }
+
   Result<> visitCatch(Index pos, Name tag) {
     return withLoc(pos, irBuilder.visitCatch(tag));
   }
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index 10894a7cf..73a7d0035 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -60,6 +60,8 @@ template<typename Ctx> MaybeResult<> block(Ctx&, bool);
 template<typename Ctx> MaybeResult<> ifelse(Ctx&, bool);
 template<typename Ctx> MaybeResult<> loop(Ctx&, bool);
 template<typename Ctx> MaybeResult<> trycatch(Ctx&, bool);
+template<typename Ctx> MaybeResult<typename Ctx::CatchT> catchinstr(Ctx&);
+template<typename Ctx> MaybeResult<> trytable(Ctx&, bool);
 template<typename Ctx> Result<> makeUnreachable(Ctx&, Index);
 template<typename Ctx> Result<> makeNop(Ctx&, Index);
 template<typename Ctx> Result<> makeBinary(Ctx&, Index, BinaryOp op);
@@ -657,7 +659,7 @@ template<typename Ctx> Result<uint32_t> tupleArity(Ctx& ctx) {
 // Instructions
 // ============
 
-// blockinstr ::= block | loop | if-else | try-catch
+// blockinstr ::= block | loop | if-else | try-catch | try_table
 template<typename Ctx> MaybeResult<> foldedBlockinstr(Ctx& ctx) {
   if (auto i = block(ctx, true)) {
     return i;
@@ -671,7 +673,9 @@ template<typename Ctx> MaybeResult<> foldedBlockinstr(Ctx& ctx) {
   if (auto i = trycatch(ctx, true)) {
     return i;
   }
-  // TODO: Other block instructions
+  if (auto i = trytable(ctx, true)) {
+    return i;
+  }
   return {};
 }
 
@@ -688,7 +692,9 @@ template<typename Ctx> MaybeResult<> unfoldedBlockinstr(Ctx& ctx) {
   if (auto i = trycatch(ctx, false)) {
     return i;
   }
-  // TODO: Other block instructions
+  if (auto i = trytable(ctx, false)) {
+    return i;
+  }
   return {};
 }
 
@@ -1159,6 +1165,81 @@ template<typename Ctx> MaybeResult<> trycatch(Ctx& ctx, bool folded) {
   return ctx.visitEnd();
 }
 
+template<typename Ctx> MaybeResult<typename Ctx::CatchT> catchinstr(Ctx& ctx) {
+  typename Ctx::CatchT result;
+  if (ctx.in.takeSExprStart("catch"sv)) {
+    auto tag = tagidx(ctx);
+    CHECK_ERR(tag);
+    auto label = labelidx(ctx);
+    CHECK_ERR(label);
+    result = ctx.makeCatch(*tag, *label);
+  } else if (ctx.in.takeSExprStart("catch_ref"sv)) {
+    auto tag = tagidx(ctx);
+    CHECK_ERR(tag);
+    auto label = labelidx(ctx);
+    CHECK_ERR(label);
+    result = ctx.makeCatchRef(*tag, *label);
+  } else if (ctx.in.takeSExprStart("catch_all"sv)) {
+    auto label = labelidx(ctx);
+    CHECK_ERR(label);
+    result = ctx.makeCatchAll(*label);
+  } else if (ctx.in.takeSExprStart("catch_all_ref"sv)) {
+    auto label = labelidx(ctx);
+    CHECK_ERR(label);
+    result = ctx.makeCatchAllRef(*label);
+  } else {
+    return {};
+  }
+
+  if (!ctx.in.takeRParen()) {
+    return ctx.in.err("expected ')' at end of catch clause");
+  }
+
+  return result;
+}
+
+// trytable ::= 'try_table' label blocktype catchinstr* instr* end id?
+//            | '(' 'try_table' label blocktype catchinstr* instr* ')'
+template<typename Ctx> MaybeResult<> trytable(Ctx& ctx, bool folded) {
+  auto pos = ctx.in.getPos();
+
+  if ((folded && !ctx.in.takeSExprStart("try_table"sv)) ||
+      (!folded && !ctx.in.takeKeyword("try_table"sv))) {
+    return {};
+  }
+
+  auto label = ctx.in.takeID();
+
+  auto type = blocktype(ctx);
+  CHECK_ERR(type);
+
+  auto catches = ctx.makeCatchList();
+  while (auto c = catchinstr(ctx)) {
+    CHECK_ERR(c);
+    ctx.appendCatch(catches, *c);
+  }
+
+  CHECK_ERR(ctx.makeTryTable(pos, label, *type, catches));
+
+  CHECK_ERR(instrs(ctx));
+
+  if (folded) {
+    if (!ctx.in.takeRParen()) {
+      return ctx.in.err("expected ')' at end of try_table");
+    }
+  } else {
+    if (!ctx.in.takeKeyword("end"sv)) {
+      return ctx.in.err("expected 'end' at end of try_table");
+    }
+
+    auto id = ctx.in.takeID();
+    if (id && id != label) {
+      return ctx.in.err("end label does not match try_table label");
+    }
+  }
+  return ctx.visitEnd();
+}
+
 template<typename Ctx> Result<> makeUnreachable(Ctx& ctx, Index pos) {
   return ctx.makeUnreachable(pos);
 }
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 4b5468014..343ef998e 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -65,6 +65,8 @@ public:
   [[nodiscard]] Result<> visitElse();
   [[nodiscard]] Result<> visitLoopStart(Loop* iff);
   [[nodiscard]] Result<> visitTryStart(Try* tryy, Name label = {});
+  [[nodiscard]] Result<> visitTryTableStart(TryTable* trytable,
+                                            Name label = {});
   [[nodiscard]] Result<> visitCatch(Name tag);
   [[nodiscard]] Result<> visitCatchAll();
   [[nodiscard]] Result<> visitDelegate(Index label);
@@ -155,7 +157,11 @@ public:
   [[nodiscard]] Result<> makeTableFill(Name table);
   [[nodiscard]] Result<> makeTableCopy(Name destTable, Name srcTable);
   [[nodiscard]] Result<> makeTry(Name label, Type type);
-  // [[nodiscard]] Result<> makeTryTable();
+  [[nodiscard]] Result<> makeTryTable(Name label,
+                                      Type type,
+                                      const std::vector<Name>& tags,
+                                      const std::vector<Index>& labels,
+                                      const std::vector<bool>& isRefs);
   [[nodiscard]] Result<> makeThrow(Name tag);
   [[nodiscard]] Result<> makeRethrow(Index label);
   // [[nodiscard]] Result<> makeThrowRef();
@@ -264,6 +270,10 @@ private:
       Try* tryy;
       Name originalLabel;
     };
+    struct TryTableScope {
+      TryTable* trytable;
+      Name originalLabel;
+    };
     using Scope = std::variant<NoScope,
                                FuncScope,
                                BlockScope,
@@ -272,7 +282,8 @@ private:
                                LoopScope,
                                TryScope,
                                CatchScope,
-                               CatchAllScope>;
+                               CatchAllScope,
+                               TryTableScope>;
 
     // The control flow structure we are building expressions for.
     Scope scope;
@@ -312,6 +323,9 @@ private:
     static ScopeCtx makeCatchAll(Try* tryy, Name originalLabel, Name label) {
       return ScopeCtx(CatchAllScope{tryy, originalLabel}, label);
     }
+    static ScopeCtx makeTryTable(TryTable* trytable, Name originalLabel = {}) {
+      return ScopeCtx(TryTableScope{trytable, originalLabel});
+    }
 
     bool isNone() { return std::get_if<NoScope>(&scope); }
     Function* getFunction() {
@@ -362,6 +376,12 @@ private:
       }
       return nullptr;
     }
+    TryTable* getTryTable() {
+      if (auto* tryTableScope = std::get_if<TryTableScope>(&scope)) {
+        return tryTableScope->trytable;
+      }
+      return nullptr;
+    }
     Type getResultType() {
       if (auto* func = getFunction()) {
         return func->type.getSignature().results;
@@ -387,6 +407,9 @@ private:
       if (auto* tryy = getCatchAll()) {
         return tryy->type;
       }
+      if (auto* trytable = getTryTable()) {
+        return trytable->type;
+      }
       WASM_UNREACHABLE("unexpected scope kind");
     }
     Name getOriginalLabel() {
@@ -414,6 +437,9 @@ private:
       if (auto* catchAllScope = std::get_if<CatchAllScope>(&scope)) {
         return catchAllScope->originalLabel;
       }
+      if (auto* tryTableScope = std::get_if<TryTableScope>(&scope)) {
+        return tryTableScope->originalLabel;
+      }
       WASM_UNREACHABLE("unexpected scope kind");
     }
   };
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index dc80d6af5..14b9046ef 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -624,6 +624,11 @@ Result<> IRBuilder::visitTryStart(Try* tryy, Name label) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitTryTableStart(TryTable* trytable, Name label) {
+  pushScope(ScopeCtx::makeTryTable(trytable, label));
+  return Ok{};
+}
+
 Result<Expression*> IRBuilder::finishScope(Block* block) {
   if (scopeStack.empty() || scopeStack.back().isNone()) {
     return Err{"unexpected end of scope"};
@@ -895,6 +900,10 @@ Result<> IRBuilder::visitEnd() {
     tryy->catchBodies.push_back(*expr);
     tryy->finalize(tryy->type);
     push(maybeWrapForLabel(tryy));
+  } else if (auto* trytable = scope.getTryTable()) {
+    trytable->body = *expr;
+    trytable->finalize(trytable->type, &wasm);
+    push(maybeWrapForLabel(trytable));
   } else {
     WASM_UNREACHABLE("unexpected scope kind");
   }
@@ -1347,7 +1356,23 @@ Result<> IRBuilder::makeTry(Name label, Type type) {
   return visitTryStart(tryy, label);
 }
 
-// Result<> IRBuilder::makeTryTable() {}
+Result<> IRBuilder::makeTryTable(Name label,
+                                 Type type,
+                                 const std::vector<Name>& tags,
+                                 const std::vector<Index>& labels,
+                                 const std::vector<bool>& isRefs) {
+  auto* trytable = wasm.allocator.alloc<TryTable>();
+  trytable->type = type;
+  trytable->catchTags.set(tags);
+  trytable->catchRefs.set(isRefs);
+  trytable->catchDests.reserve(labels.size());
+  for (auto label : labels) {
+    auto name = getLabelName(label);
+    CHECK_ERR(name);
+    trytable->catchDests.push_back(*name);
+  }
+  return visitTryTableStart(trytable, label);
+}
 
 Result<> IRBuilder::makeThrow(Name tag) {
   Throw curr(wasm.allocator);