[Outlining] Fixes break reconstruction (#6352)

Adds new visitBreakWithType and visitSwitchWithType functions to the IRBuilder API. These functions work around an assumption in IRBuilder that the module is being traversed in the fully nested format, i.e., that the destination scope of a break or switch has been visited before visiting the break or switch. Instead, the type of the destination scope is passed to IRBuilder.

1 files changed, 43 insertions, 3 deletions

diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index b1bf8c855..8d87d0f33 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -419,8 +419,14 @@ Result<> IRBuilder::visitArrayNewFixed(ArrayNewFixed* curr) {
   return Ok{};
 }
 
-Result<Expression*> IRBuilder::getBranchValue(Name labelName,
+Result<Expression*> IRBuilder::getBranchValue(Expression* curr,
+                                              Name labelName,
                                               std::optional<Index> label) {
+  // As new branch instructions are added, one of the existing branch visit*
+  // functions is likely to be copied, along with its call to getBranchValue().
+  // This assert serves as a reminder to also add an implementation of
+  // visit*WithType() for new branch instructions.
+  assert(curr->is<Break>() || curr->is<Switch>());
   if (!label) {
     auto index = getLabelIndex(labelName);
     CHECK_ERR(index);
@@ -440,23 +446,57 @@ Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
     CHECK_ERR(cond);
     curr->condition = *cond;
   }
-  auto value = getBranchValue(curr->name, label);
+  auto value = getBranchValue(curr, curr->name, label);
   CHECK_ERR(value);
   curr->value = *value;
   return Ok{};
 }
 
+Result<> IRBuilder::visitBreakWithType(Break* curr, Type type) {
+  if (curr->condition) {
+    auto cond = pop();
+    CHECK_ERR(cond);
+    curr->condition = *cond;
+  }
+  if (type == Type::none) {
+    curr->value = nullptr;
+  } else {
+    auto value = pop(type.size());
+    CHECK_ERR(value)
+    curr->value = *value;
+  }
+  curr->finalize();
+  push(curr);
+  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);
+  auto value = getBranchValue(curr, curr->default_, defaultLabel);
   CHECK_ERR(value);
   curr->value = *value;
   return Ok{};
 }
 
+Result<> IRBuilder::visitSwitchWithType(Switch* curr, Type type) {
+  auto cond = pop();
+  CHECK_ERR(cond);
+  curr->condition = *cond;
+  if (type == Type::none) {
+    curr->value = nullptr;
+  } else {
+    auto value = pop(type.size());
+    CHECK_ERR(value)
+    curr->value = *value;
+  }
+  curr->finalize();
+  push(curr);
+  return Ok{};
+}
+
 Result<> IRBuilder::visitCall(Call* curr) {
   auto numArgs = wasm.getFunction(curr->target)->getNumParams();
   curr->operands.resize(numArgs);