[Parser] Parse br_table (#6098)

4 files changed, 68 insertions, 8 deletions

diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 6a3ab9614..485c596e7 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -383,6 +383,9 @@ struct NullInstrParserCtx {
   Result<> makeMemoryFill(Index, MemoryIdxT*) { return Ok{}; }
   Result<> makeCall(Index, FuncIdxT, bool) { return Ok{}; }
   Result<> makeBreak(Index, LabelIdxT) { return Ok{}; }
+  Result<> makeSwitch(Index, const std::vector<LabelIdxT>&, LabelIdxT) {
+    return Ok{};
+  }
   Result<> makeReturn(Index) { return Ok{}; }
   template<typename HeapTypeT> Result<> makeRefNull(Index, HeapTypeT) {
     return Ok{};
@@ -1246,6 +1249,11 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     return withLoc(pos, irBuilder.makeBreak(label));
   }
 
+  Result<>
+  makeSwitch(Index pos, const std::vector<Index> labels, Index defaultLabel) {
+    return withLoc(pos, irBuilder.makeSwitch(labels, defaultLabel));
+  }
+
   Result<> makeReturn(Index pos) {
     return withLoc(pos, irBuilder.makeReturn());
   }
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index eb93fa717..0641e941c 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -1200,7 +1200,20 @@ template<typename Ctx> Result<> makeBreak(Ctx& ctx, Index pos) {
 }
 
 template<typename Ctx> Result<> makeBreakTable(Ctx& ctx, Index pos) {
-  return ctx.in.err("unimplemented instruction");
+  std::vector<typename Ctx::LabelIdxT> labels;
+  while (true) {
+    // Parse at least one label; return an error only if we parse none.
+    auto label = labelidx(ctx);
+    if (labels.empty()) {
+      CHECK_ERR(label);
+    } else if (label.getErr()) {
+      break;
+    }
+    labels.push_back(*label);
+  }
+  auto defaultLabel = labels.back();
+  labels.pop_back();
+  return ctx.makeSwitch(pos, labels, defaultLabel);
 }
 
 template<typename Ctx> Result<> makeReturn(Ctx& ctx, Index pos) {
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 2fc03326b..6d4fca3c3 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -76,7 +76,8 @@ public:
   [[nodiscard]] Result<> makeIf(Name label, Type type);
   [[nodiscard]] Result<> makeLoop(Name label, Type type);
   [[nodiscard]] Result<> makeBreak(Index label);
-  // [[nodiscard]] Result<> makeSwitch();
+  [[nodiscard]] Result<> makeSwitch(const std::vector<Index>& labels,
+                                    Index defaultLabel);
   // Unlike Builder::makeCall, this assumes the function already exists.
   [[nodiscard]] Result<> makeCall(Name func, bool isReturn);
   // [[nodiscard]] Result<> makeCallIndirect();
@@ -191,6 +192,8 @@ public:
   [[nodiscard]] Result<> visitArrayNew(ArrayNew*);
   [[nodiscard]] Result<> visitBreak(Break*,
                                     std::optional<Index> label = std::nullopt);
+  [[nodiscard]] Result<>
+  visitSwitch(Switch*, std::optional<Index> defaultLabel = std::nullopt);
   [[nodiscard]] Result<> visitCall(Call*);
 
 private:
@@ -392,6 +395,9 @@ private:
   // the value, if they are different. May only be called directly after
   // hoistLastValue().
   [[nodiscard]] Result<> packageHoistedValue(const HoistedVal&);
+
+  [[nodiscard]] Result<Expression*> getBranchValue(Name labelName,
+                                                   std::optional<Index> label);
 };
 
 } // namespace wasm
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 0c5f0830e..cbb16ff68 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -290,9 +290,10 @@ Result<> IRBuilder::visitArrayNew(ArrayNew* curr) {
   return Ok{};
 }
 
-Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
+Result<Expression*> IRBuilder::getBranchValue(Name labelName,
+                                              std::optional<Index> label) {
   if (!label) {
-    auto index = getLabelIndex(curr->name);
+    auto index = getLabelIndex(labelName);
     CHECK_ERR(index);
     label = *index;
   }
@@ -305,12 +306,29 @@ Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
     values[size - 1 - i] = *val;
   }
   if (values.size() == 0) {
-    curr->value = nullptr;
+    return nullptr;
   } else if (values.size() == 1) {
-    curr->value = values[0];
+    return values[0];
   } else {
-    curr->value = builder.makeTupleMake(values);
+    return builder.makeTupleMake(values);
   }
+}
+
+Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
+  auto value = getBranchValue(curr->name, label);
+  CHECK_ERR(value);
+  curr->value = *value;
+  return Ok{};
+}
+
+Result<> IRBuilder::visitSwitch(Switch* curr,
+                                std::optional<Index> defaultLabel) {
+  auto cond = pop();
+  CHECK_ERR(cond);
+  curr->condition = *cond;
+  auto value = getBranchValue(curr->default_, defaultLabel);
+  CHECK_ERR(value);
+  curr->value = *value;
   return Ok{};
 }
 
@@ -556,7 +574,22 @@ Result<> IRBuilder::makeBreak(Index label) {
   return Ok{};
 }
 
-// Result<> IRBuilder::makeSwitch() {}
+Result<> IRBuilder::makeSwitch(const std::vector<Index>& labels,
+                               Index defaultLabel) {
+  std::vector<Name> names;
+  names.reserve(labels.size());
+  for (auto label : labels) {
+    auto name = getLabelName(label);
+    CHECK_ERR(name);
+    names.push_back(*name);
+  }
+  auto defaultName = getLabelName(defaultLabel);
+  CHECK_ERR(defaultName);
+  Switch curr(wasm.allocator);
+  CHECK_ERR(visitSwitch(&curr, defaultLabel));
+  push(builder.makeSwitch(names, *defaultName, curr.condition, curr.value));
+  return Ok{};
+}
 
 Result<> IRBuilder::makeCall(Name func, bool isReturn) {
   Call curr(wasm.allocator);