Reset function context when ending a function in IRBuilder (#7081)

IRBuilder contains a pointer to the current function that is used to
create scratch locals, look up the operand types for returns, etc. This
pointer is nullable because IRBuilder can also be used in non-function
contexts such as global initializers. Visiting the start of a function
sets the function pointer, and after this change visiting the end of a
function resets the pointer to null. This avoids potential problems
where code outside a function would be able to incorrectly use scratch
locals and returns if the IRBuilder had previously been used to build a
function.

This change requires some adjustments to Outlining, which visits code
out of order, so ends up visiting code from inside a function after
visiting the end of the function. To support this use case, add a
`setFunction` method to IRBuilder that lets the user explicitly control
its function context. Also remove the optional function pointer
parameter to the IRBuilder constructor since it is less flexible and not
used.

4 files changed, 20 insertions, 10 deletions

diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index b6699aab6..4121788d7 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -1143,7 +1143,7 @@ ifelse(Ctx& ctx, const std::vector<Annotation>& annotations, bool folded) {
     ctx.setSrcLoc(annotations);
   }
 
-  ctx.makeIf(pos, annotations, label, *type);
+  CHECK_ERR(ctx.makeIf(pos, annotations, label, *type));
 
   if (folded && !ctx.in.takeSExprStart("then"sv)) {
     return ctx.in.err("expected 'then' before if instructions");
diff --git a/src/passes/Outlining.cpp b/src/passes/Outlining.cpp
index 68c3d0397..429511557 100644
--- a/src/passes/Outlining.cpp
+++ b/src/passes/Outlining.cpp
@@ -42,8 +42,8 @@ namespace wasm {
 struct ReconstructStringifyWalker
   : public StringifyWalker<ReconstructStringifyWalker> {
 
-  ReconstructStringifyWalker(Module* wasm)
-    : existingBuilder(*wasm), outlinedBuilder(*wasm) {
+  ReconstructStringifyWalker(Module* wasm, Function* func)
+    : existingBuilder(*wasm), outlinedBuilder(*wasm), func(func) {
     this->setModule(wasm);
     DBG(std::cerr << "\nexistingBuilder: " << &existingBuilder
                   << " outlinedBuilder: " << &outlinedBuilder << "\n");
@@ -77,6 +77,9 @@ struct ReconstructStringifyWalker
   // contain repeat sequences found in the program.
   IRBuilder outlinedBuilder;
 
+  // The function we are outlining from.
+  Function* func;
+
   void addUniqueSymbol(SeparatorReason reason) {
     if (auto curr = reason.getFuncStart()) {
       startExistingFunction(curr->func);
@@ -108,6 +111,8 @@ struct ReconstructStringifyWalker
       DBG(desc = "Loop Start at ");
     } else if (reason.getEnd()) {
       ASSERT_OK(existingBuilder.visitEnd());
+      // Reset the function in case we just ended the function scope.
+      existingBuilder.setFunction(func);
       // Outlining performs an unnested walk of the Wasm module, visiting
       // each scope one at a time. IRBuilder, in contrast, expects to
       // visit several nested scopes at a time. Thus, calling end() finalizes
@@ -346,15 +351,16 @@ struct Outlining : public Pass {
 
   void outline(Module* module, Sequences seqByFunc) {
     // TODO: Make this a function-parallel sub-pass.
-    ReconstructStringifyWalker reconstruct(module);
     std::vector<Name> keys(seqByFunc.size());
     std::transform(seqByFunc.begin(),
                    seqByFunc.end(),
                    keys.begin(),
                    [](auto pair) { return pair.first; });
-    for (auto func : keys) {
-      reconstruct.sequences = std::move(seqByFunc[func]);
-      reconstruct.doWalkFunction(module->getFunction(func));
+    for (auto funcName : keys) {
+      auto* func = module->getFunction(funcName);
+      ReconstructStringifyWalker reconstruct(module, func);
+      reconstruct.sequences = std::move(seqByFunc[funcName]);
+      reconstruct.doWalkFunction(func);
     }
   }
 
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 40689a879..0210c55e5 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -40,8 +40,7 @@ namespace wasm {
 // globals, tables, functions, etc.) to already exist in the module.
 class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
 public:
-  IRBuilder(Module& wasm, Function* func = nullptr)
-    : wasm(wasm), func(func), builder(wasm) {}
+  IRBuilder(Module& wasm) : wasm(wasm), builder(wasm) {}
 
   // Get the valid Binaryen IR expression representing the sequence of visited
   // instructions. The IRBuilder is reset and can be used with a fresh sequence
@@ -60,6 +59,10 @@ public:
   // pushed instruction.
   void setDebugLocation(const std::optional<Function::DebugLocation>&);
 
+  // Set the function used to add scratch locals when constructing an isolated
+  // sequence of IR.
+  void setFunction(Function* func) { this->func = func; }
+
   // 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.
@@ -234,7 +237,7 @@ public:
 
 private:
   Module& wasm;
-  Function* func;
+  Function* func = nullptr;
   Builder builder;
 
   // The location lacks debug info as it was marked as not having it.
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 5e5decca4..398db68ff 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -981,6 +981,7 @@ Result<> IRBuilder::visitEnd() {
     if (scope.needsPopFixup()) {
       EHUtils::handleBlockNestedPops(func, wasm);
     }
+    this->func = nullptr;
   } else if (auto* block = scope.getBlock()) {
     assert(*expr == block);
     block->name = scope.label;