Fix optimized grouping of locals (#8577)

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.

Author: tlively

src/wasm-stack.h

@@ -94,7 +94,9 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
                    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<BinaryInstWriter> {
 
   std::vector<Name> 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<Type, Index> {
+    FeatureSet feats;
+
+  public:
+    using InsertOrderedMap<Type, Index>::iterator;
+    using InsertOrderedMap<Type, Index>::const_iterator;
+    using InsertOrderedMap<Type, Index>::begin;
+    using InsertOrderedMap<Type, Index>::end;
+    using InsertOrderedMap<Type, Index>::size;
+
+    TypeIndexMap(FeatureSet feats)
+      : InsertOrderedMap<Type, Index>(), 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<Type> localTypes;
   // type => number of locals of that type in the compact form
-  std::unordered_map<Type, size_t> 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<Type, Index> scratchLocals;
+  TypeIndexMap scratchLocals;
   // Return the type and number of required scratch locals.
-  InsertOrderedMap<Type, Index> 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

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 {
 
@@ -3245,11 +3243,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<Type, Index> 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.
@@ -3267,40 +3265,26 @@ 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);
   }
   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<Type, Index> BinaryInstWriter::countScratchLocals() {
+BinaryInstWriter::TypeIndexMap BinaryInstWriter::countScratchLocals() {
   struct ScratchLocalFinder : PostWalker<ScratchLocalFinder> {
     BinaryInstWriter& parent;
-    InsertOrderedMap<Type, Index> 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) {