Group locals by written type in binary writer (#8568)

We have logic to group locals by type in the binary writer to take
advantage of the run-length encoding of locals. But that logic
previously grouped locals by their IR types, rather than the types that
would actually be written to the binary. These can differ in when the IR
uses more precise types than are allowed to be written given the enabled
feature set. For example, the IR might use exact types but have to write
inexact types because custom descriptors are not enabled. In such cases,
it is possible that different groups of locals would be written with the
same type, which is suboptimal. Fix the problem by grouping locals by
their written types given the enabled features.

Fixes #7934.

Author: tlively

src/wasm-stack.h

@@ -150,6 +150,7 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
   std::unordered_map<Type, size_t> 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

src/wasm/wasm-stack.cpp

@@ -3245,7 +3245,7 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() {
   Index baseIndex = func->getVarIndexBase();
   for (auto& type : localTypes) {
     nextFreeIndex[type] = baseIndex;
-    baseIndex += numLocalsByType[type];
+    baseIndex += getNumLocalsForType(type);
   }
 
   // Map the IR index pairs to indices.
@@ -3261,21 +3261,36 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() {
     scratchLocals[type] = nextFreeIndex[type];
   }
 
-  o << U32LEB(numLocalsByType.size());
+  o << U32LEB(localTypes.size());
   for (auto& localType : localTypes) {
-    o << U32LEB(numLocalsByType.at(localType));
+    o << U32LEB(getNumLocalsForType(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() {
   struct ScratchLocalFinder : PostWalker<ScratchLocalFinder> {
     BinaryInstWriter& parent;