[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, 13 insertions, 1 deletions

diff --git a/src/passes/Outlining.cpp b/src/passes/Outlining.cpp
index 345b3f9a2..68c3d0397 100644
--- a/src/passes/Outlining.cpp
+++ b/src/passes/Outlining.cpp
@@ -132,7 +132,19 @@ struct ReconstructStringifyWalker
                          : 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 ensures new unhandled branch instructions
+        // will quickly cause an error. Serves as a reminder to
+        // implement a new special-case visit*WithType.
+        assert(curr->is<BrOn>() || !Properties::isBranch(curr));
+        ASSERT_OK(builder->visit(curr));
+      }
     }
     DBG(printVisitExpression(curr));