Group locals by written type in binary writer

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;