label fix

src/passes/Outlining.cpp

diff --git a/src/passes/Outlining.cpp b/src/passes/Outlining.cpp
index 345b3f9..4a7b9bf 100644
--- a/src/passes/Outlining.cpp
+++ b/src/passes/Outlining.cpp
@@ -132,7 +132,15 @@
                          : state == NotInSeq ? &existingBuilder
                                              : nullptr;
     if (builder) {
-      ASSERT_OK(builder->visit(curr));
+      if (auto* expr = curr->dynCast<Break>()) {
+        Type type = expr->value ? expr->value->type : Type::none;
+        ASSERT_OK(builder->visitBreakWithType(expr, type));
+      } else if (auto* expr = curr->dynCast<Switch>()) {
+        Type type = expr->value ? expr->value->type : Type::none;
+        ASSERT_OK(builder->visitSwitchWithType(expr, type));
+      } else {
+        ASSERT_OK(builder->visit(curr));
+      }
     }
     DBG(printVisitExpression(curr));
 

src/wasm-ir-builder.h

diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 05c59ca..d647ecc 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -222,8 +222,10 @@
   [[nodiscard]] Result<> visitArrayNewFixed(ArrayNewFixed*);
   [[nodiscard]] Result<> visitBreak(Break*,
                                     std::optional<Index> label = std::nullopt);
+  [[nodiscard]] Result<> visitBreakWithType(Break*, Type);
   [[nodiscard]] Result<>
   visitSwitch(Switch*, std::optional<Index> defaultLabel = std::nullopt);
+  [[nodiscard]] Result<> visitSwitchWithType(Switch*, Type);
   [[nodiscard]] Result<> visitCall(Call*);
   [[nodiscard]] Result<> visitCallIndirect(CallIndirect*);
   [[nodiscard]] Result<> visitCallRef(CallRef*);

src/wasm/wasm-ir-builder.cpp

diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 1df441d..102551b 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -431,6 +431,26 @@
   return Ok{};
 }
 
+Result<> IRBuilder::visitBreakWithType(Break* curr, Type type) {
+  if (curr->condition) {
+    auto cond = pop();
+    CHECK_ERR(cond);
+    curr->condition = *cond;
+  }
+  if (!curr->value) {
+    curr->value = nullptr;
+  } else {
+    auto value = pop(type.size());
+    CHECK_ERR(value)
+    curr->value = *value;
+  }
+  // TODO: Call more efficient versions of finalize() that take the known type
+  // for other kinds of nodes as well, as done above.
+  ReFinalizeNode{}.visit(curr);
+  push(curr);
+  return Ok{};
+}
+
 Result<> IRBuilder::visitSwitch(Switch* curr,
                                 std::optional<Index> defaultLabel) {
   auto cond = pop();
@@ -442,6 +462,24 @@
   return Ok{};
 }
 
+Result<> IRBuilder::visitSwitchWithType(Switch* curr, Type type) {
+  auto cond = pop();
+  CHECK_ERR(cond);
+  curr->condition = *cond;
+  if (!curr->value) {
+    curr->value = nullptr;
+  } else {
+    auto value = pop(type.size());
+    CHECK_ERR(value)
+    curr->value = *value;
+  }
+  // TODO: Call more efficient versions of finalize() that take the known type
+  // for other kinds of nodes as well, as done above.
+  ReFinalizeNode{}.visit(curr);
+  push(curr);
+  return Ok{};
+}
+
 Result<> IRBuilder::visitCall(Call* curr) {
   auto numArgs = wasm.getFunction(curr->target)->getNumParams();
   curr->operands.resize(numArgs);

test/lit/passes/outlining.wast

diff --git a/test/lit/passes/outlining.wast b/test/lit/passes/outlining.wast
index befce75..19ec4ce 100644
--- a/test/lit/passes/outlining.wast
+++ b/test/lit/passes/outlining.wast
@@ -614,6 +614,48 @@
   )
 )
 
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (func $outline$ (type $0)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 1)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+
+  ;; CHECK:      (func $a (type $0)
+  ;; CHECK-NEXT:  (block $label1
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:   (loop $loop-in
+  ;; CHECK-NEXT:    (br $label1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $outline$)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $a
+    (block $label1
+      (drop
+        (i32.const 2)
+      )
+      (drop
+        (i32.const 1)
+      )
+      (loop
+        (br $label1)
+      )
+      (drop
+        (i32.const 2)
+      )
+      (drop
+        (i32.const 1)
+      )
+    )
+  )
+)
+
 ;; Tests return instructions are correctly filtered from being outlined.
 (module
   ;; CHECK:      (type $0 (func (result i32)))