From 8c0f469b2c43ebd4151b1bc4fdb1c3f451825b95 Mon Sep 17 00:00:00 2001 From: Thomas Lively Date: Wed, 21 Oct 2026 23:15:47 -0700 Subject: [PATCH 1/2] [StackIR] Support control flow input parameters Add a params field to StackInst on control-flow begin instructions (BlockBegin, IfBegin, LoopBegin, TryBegin, and TryTableBegin), and update ModuleUtils::collectHeapTypeInfo, BinaryInstWriter, WasmBinaryWriter, PrintSExpression, and StackIROptimizer to register, emit, print, and account for control-flow input parameter signatures. The optimizer does not yet use the input parameters, so add gtests that artificially test printing and binary writing for now. These tests will be replaced with normal lit tests in future commits that start using the block parameters in the optimizer. Using control flow input parameters in Stack IR will ultimately improve code size by up to 1.2% geomean across many interesting Wasm modules, depending on what other optimizations run first. --- src/ir/module-utils.cpp | 67 ++++++++--- src/ir/module-utils.h | 11 +- src/passes/Print.cpp | 102 +++++++++------- src/wasm-binary.h | 9 +- src/wasm-stack.h | 29 +++-- src/wasm/wasm-binary.cpp | 20 +++- src/wasm/wasm-stack-opts.cpp | 9 +- src/wasm/wasm-stack.cpp | 9 +- test/gtest/printing.cpp | 219 ++++++++++++++++++++++++++++++++++- 9 files changed, 390 insertions(+), 85 deletions(-) diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp index af8f3e9ff66..22dd8c0d65f 100644 --- a/src/ir/module-utils.cpp +++ b/src/ir/module-utils.cpp @@ -24,6 +24,7 @@ #include "support/topological_sort.h" #include "support/utilities.h" #include "wasm-builder.h" +#include "wasm-stack.h" namespace wasm::ModuleUtils { @@ -398,9 +399,10 @@ struct TypeInfos { } } void noteControlFlow(Signature sig) { - // TODO: support control flow input parameters. - assert(sig.params.size() == 0); - if (sig.results.isTuple()) { + if (sig.results == Type::unreachable) { + sig.results = Type::none; + } + if (sig.params != Type::none || sig.results.isTuple()) { // We have to use a function type. ++controlFlowSignatures[sig]; } else if (sig.results != Type::none) { @@ -417,9 +419,13 @@ using ReferencedFuncs = std::unordered_map>; struct CodeScanner : PostWalker { TypeInfos& info; ReferencedFuncs& referencedFuncs; + bool hasStackIR; - CodeScanner(Module& wasm, TypeInfos& info, ReferencedFuncs& referencedFuncs) - : info(info), referencedFuncs(referencedFuncs) { + CodeScanner(Module& wasm, + TypeInfos& info, + ReferencedFuncs& referencedFuncs, + bool hasStackIR = false) + : info(info), referencedFuncs(referencedFuncs), hasStackIR(hasStackIR) { setModule(&wasm); } @@ -475,19 +481,29 @@ struct CodeScanner : PostWalker { info.note(curr->type); } void visitBlock(Block* curr) { - info.noteControlFlow(Signature(Type::none, curr->type)); + if (!hasStackIR) { + info.noteControlFlow(Signature(Type::none, curr->type)); + } } void visitIf(If* curr) { - info.noteControlFlow(Signature(Type::none, curr->type)); + if (!hasStackIR) { + info.noteControlFlow(Signature(Type::none, curr->type)); + } } void visitLoop(Loop* curr) { - info.noteControlFlow(Signature(Type::none, curr->type)); + if (!hasStackIR) { + info.noteControlFlow(Signature(Type::none, curr->type)); + } } void visitTry(Try* curr) { - info.noteControlFlow(Signature(Type::none, curr->type)); + if (!hasStackIR) { + info.noteControlFlow(Signature(Type::none, curr->type)); + } } void visitTryTable(TryTable* curr) { - info.noteControlFlow(Signature(Type::none, curr->type)); + if (!hasStackIR) { + info.noteControlFlow(Signature(Type::none, curr->type)); + } } void visitRefFunc(RefFunc* curr) { referencedFuncs.at(curr->func) = true; } }; @@ -503,7 +519,8 @@ InsertOrderedMap collectHeapTypeInfo(Module& wasm, WorldMode worldMode, TypeInclusion inclusion, - VisibilityHandling visibility) { + VisibilityHandling visibility, + const ModuleStackIR* stackIR) { // Collect module-level info. TypeInfos info; ReferencedFuncs referencedFuncs; @@ -531,11 +548,27 @@ collectHeapTypeInfo(Module& wasm, for (auto type : func->vars) { info.note(type); } + auto* funcStackIR = stackIR ? stackIR->getStackIROrNull(func) : nullptr; + if (funcStackIR) { + for (auto* inst : *funcStackIR) { + if (!inst) { + continue; + } + if (inst->op == StackInst::BlockBegin || + inst->op == StackInst::IfBegin || + inst->op == StackInst::LoopBegin || + inst->op == StackInst::TryBegin || + inst->op == StackInst::TryTableBegin) { + info.noteControlFlow(Signature(inst->params, inst->origin->type)); + } + } + } // Don't just use `func->imported()` here because we also might be // printing an error message on a partially parsed module whose declared // function bodies have not all been parsed yet. if (func->body) { - CodeScanner(wasm, info, referencedFuncs).walk(func->body); + CodeScanner(wasm, info, referencedFuncs, funcStackIR != nullptr) + .walk(func->body); } }); @@ -975,9 +1008,13 @@ std::vector getPrivateHeapTypes(Module& wasm, WorldMode worldMode) { return types; } -IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm) { - auto counts = - collectHeapTypeInfo(wasm, WorldMode::Open, TypeInclusion::BinaryTypes); +IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm, + const ModuleStackIR* stackIR) { + auto counts = collectHeapTypeInfo(wasm, + WorldMode::Open, + TypeInclusion::BinaryTypes, + VisibilityHandling::NoVisibility, + stackIR); // Collect the rec groups. std::unordered_map groupIndices; diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h index 70297c7cee2..03bcd5275b4 100644 --- a/src/ir/module-utils.h +++ b/src/ir/module-utils.h @@ -22,6 +22,10 @@ #include "support/unique_deferring_queue.h" #include "wasm.h" +namespace wasm { +class ModuleStackIR; +} // namespace wasm + namespace wasm::ModuleUtils { // Copies a function into a module. If newName is provided it is used as the @@ -475,7 +479,8 @@ InsertOrderedMap collectHeapTypeInfo( Module& wasm, WorldMode worldMode, TypeInclusion inclusion = TypeInclusion::AllTypes, - VisibilityHandling visibility = VisibilityHandling::NoVisibility); + VisibilityHandling visibility = VisibilityHandling::NoVisibility, + const ModuleStackIR* stackIR = nullptr); // Helper function for collecting all the non-basic heap types used in the // module, i.e. the types that would appear in the type section. @@ -503,7 +508,9 @@ struct IndexedHeapTypes { // Similar to `collectHeapTypes`, but provides fast lookup of the index for each // type as well. Also orders the types to be valid and sorts the types by // frequency of use to minimize code size. -IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm); +IndexedHeapTypes +getOptimizedIndexedHeapTypes(Module& wasm, + const ModuleStackIR* stackIR = nullptr); } // namespace wasm::ModuleUtils diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp index 2140e46adb0..d6b2b58137c 100644 --- a/src/passes/Print.cpp +++ b/src/passes/Print.cpp @@ -252,18 +252,21 @@ struct PrintSExpression : public UnifiedExpressionVisitor { } std::ostream& printBlockType(Signature sig) { - assert(sig.params == Type::none); - if (sig.results == Type::none) { - return o; - } - if (sig.results.isTuple()) { + if (sig.params != Type::none || sig.results.isTuple()) { if (auto it = signatureTypes.find(sig); it != signatureTypes.end()) { - o << "(type "; + o << " (type "; printHeapTypeName(it->second); - o << ") "; + o << ')'; } } - printResultType(sig.results); + if (sig.params != Type::none) { + o << ' '; + printParamType(sig.params); + } + if (sig.results != Type::none) { + o << ' '; + printResultType(sig.results); + } return o; } @@ -362,6 +365,33 @@ struct PrintExpressionContents currFunction(parent.currFunction), o(parent.o), features(wasm ? wasm->features : FeatureSet::All), full(isFullForced()) {} + using OverriddenVisitor::visit; + + void visit(StackInst* inst) { + switch (inst->op) { + case StackInst::Basic: + visit(inst->origin); + break; + case StackInst::BlockBegin: + visitBlock(inst->origin->cast(), inst->params); + break; + case StackInst::IfBegin: + visitIf(inst->origin->cast(), inst->params); + break; + case StackInst::LoopBegin: + visitLoop(inst->origin->cast(), inst->params); + break; + case StackInst::TryBegin: + visitTry(inst->origin->cast(), inst->params); + break; + case StackInst::TryTableBegin: + visitTryTable(inst->origin->cast(), inst->params); + break; + default: + WASM_UNREACHABLE("unexpected op"); + } + } + std::ostream& printType(Type type) { return parent.printType(type); } std::ostream& printHeapTypeName(HeapType type) { @@ -376,8 +406,11 @@ struct PrintExpressionContents return parent.printParamType(type); } - std::ostream& printBlockType(Signature sig) { - return parent.printBlockType(sig); + void printBlockType(Type results, Type params = Type::none) { + if (!results.isConcrete()) { + results = Type::none; + } + parent.printBlockType(Signature(params, results)); } void printMemoryPostfix(uint8_t bytes, Type type) { @@ -407,18 +440,15 @@ struct PrintExpressionContents return o; } - void visitBlock(Block* curr) { + void visitBlock(Block* curr, Type params = Type::none) { printMedium(o, "block"); if (curr->name.is()) { o << ' '; curr->name.print(o); } - if (curr->type.isConcrete()) { - o << ' '; - printBlockType(Signature(Type::none, curr->type)); - } + printBlockType(curr->type, params); } - void visitIf(If* curr) { + void visitIf(If* curr, Type params = Type::none) { printMedium(o, "if"); // Ifs are unreachable if their condition is unreachable, but in that case // the arms might have some concrete type we have to account for to produce @@ -427,21 +457,15 @@ struct PrintExpressionContents if (curr->condition->type == Type::unreachable && curr->ifFalse) { type = Type::getLeastUpperBound(curr->ifTrue->type, curr->ifFalse->type); } - if (type.isConcrete()) { - o << ' '; - printBlockType(Signature(Type::none, type)); - } + printBlockType(type, params); } - void visitLoop(Loop* curr) { + void visitLoop(Loop* curr, Type params = Type::none) { printMedium(o, "loop"); if (curr->name.is()) { o << ' '; curr->name.print(o); } - if (curr->type.isConcrete()) { - o << ' '; - printBlockType(Signature(Type::none, curr->type)); - } + printBlockType(curr->type, params); } void visitBreak(Break* curr) { if (curr->condition) { @@ -2136,23 +2160,17 @@ struct PrintExpressionContents printMedium(o, "elem.drop "); curr->segment.print(o); } - void visitTry(Try* curr) { + void visitTry(Try* curr, Type params = Type::none) { printMedium(o, "try"); if (curr->name.is()) { o << ' '; curr->name.print(o); } - if (curr->type.isConcrete()) { - o << ' '; - printBlockType(Signature(Type::none, curr->type)); - } + printBlockType(curr->type, params); } - void visitTryTable(TryTable* curr) { + void visitTryTable(TryTable* curr, Type params = Type::none) { printMedium(o, "try_table"); - if (curr->type.isConcrete()) { - o << ' '; - printBlockType(Signature(Type::none, curr->type)); - } + printBlockType(curr->type, params); for (Index i = 0; i < curr->catchTags.size(); i++) { o << " ("; if (curr->catchTags[i]) { @@ -2713,7 +2731,9 @@ struct PrintExpressionContents void PrintSExpression::setModule(Module* module) { currModule = module; if (module) { - heapTypes = ModuleUtils::getOptimizedIndexedHeapTypes(*module).types; + heapTypes = ModuleUtils::getOptimizedIndexedHeapTypes( + *module, moduleStackIR ? &*moduleStackIR : nullptr) + .types; for (auto type : heapTypes) { if (type.isSignature()) { signatureTypes.insert({type.getSignature(), type}); @@ -3842,14 +3862,16 @@ printStackInst(StackInst* inst, std::ostream& o, Function* func) { case StackInst::BlockBegin: case StackInst::IfBegin: case StackInst::LoopBegin: - case StackInst::TryBegin: { - PrintExpressionContents(printer).visit(inst->origin); + case StackInst::TryBegin: + case StackInst::TryTableBegin: { + PrintExpressionContents(printer).visit(inst); break; } case StackInst::BlockEnd: case StackInst::IfEnd: case StackInst::LoopEnd: - case StackInst::TryEnd: { + case StackInst::TryEnd: + case StackInst::TryTableEnd: { printMedium(o, "end"); o << " ;; type: "; printer.printType(inst->type); @@ -3914,7 +3936,7 @@ static std::ostream& printStackIR(StackIR* ir, PrintSExpression& printer) { case StackInst::TryTableBegin: { controlFlowDepth++; doIndent(); - PrintExpressionContents(printer).visit(inst->origin); + PrintExpressionContents(printer).visit(inst); indent++; break; } diff --git a/src/wasm-binary.h b/src/wasm-binary.h index 470f58be2e7..02ae5e612fa 100644 --- a/src/wasm-binary.h +++ b/src/wasm-binary.h @@ -1309,6 +1309,8 @@ enum FeaturePrefix { FeatureUsed = '+', FeatureDisallowed = '-' }; // (local index in IR, tuple index) => binary local index using MappedLocals = std::unordered_map, size_t>; +class ModuleStackIR; + // Writes out wasm to the binary format class WasmBinaryWriter { @@ -1383,10 +1385,8 @@ class WasmBinaryWriter { public: WasmBinaryWriter(Module* input, BufferWithRandomAccess& o, - const PassOptions& options) - : wasm(input), o(o), options(options), indexes(*input) { - prepare(); - } + const PassOptions& options); + ~WasmBinaryWriter(); // locations in the output binary for the various parts of the module struct TableOfContents { @@ -1552,6 +1552,7 @@ class WasmBinaryWriter { Function::DebugLocation lastDebugLocation; std::unique_ptr importInfo; + std::unique_ptr moduleStackIR; // General debugging info: track locations as we write. BinaryLocations binaryLocations; diff --git a/src/wasm-stack.h b/src/wasm-stack.h index 7d772e0d85f..0e74360f2db 100644 --- a/src/wasm-stack.h +++ b/src/wasm-stack.h @@ -23,6 +23,7 @@ #include "ir/properties.h" #include "pass.h" #include "support/insert_ordered.h" +#include "support/utilities.h" #include "wasm-binary.h" #include "wasm-traversal.h" #include "wasm.h" @@ -85,6 +86,11 @@ class StackInst { // the type - usually identical to the origin type, but e.g. wasm has no // unreachable blocks, they must be none Type type; + + // For control-flow begin instructions (BlockBegin, IfBegin, LoopBegin, + // TryBegin, TryTableBegin), the input parameter types of the control-flow + // block signature. + Type params = Type::none; }; using StackIR = std::vector; @@ -99,16 +105,20 @@ class BinaryInstWriter : public OverriddenVisitor { numLocalsByType(parent.getModule()->features), scratchLocals(parent.getModule()->features) {} - void visit(Expression* curr) { + void visit(Expression* curr, Type params = Type::none) { if (func) { parent.trackExpressionStart(curr, func); } + currParams = params; OverriddenVisitor::visit(curr); + currParams = Type::none; if (func) { parent.trackExpressionEnd(curr, func); } } + void visit(StackInst* inst) { visit(inst->origin, inst->params); } + #define DELEGATE(CLASS_TO_VISIT) \ void visit##CLASS_TO_VISIT(CLASS_TO_VISIT* curr); @@ -156,6 +166,7 @@ class BinaryInstWriter : public OverriddenVisitor { BufferWithRandomAccess& o; Function* func = nullptr; bool DWARF; + Type currParams = Type::none; std::vector breakStack; @@ -598,11 +609,11 @@ class ModuleStackIR { // Get StackIR for a function, if it exists. (This allows some functions to // have it and others not, if we add such capability in the future.) StackIR* getStackIROrNull(Function* func) { - auto iter = analysis.map.find(func); - if (iter == analysis.map.end()) { - return nullptr; - } - return &iter->second; + return find_or_null(analysis.map, func); + } + + const StackIR* getStackIROrNull(Function* func) const { + return find_or_null(analysis.map, func); } }; @@ -668,10 +679,8 @@ class StackIROptimizer { std::ostream& printStackIR(std::ostream& o, Module* module, const PassOptions& options); -} // namespace wasm - -namespace std { std::ostream& operator<<(std::ostream& o, wasm::StackInst& inst); -} // namespace std + +} // namespace wasm #endif // wasm_wasm_stack_h diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp index aede8e67b0d..42dab5f9a21 100644 --- a/src/wasm/wasm-binary.cpp +++ b/src/wasm/wasm-binary.cpp @@ -38,10 +38,23 @@ namespace wasm { +WasmBinaryWriter::WasmBinaryWriter(Module* input, + BufferWithRandomAccess& o, + const PassOptions& options) + : wasm(input), o(o), options(options), indexes(*input) { + prepare(); +} + +WasmBinaryWriter::~WasmBinaryWriter() = default; + void WasmBinaryWriter::prepare() { + if (options.generateStackIR) { + moduleStackIR = std::make_unique(*wasm, options); + } // Collect function types and their frequencies. Collect information in each // function in parallel, then merge. - indexedTypes = ModuleUtils::getOptimizedIndexedHeapTypes(*wasm); + indexedTypes = + ModuleUtils::getOptimizedIndexedHeapTypes(*wasm, moduleStackIR.get()); for (Index i = 0, size = indexedTypes.types.size(); i < size; ++i) { if (indexedTypes.types[i].isSignature()) { signatureIndexes.insert({indexedTypes.types[i].getSignature(), i}); @@ -524,11 +537,6 @@ void WasmBinaryWriter::writeFunctions() { return; } - std::optional moduleStackIR; - if (options.generateStackIR) { - moduleStackIR.emplace(*wasm, options); - } - auto sectionStart = startSection(BinaryConsts::Section::Code); o << U32LEB(importInfo->getNumDefinedFunctions()); bool DWARF = Debug::hasDWARFSections(*getModule()); diff --git a/src/wasm/wasm-stack-opts.cpp b/src/wasm/wasm-stack-opts.cpp index e3ec8f4ac18..76d3d4ce1a3 100644 --- a/src/wasm/wasm-stack-opts.cpp +++ b/src/wasm/wasm-stack-opts.cpp @@ -216,7 +216,7 @@ void StackIROptimizer::local2Stack() { if (isControlFlowBegin(inst)) { // Save the stack for when we end this control flow. savedValues.push_back(values); // TODO: optimize copies - values.clear(); + values.assign(inst->params.size(), null); } else if (isControlFlowEnd(inst)) { assert(!savedValues.empty()); values = savedValues.back(); @@ -491,9 +491,12 @@ void StackIROptimizer::removeAt(Index i) { Index StackIROptimizer::getNumConsumedValues(StackInst* inst) { if (isControlFlow(inst)) { - // If consumes 1; that's it. + // If consumes 1 (condition) plus any input parameters. if (inst->op == StackInst::IfBegin) { - return 1; + return 1 + inst->params.size(); + } + if (isControlFlowBegin(inst)) { + return inst->params.size(); } return 0; } diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp index 9009a425c7d..0f017467320 100644 --- a/src/wasm/wasm-stack.cpp +++ b/src/wasm/wasm-stack.cpp @@ -24,9 +24,10 @@ static Name IMPOSSIBLE_CONTINUE("impossible-continue"); void BinaryInstWriter::emitResultType(Type type) { if (type == Type::unreachable) { - parent.writeType(Type::none); - } else if (type.isTuple()) { - o << S32LEB(parent.getSignatureIndex(Signature(Type::none, type))); + type = Type::none; + } + if (currParams != Type::none || type.isTuple()) { + o << S32LEB(parent.getSignatureIndex(Signature(currParams, type))); } else { parent.writeType(type); } @@ -3750,7 +3751,7 @@ void StackIRToBinaryWriter::write() { if (sourceMap) { parent.writeSourceMapLocation(inst->origin, func); } - writer.visit(inst->origin); + writer.visit(inst); break; } case StackInst::TryEnd: diff --git a/test/gtest/printing.cpp b/test/gtest/printing.cpp index 31259788618..083d87edcc7 100644 --- a/test/gtest/printing.cpp +++ b/test/gtest/printing.cpp @@ -1,5 +1,8 @@ +#include "ir/module-utils.h" +#include "pass.h" #include "print-test.h" - +#include "wasm-binary.h" +#include "wasm-stack.h" #include "wasm.h" using namespace wasm; @@ -42,3 +45,217 @@ TEST_F(PrintingTest, Print) { )print"); } } + +TEST_F(PrintingTest, StackIRParams) { + auto moduleText = R"wasm( + (module + (func $f (param $p i32) (result i32) + (local $x i32) + (block $b (result i32) + (loop $l + (try $t + (do + (try_table (catch_all $b) + (br_if $l (local.get $p)) + ) + ) + (catch_all + (nop) + ) + ) + ) + (if (result i32) + (local.get $p) + (then + (i32.const 1) + ) + (else + (i32.const 2) + ) + ) + ) + ) + ) + )wasm"; + + Module wasm; + wasm.features = FeatureSet::All; + parseWast(wasm, moduleText); + + PassOptions options; + ModuleStackIR moduleStackIR(wasm, options); + auto* func = wasm.getFunction("f"); + auto* stackIR = moduleStackIR.getStackIROrNull(func); + ASSERT_NE(stackIR, nullptr); + + TypeBuilder builder(1); + builder[0] = Struct({Field(Type::i32, Mutable)}); + auto built = builder.build(); + ASSERT_TRUE(built); + HeapType structType = (*built)[0]; + Type structRef(structType, Nullable); + + for (auto* inst : *stackIR) { + if (!inst) { + continue; + } + switch (inst->op) { + case StackInst::BlockBegin: { + inst->params = structRef; + std::stringstream ss; + ss << *inst; + EXPECT_EQ(ss.str(), + "block $b (param (ref null $struct.0)) (result i32)"); + break; + } + case StackInst::LoopBegin: { + inst->params = Type::i64; + std::stringstream ss; + ss << *inst; + EXPECT_EQ(ss.str(), "loop $l (param i64)"); + break; + } + case StackInst::TryBegin: { + inst->params = Type::f32; + std::stringstream ss; + ss << *inst; + EXPECT_EQ(ss.str(), "try $t (param f32)"); + break; + } + case StackInst::TryTableBegin: { + inst->params = Type::i32; + std::stringstream ss; + ss << *inst; + EXPECT_EQ(ss.str(), "try_table (param i32) (catch_all $b)"); + break; + } + case StackInst::IfBegin: { + inst->params = Type({Type::i32, Type::f64}); + std::stringstream ss; + ss << *inst; + EXPECT_EQ(ss.str(), "if (param i32 f64) (result i32)"); + break; + } + default: + break; + } + } + + // Verify that getOptimizedIndexedHeapTypes collects control flow signatures + // from StackIR as well as child heap types only referenced in + // StackInst::params. + auto indexedTypes = + ModuleUtils::getOptimizedIndexedHeapTypes(wasm, &moduleStackIR); + EXPECT_TRUE(indexedTypes.indices.contains(Signature(structRef, Type::i32))); + EXPECT_TRUE(indexedTypes.indices.contains(structType)); + EXPECT_TRUE(indexedTypes.indices.contains(Signature(Type::i64, Type::none))); + EXPECT_TRUE(indexedTypes.indices.contains(Signature(Type::f32, Type::none))); + EXPECT_TRUE(indexedTypes.indices.contains(Signature(Type::i32, Type::none))); + EXPECT_TRUE(indexedTypes.indices.contains( + Signature(Type({Type::i32, Type::f64}), Type::i32))); +} + +TEST_F(PrintingTest, StackIRParamBinaryAndOptimizer) { + auto moduleText = R"wasm( + (module + (func $id (param $x i32) (result i32) + (local $tmp i32) + (local.set $tmp (local.get $x)) + (block $b (result i32) + (i32.add + (local.get $tmp) + (i32.const 1) + ) + ) + ) + ) + )wasm"; + + Module wasm; + wasm.features = FeatureSet::All; + parseWast(wasm, moduleText); + + PassOptions options; + options.optimizeLevel = 3; + ModuleStackIR moduleStackIR(wasm, options); + auto* func = wasm.getFunction("id"); + auto* stackIR = moduleStackIR.getStackIROrNull(func); + ASSERT_NE(stackIR, nullptr); + + // Build a StackIR sequence where $b takes an i32 param from the stack and + // adds it to a local.get inside the block: + // local.get $x + // local.set $tmp + // i32.const 10 ;; consumed as block param + // block $b (param i32) (result i32) + // local.get $tmp + // i32.add ;; consumes block param + $tmp + // i32.const 0 + // br_if $b + // end + Builder irBuilder(wasm); + auto* getX = irBuilder.makeLocalGet(0, Type::i32); + auto* setTmp = irBuilder.makeLocalSet(1, getX); + auto* c10 = irBuilder.makeConst(Literal(int32_t(10))); + auto* getTmp = irBuilder.makeLocalGet(1, Type::i32); + auto* add = irBuilder.makeBinary(AddInt32, c10, getTmp); + auto* c0 = irBuilder.makeConst(Literal(int32_t(0))); + auto* brIf = irBuilder.makeBreak("b", add, c0); + auto* block = irBuilder.makeBlock("b", {brIf}, Type::i32); + func->body = irBuilder.makeSequence(setTmp, block); + + auto makeInst = [&](StackInst::Op op, Expression* origin) { + auto* inst = wasm.allocator.alloc(); + inst->op = op; + inst->origin = origin; + inst->type = origin->type; + return inst; + }; + + auto* blockBegin = makeInst(StackInst::BlockBegin, block); + blockBegin->params = Type::i32; + auto* blockEnd = makeInst(StackInst::BlockEnd, block); + + *stackIR = { + makeInst(StackInst::Basic, getX), + makeInst(StackInst::Basic, setTmp), + makeInst(StackInst::Basic, c10), + blockBegin, + makeInst(StackInst::Basic, getTmp), + makeInst(StackInst::Basic, add), + makeInst(StackInst::Basic, c0), + makeInst(StackInst::Basic, brIf), + blockEnd, + }; + + // Run StackIROptimizer: local2Stack must see that blockBegin consumes c10 and + // pushes 1 parameter inside the block, so setTmp/getTmp cannot be eliminated + // across the block boundary, and i32.add consumes both the block param and + // getTmp without stack underflow. + StackIROptimizer optimizer(func, *stackIR, options, wasm.features); + optimizer.run(); + EXPECT_NE((*stackIR)[1], nullptr); + EXPECT_NE((*stackIR)[4], nullptr); + + // Also test binary emission via StackIRToBinaryWriter: the block header must + // emit the signature index for (param i32) (result i32). + BufferWithRandomAccess moduleBuf; + WasmBinaryWriter binaryWriter(&wasm, moduleBuf, options); + BufferWithRandomAccess funcBuf; + StackIRToBinaryWriter stackWriter(binaryWriter, funcBuf, func, *stackIR); + stackWriter.write(); + + // Verify the emitted bytes contain BinaryConsts::Block followed by the + // signature index of (param i32) (result i32). + auto sigIndex = + binaryWriter.getSignatureIndex(Signature(Type::i32, Type::i32)); + bool foundParameterizedBlock = false; + for (size_t i = 0; i + 1 < funcBuf.size(); ++i) { + if (funcBuf[i] == uint8_t(BinaryConsts::Block) && + funcBuf[i + 1] == uint8_t(sigIndex)) { + foundParameterizedBlock = true; + break; + } + } + EXPECT_TRUE(foundParameterizedBlock); +} From 08006343ef568a6412e4f7fbc7bb438235ef7d73 Mon Sep 17 00:00:00 2001 From: Thomas Lively Date: Mon, 26 Oct 2026 15:48:29 -0700 Subject: [PATCH 2/2] update tests --- test/lit/passes/constraint-analysis.wast | 2 +- test/lit/passes/translate-to-exnref.wast | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/lit/passes/constraint-analysis.wast b/test/lit/passes/constraint-analysis.wast index b5bba9857d5..d444460146d 100644 --- a/test/lit/passes/constraint-analysis.wast +++ b/test/lit/passes/constraint-analysis.wast @@ -7267,7 +7267,7 @@ (local.get $x) ) ) - + ;; CHECK: (func $asymmetric-contradiction (type $2) (param $0 i32) (param $1 i32) ;; CHECK-NEXT: (if ;; CHECK-NEXT: (i32.and diff --git a/test/lit/passes/translate-to-exnref.wast b/test/lit/passes/translate-to-exnref.wast index b4168505256..60497857a9c 100644 --- a/test/lit/passes/translate-to-exnref.wast +++ b/test/lit/passes/translate-to-exnref.wast @@ -19,10 +19,10 @@ ;; CHECK: (type $6 (func (param i32 i64))) ;; CHECK: (tag $e-empty (type $1)) - ;; STACKIR-OPT: (type $0 (func (result i32 i64))) - ;; STACKIR-OPT: (type $1 (func)) + ;; STACKIR-OPT: (type $0 (func (result i32 i64))) + ;; STACKIR-OPT: (type $2 (func (result i32))) ;; STACKIR-OPT: (type $3 (func (result i32 (ref exn))))