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.

5 files changed, 57 insertions, 42 deletions

diff --git a/src/parser/context-defs.cpp b/src/parser/context-defs.cpp
index b43b7049f..76e61d4fb 100644
--- a/src/parser/context-defs.cpp
+++ b/src/parser/context-defs.cpp
@@ -57,9 +57,6 @@ Result<> ParseDefsCtx::addFunc(Name,
                                std::optional<LocalsT>,
                                Index pos) {
   CHECK_ERR(withLoc(pos, irBuilder.visitEnd()));
-  auto body = irBuilder.build();
-  CHECK_ERR(withLoc(pos, body));
-  wasm.functions[index]->body = *body;
   return Ok{};
 }
 
diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index f2aeb2489..7eedcccc2 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -816,9 +816,10 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
 
   IRBuilder irBuilder;
 
-  void setFunction(Function* func) {
+  Result<> visitFunctionStart(Function* func) {
     this->func = func;
-    irBuilder.setFunction(func);
+    CHECK_ERR(irBuilder.visitFunctionStart(func));
+    return Ok{};
   }
 
   ParseDefsCtx(std::string_view in,
diff --git a/src/parser/wat-parser.cpp b/src/parser/wat-parser.cpp
index 7b58be4d5..7e106d3b9 100644
--- a/src/parser/wat-parser.cpp
+++ b/src/parser/wat-parser.cpp
@@ -157,8 +157,7 @@ Result<> parseModule(Module& wasm, std::string_view input) {
 
     for (Index i = 0; i < decls.funcDefs.size(); ++i) {
       ctx.index = i;
-      ctx.setFunction(wasm.functions[i].get());
-      CHECK_ERR(ctx.irBuilder.makeBlock(Name{}, ctx.func->getResults()));
+      CHECK_ERR(ctx.visitFunctionStart(wasm.functions[i].get()));
       WithPosition with(ctx, decls.funcDefs[i].pos);
       auto parsed = func(ctx);
       CHECK_ERR(parsed);
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index a6afc63c4..2da69fff7 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -51,6 +51,7 @@ public:
   // Handle the boundaries of control flow structures. Users may choose to use
   // the corresponding `makeXYZ` function below instead of `visitXYZStart`, but
   // either way must call `visitEnd` and friends at the appropriate times.
+  [[nodiscard]] Result<> visitFunctionStart(Function* func);
   [[nodiscard]] Result<> visitBlockStart(Block* block);
   [[nodiscard]] Result<> visitIfStart(If* iff, Name label = {});
   [[nodiscard]] Result<> visitElse();
@@ -170,8 +171,6 @@ public:
   // [[nodiscard]] Result<> makeStringSliceWTF();
   // [[nodiscard]] Result<> makeStringSliceIter();
 
-  void setFunction(Function* func) { this->func = func; }
-
   // Private functions that must be public for technical reasons.
   [[nodiscard]] Result<> visitExpression(Expression*);
   [[nodiscard]] Result<> visitBlock(Block*);
@@ -189,6 +188,9 @@ private:
   // to have.
   struct ScopeCtx {
     struct NoScope {};
+    struct FuncScope {
+      Function* func;
+    };
     struct BlockScope {
       Block* block;
     };
@@ -203,8 +205,8 @@ private:
     struct LoopScope {
       Loop* loop;
     };
-    using Scope =
-      std::variant<NoScope, BlockScope, IfScope, ElseScope, LoopScope>;
+    using Scope = std::
+      variant<NoScope, FuncScope, BlockScope, IfScope, ElseScope, LoopScope>;
 
     // The control flow structure we are building expressions for.
     Scope scope;
@@ -217,6 +219,9 @@ private:
     ScopeCtx() : scope(NoScope{}) {}
     ScopeCtx(Scope scope) : scope(scope) {}
 
+    static ScopeCtx makeFunc(Function* func) {
+      return ScopeCtx(FuncScope{func});
+    }
     static ScopeCtx makeBlock(Block* block) {
       return ScopeCtx(BlockScope{block});
     }
@@ -229,6 +234,12 @@ private:
     static ScopeCtx makeLoop(Loop* loop) { return ScopeCtx(LoopScope{loop}); }
 
     bool isNone() { return std::get_if<NoScope>(&scope); }
+    Function* getFunction() {
+      if (auto* funcScope = std::get_if<FuncScope>(&scope)) {
+        return funcScope->func;
+      }
+      return nullptr;
+    }
     Block* getBlock() {
       if (auto* blockScope = std::get_if<BlockScope>(&scope)) {
         return blockScope->block;
@@ -254,6 +265,9 @@ private:
       return nullptr;
     }
     Type getResultType() {
+      if (auto* func = getFunction()) {
+        return func->type.getSignature().results;
+      }
       if (auto* block = getBlock()) {
         return block->type;
       }
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{};
 }