Support function contexts in IRBuilder (#5967)

Add a `visitFunctionStart` function to IRBuilder and make it responsible for
setting the function's body when the context is closed. This will simplify
outlining, will be necessary to support branches to function scope properly, and
removes an extra block around function bodies in the new wat parser.

1 files changed, 35 insertions, 31 deletions

diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 585f9be49..4312fefe3 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -281,6 +281,15 @@ Result<> IRBuilder::visitArrayNew(ArrayNew* curr) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitFunctionStart(Function* func) {
+  if (!scopeStack.empty()) {
+    return Err{"unexpected start of function"};
+  }
+  scopeStack.push_back(ScopeCtx::makeFunc(func));
+  this->func = func;
+  return Ok{};
+}
+
 Result<> IRBuilder::visitBlockStart(Block* curr) {
   scopeStack.push_back(ScopeCtx::makeBlock(curr));
   return Ok{};
@@ -327,12 +336,12 @@ Result<Expression*> IRBuilder::finishScope(Block* block) {
       auto hoisted = hoistLastValue();
       CHECK_ERR(hoisted);
       auto hoistedType = scope.exprStack.back()->type;
-      if (hoistedType.size() != block->type.size()) {
+      if (hoistedType.size() != type.size()) {
         // We cannot propagate the hoisted value directly because it does not
         // have the correct number of elements. Break it up if necessary and
         // construct our returned tuple from parts.
         CHECK_ERR(packageHoistedValue(*hoisted));
-        std::vector<Expression*> elems(block->type.size());
+        std::vector<Expression*> elems(type.size());
         for (size_t i = 0; i < elems.size(); ++i) {
           auto elem = pop();
           CHECK_ERR(elem);
@@ -369,11 +378,11 @@ Result<Expression*> IRBuilder::finishScope(Block* block) {
   } else {
     // More than one expression, so we need a block. Allocate one if we weren't
     // already given one.
-    if (!block) {
-      block = wasm.allocator.alloc<Block>();
-      block->type = type;
+    if (block) {
+      block->list.set(scope.exprStack);
+    } else {
+      block = builder.makeBlock(scope.exprStack, type);
     }
-    block->list.set(scope.exprStack);
     ret = block;
   }
   scopeStack.pop_back();
@@ -395,50 +404,45 @@ Result<> IRBuilder::visitElse() {
 }
 
 Result<> IRBuilder::visitEnd() {
-  auto& scope = getScope();
+  auto scope = getScope();
   if (scope.isNone()) {
     return Err{"unexpected end"};
   }
-  if (auto* block = scope.getBlock()) {
-    auto expr = finishScope(block);
-    CHECK_ERR(expr);
+  auto expr = finishScope(scope.getBlock());
+  CHECK_ERR(expr);
+
+  // If the scope expression cannot be directly labeled, we may need to wrap it
+  // in a block.
+  auto maybeWrapForLabel = [&](Expression* curr) -> Expression* {
+    if (auto label = scope.getLabel()) {
+      return builder.makeBlock(label, {curr}, curr->type);
+    }
+    return curr;
+  };
+
+  if (auto* func = scope.getFunction()) {
+    func->body = *expr;
+  } else if (auto* block = scope.getBlock()) {
     assert(*expr == block);
     // TODO: Track branches so we can know whether this block is a target and
     // finalize more efficiently.
     block->finalize(block->type);
     push(block);
-    return Ok{};
   } else if (auto* loop = scope.getLoop()) {
-    auto expr = finishScope();
-    CHECK_ERR(expr);
     loop->body = *expr;
     loop->finalize(loop->type);
     push(loop);
-    return Ok{};
-  }
-  auto label = scope.getLabel();
-  Expression* scopeExpr = nullptr;
-  if (auto* iff = scope.getIf()) {
-    auto expr = finishScope();
-    CHECK_ERR(expr);
+  } else if (auto* iff = scope.getIf()) {
     iff->ifTrue = *expr;
     iff->ifFalse = nullptr;
     iff->finalize(iff->type);
-    scopeExpr = iff;
+    push(maybeWrapForLabel(iff));
   } else if (auto* iff = scope.getElse()) {
-    auto expr = finishScope();
-    CHECK_ERR(expr);
     iff->ifFalse = *expr;
     iff->finalize(iff->type);
-    scopeExpr = iff;
-  }
-  assert(scopeExpr && "unexpected scope kind");
-  if (label) {
-    // We cannot directly name an If in Binaryen IR, so we need to wrap it in
-    // a block.
-    push(builder.makeBlock(label, {scopeExpr}, scopeExpr->type));
+    push(maybeWrapForLabel(iff));
   } else {
-    push(scopeExpr);
+    WASM_UNREACHABLE("unexpected scope kind");
   }
   return Ok{};
 }