Fix visitBlock and add visitBlockStart in IRBuilder (#5959)

Visiting a block should push it onto the stack just like visiting any other
expression, but we previously had a `visitBlock` that introduced a new scope
instead. Fix `visitBlock` to behave as expected and introduce a new
`visitBlockStart` method to introduce a new scope.

Unfortunately this cannot be fully tested yet because the wat parser uses the
`makeXYZ` API intead of the `visit` API, but at least I updated `makeBlock` to
call `visitBlockStart`, so that is tested.

2 files changed, 19 insertions, 9 deletions

diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 50c034897..1a19f62e8 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -45,15 +45,13 @@ public:
   [[nodiscard]] Result<Expression*> build();
 
   // Call visit() on an existing Expression with its non-child fields
-  // initialized to initialize the child fields and refinalize it. The specific
-  // visitors are internal implementation details.
+  // initialized to initialize the child fields and refinalize it.
   [[nodiscard]] Result<> visit(Expression*);
-  [[nodiscard]] Result<> visitExpression(Expression*);
-  [[nodiscard]] Result<> visitBlock(Block*);
-  [[nodiscard]] Result<> visitReturn(Return*);
-  [[nodiscard]] Result<> visitStructNew(StructNew*);
-  [[nodiscard]] Result<> visitArrayNew(ArrayNew*);
 
+  // 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<> visitBlockStart(Block* block);
   [[nodiscard]] Result<> visitEnd();
 
   // Alternatively, call makeXYZ to have the IRBuilder allocate the nodes. This
@@ -171,6 +169,13 @@ public:
 
   void setFunction(Function* func) { this->func = func; }
 
+  // Private functions that must be public for technical reasons.
+  [[nodiscard]] Result<> visitExpression(Expression*);
+  [[nodiscard]] Result<> visitBlock(Block*);
+  [[nodiscard]] Result<> visitReturn(Return*);
+  [[nodiscard]] Result<> visitStructNew(StructNew*);
+  [[nodiscard]] Result<> visitArrayNew(ArrayNew*);
+
 private:
   Module& wasm;
   Function* func;
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 6445cc3c9..bc3180423 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -233,7 +233,7 @@ Result<> IRBuilder::visitExpression(Expression* curr) {
 }
 
 Result<> IRBuilder::visitBlock(Block* curr) {
-  scopeStack.push_back({{}, curr});
+  // No children; pushing and finalizing will be handled by `visit`.
   return Ok{};
 }
 
@@ -281,6 +281,11 @@ Result<> IRBuilder::visitArrayNew(ArrayNew* curr) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitBlockStart(Block* curr) {
+  scopeStack.push_back({{}, curr});
+  return Ok{};
+}
+
 Result<> IRBuilder::visitEnd() {
   if (scopeStack.empty() || !scopeStack.back().block) {
     return Err{"unexpected end"};
@@ -347,7 +352,7 @@ Result<> IRBuilder::makeBlock(Name label, Type type) {
   auto* block = wasm.allocator.alloc<Block>();
   block->name = label;
   block->type = type;
-  scopeStack.push_back({{}, block});
+  CHECK_ERR(visitBlockStart(block));
   return Ok{};
 }