Fix optimized grouping of locals

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 {
 
@@ -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<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.
@@ -3263,40 +3261,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) {