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.

Author: tlively

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{};
 }
 

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,

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);

src/wasm-ir-builder.h

@@ -51,6 +51,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // 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 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // [[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 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // to have.
   struct ScopeCtx {
     struct NoScope {};
+    struct FuncScope {
+      Function* func;
+    };
     struct BlockScope {
       Block* block;
     };
@@ -203,8 +205,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     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 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     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 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     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 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
       return nullptr;
     }
     Type getResultType() {
+      if (auto* func = getFunction()) {
+        return func->type.getSignature().results;
+      }
       if (auto* block = getBlock()) {
         return block->type;
       }

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{};
 }