From da00140fbbc3fa5da5c3d19603263b6a4161c201 Mon Sep 17 00:00:00 2001 From: Thomas Lively Date: Tue, 7 Apr 2026 09:51:17 -0700 Subject: [PATCH] Fix optimized grouping of locals In #8568 we optimized the grouping of locals in the binary writer to account for how types will be written given the enabled features. However, that change did not properly update the handling of scratch locals accordingly, leading to inconsistencies in the indices assigned to local types in different locations. Fix the problem by reverting the changes from #8568 and handling the mapping from IR types to written types at a lower level; specifically, create a new `TypeIndexMap` type that extends `InsertOrderedMap` but always applies `asWrittenGivenFeatures` to its keys. Use this new map type both for the `numLocalsByType` map and the `scratchLocals` map. --- src/wasm-stack.h | 37 ++++++++++++++++++++++++++++++++----- src/wasm/wasm-stack.cpp | 30 +++++++----------------------- 2 files changed, 39 insertions(+), 28 deletions(-) diff --git a/src/wasm-stack.h b/src/wasm-stack.h index 05898623187..6a6649d0aa4 100644 --- a/src/wasm-stack.h +++ b/src/wasm-stack.h @@ -94,7 +94,9 @@ class BinaryInstWriter : public OverriddenVisitor { BufferWithRandomAccess& o, Function* func, bool DWARF) - : parent(parent), o(o), func(func), DWARF(DWARF) {} + : parent(parent), o(o), func(func), DWARF(DWARF), + numLocalsByType(parent.getModule()->features), + scratchLocals(parent.getModule()->features) {} void visit(Expression* curr) { if (func) { @@ -144,20 +146,45 @@ class BinaryInstWriter : public OverriddenVisitor { std::vector breakStack; + // Map types to indices or counts, but transparently convert the key types to + // their written versions given the enabled features. + struct TypeIndexMap : private InsertOrderedMap { + FeatureSet feats; + + public: + using InsertOrderedMap::iterator; + using InsertOrderedMap::const_iterator; + using InsertOrderedMap::begin; + using InsertOrderedMap::end; + using InsertOrderedMap::size; + + TypeIndexMap(FeatureSet feats) + : InsertOrderedMap(), feats(feats) {} + + Index& operator[](Type type) { + return InsertOrderedMap::operator[](type.asWrittenGivenFeatures(feats)); + } + iterator find(Type type) { + return InsertOrderedMap::find(type.asWrittenGivenFeatures(feats)); + } + const_iterator find(Type type) const { + return InsertOrderedMap::find(type.asWrittenGivenFeatures(feats)); + } + }; + // The types of locals in the compact form, in order. std::vector localTypes; // type => number of locals of that type in the compact form - std::unordered_map numLocalsByType; + TypeIndexMap numLocalsByType; void noteLocalType(Type type, Index count = 1); - Index getNumLocalsForType(Type type); // Keeps track of the binary index of the scratch locals used to lower // tuple.extract. If there are multiple scratch locals of the same type, they // are contiguous and this map holds the index of the first. - InsertOrderedMap scratchLocals; + TypeIndexMap scratchLocals; // Return the type and number of required scratch locals. - InsertOrderedMap countScratchLocals(); + TypeIndexMap countScratchLocals(); // local.get, local.tee, and global.get expressions that will be followed by // tuple.extracts. We can optimize these by getting only the local for the diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp index ea932dc175c..53675586eed 100644 --- a/src/wasm/wasm-stack.cpp +++ b/src/wasm/wasm-stack.cpp @@ -15,10 +15,8 @@ */ #include "wasm-stack.h" -#include "ir/find_all.h" #include "ir/properties.h" #include "wasm-binary.h" -#include "wasm-debug.h" namespace wasm { @@ -3241,11 +3239,11 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() { // Map IR (local index, tuple index) pairs to binary local indices. Since // locals are grouped by type, start by calculating the base indices for each // type. - std::unordered_map nextFreeIndex; + TypeIndexMap nextFreeIndex(parent.getModule()->features); Index baseIndex = func->getVarIndexBase(); for (auto& type : localTypes) { nextFreeIndex[type] = baseIndex; - baseIndex += getNumLocalsForType(type); + baseIndex += numLocalsByType[type]; } // Map the IR index pairs to indices. @@ -3263,18 +3261,12 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() { o << U32LEB(localTypes.size()); for (auto& localType : localTypes) { - o << U32LEB(getNumLocalsForType(localType)); + o << U32LEB(numLocalsByType[localType]); parent.writeType(localType); } } void BinaryInstWriter::noteLocalType(Type type, Index count) { - // Group locals by the type they will eventually be written out as. For - // example, we do not need to differentiate exact and inexact versions of the - // same reference type if custom descriptors is not enabled and the type will - // be written as inexact either way. - auto feats = parent.getModule()->features; - type = type.asWrittenGivenFeatures(feats); auto& num = numLocalsByType[type]; if (num == 0) { localTypes.push_back(type); @@ -3282,21 +3274,13 @@ void BinaryInstWriter::noteLocalType(Type type, Index count) { num += count; } -Index BinaryInstWriter::getNumLocalsForType(Type type) { - auto feats = parent.getModule()->features; - type = type.asWrittenGivenFeatures(feats); - if (auto it = numLocalsByType.find(type); it != numLocalsByType.end()) { - return it->second; - } - return 0; -} - -InsertOrderedMap BinaryInstWriter::countScratchLocals() { +BinaryInstWriter::TypeIndexMap BinaryInstWriter::countScratchLocals() { struct ScratchLocalFinder : PostWalker { BinaryInstWriter& parent; - InsertOrderedMap scratches; + TypeIndexMap scratches; - ScratchLocalFinder(BinaryInstWriter& parent) : parent(parent) {} + ScratchLocalFinder(BinaryInstWriter& parent) + : parent(parent), scratches(parent.parent.getModule()->features) {} void visitTupleExtract(TupleExtract* curr) { if (curr->type == Type::unreachable) {