From ce1a2b480ae89e65b4d94e1bb3332c5980e2479f Mon Sep 17 00:00:00 2001 From: Thomas Lively Date: Fri, 15 Nov 2024 14:44:20 -0500 Subject: [PATCH] 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. --- src/parser/parsers.h | 2 +- src/passes/Outlining.cpp | 18 ++++++++++++------ src/wasm-ir-builder.h | 9 ++++++--- src/wasm/wasm-ir-builder.cpp | 1 + 4 files changed, 20 insertions(+), 10 deletions(-) diff --git a/src/parser/parsers.h b/src/parser/parsers.h index b6699aab6a9..4121788d75e 100644 --- a/src/parser/parsers.h +++ b/src/parser/parsers.h @@ -1143,7 +1143,7 @@ ifelse(Ctx& ctx, const std::vector& 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 68c3d039707..429511557e1 100644 --- a/src/passes/Outlining.cpp +++ b/src/passes/Outlining.cpp @@ -42,8 +42,8 @@ namespace wasm { struct ReconstructStringifyWalker : public StringifyWalker { - 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 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 40689a879e6..0210c55e50b 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> { 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 @@ class IRBuilder : public UnifiedExpressionVisitor> { // pushed instruction. void setDebugLocation(const std::optional&); + // 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 @@ class IRBuilder : public UnifiedExpressionVisitor> { 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 5e5decca42b..398db68ff4f 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;