Fix binary writing of local name indexes (#3649)

When writing a binary, we take the local indexes in the IR and turn
them into the format in the binary, which clumps them by type. When
writing the names section we should be aware of that ordering, but
we never were, as noticed in #3499

This fixes that by saving the mapping of locals when we are emitting
the name section, then using it when emitting the local names.

This also fixes the order of the types themselves as part of the
refactoring. We used to depend on the ordering of types to decide
which to emit first, but that isn't good for at least two reasons. First,
it hits #3648 - that order is not fully
defined for recursive types. Also, it's not good for code size - we've
ordered the locals in a way we think is best already (ReorderLocals pass).
This PR makes us pick an order of types based on that, as much as
possible, that is, when we see a type for the first time we append it to
a list whose order we use.

Test changes: Some are just because we use a different order than
before, as in atomics64. But some are actual fixes, e.g. in heap-types
where we now have (local $tv (ref null $vector)) which is indeed
right - v there is for vector, and likewise m for matrix etc. - we
just had wrong names before. Another example, we now have
(local $local_externref externref) whereas before the name was
funcref, and which was wrong... seems like the incorrectness was
more common on reference types and GC types, which is why this was
not noticed before.

Fixes #3499

Makes part of #3648 moot.

4 files changed, 64 insertions, 25 deletions

diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index 9f298662b..d490a2a96 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -1090,6 +1090,9 @@ enum FeaturePrefix {
 
 } // namespace BinaryConsts
 
+// (local index in IR, tuple index) => binary local index
+using MappedLocals = std::unordered_map<std::pair<Index, Index>, size_t>;
+
 // Writes out wasm to the binary format
 
 class WasmBinaryWriter {
@@ -1275,6 +1278,11 @@ private:
   // the function is written out.
   std::vector<Expression*> binaryLocationTrackedExpressionsForFunc;
 
+  // Maps function names to their mapped locals. This is used when we emit the
+  // local names section: we map the locals when writing the function, save that
+  // info here, and then use it when writing the names.
+  std::unordered_map<Name, MappedLocals> funcMappedLocals;
+
   void prepare();
 };
 
diff --git a/src/wasm-stack.h b/src/wasm-stack.h
index 5af9adbce..38b64dff4 100644
--- a/src/wasm-stack.h
+++ b/src/wasm-stack.h
@@ -118,6 +118,8 @@ public:
   void emitUnreachable();
   void mapLocalsAndEmitHeader();
 
+  MappedLocals mappedLocals;
+
 private:
   void emitMemoryAccess(size_t alignment, size_t bytes, uint32_t offset);
   int32_t getBreakIndex(Name name);
@@ -130,10 +132,12 @@ private:
 
   std::vector<Name> breakStack;
 
+  // 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::map<Type, size_t> numLocalsByType;
-  // (local index, tuple index) => binary local index
-  std::map<std::pair<Index, Index>, size_t> mappedLocals;
+  std::unordered_map<Type, size_t> numLocalsByType;
+
+  void noteLocalType(Type type);
 
   // Keeps track of the binary index of the scratch locals used to lower
   // tuple.extract.
@@ -401,6 +405,8 @@ public:
     }
   }
 
+  MappedLocals& getMappedLocals() { return writer.mappedLocals; }
+
 private:
   WasmBinaryWriter& parent;
   BinaryInstWriter writer;
@@ -458,6 +464,8 @@ public:
 
   void write();
 
+  MappedLocals& getMappedLocals() { return writer.mappedLocals; }
+
 private:
   BinaryInstWriter writer;
   Function* func;
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 0993eb124..0f40cfd58 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -336,10 +336,18 @@ void WasmBinaryWriter::writeFunctions() {
     // Emit Stack IR if present, and if we can
     if (func->stackIR && !sourceMap && !DWARF) {
       BYN_TRACE("write Stack IR\n");
-      StackIRToBinaryWriter(*this, o, func).write();
+      StackIRToBinaryWriter writer(*this, o, func);
+      writer.write();
+      if (debugInfo) {
+        funcMappedLocals[func->name] = std::move(writer.getMappedLocals());
+      }
     } else {
       BYN_TRACE("write Binaryen IR\n");
-      BinaryenIRToBinaryWriter(*this, o, func, sourceMap, DWARF).write();
+      BinaryenIRToBinaryWriter writer(*this, o, func, sourceMap, DWARF);
+      writer.write();
+      if (debugInfo) {
+        funcMappedLocals[func->name] = std::move(writer.getMappedLocals());
+      }
     }
     size_t size = o.size() - start;
     assert(size <= std::numeric_limits<uint32_t>::max());
@@ -678,20 +686,28 @@ void WasmBinaryWriter::writeNames() {
         startSubsection(BinaryConsts::UserSections::Subsection::NameLocal);
       o << U32LEB(functionsWithLocalNames.size());
       Index emitted = 0;
-      for (auto& indexedFunc : functionsWithLocalNames) {
+      for (auto& kv : functionsWithLocalNames) {
+        auto index = kv.first;
+        auto* func = kv.second;
+        // Pairs of (local index in IR, name).
         std::vector<std::pair<Index, Name>> localsWithNames;
-        auto numLocals = indexedFunc.second->getNumLocals();
+        auto numLocals = func->getNumLocals();
         for (Index i = 0; i < numLocals; ++i) {
-          if (indexedFunc.second->hasLocalName(i)) {
-            localsWithNames.push_back({i, indexedFunc.second->getLocalName(i)});
+          if (func->hasLocalName(i)) {
+            localsWithNames.push_back({i, func->getLocalName(i)});
           }
         }
         assert(localsWithNames.size());
-        o << U32LEB(indexedFunc.first);
+        o << U32LEB(index);
         o << U32LEB(localsWithNames.size());
         for (auto& indexedLocal : localsWithNames) {
-          o << U32LEB(indexedLocal.first);
-          writeEscapedName(indexedLocal.second.str);
+          auto indexInFunc = indexedLocal.first;
+          auto name = indexedLocal.second;
+          // TODO: handle multivalue
+          auto indexInBinary =
+            funcMappedLocals.at(func->name)[{indexInFunc, 0}];
+          o << U32LEB(indexInBinary);
+          writeEscapedName(name.str);
         }
         emitted++;
       }
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index 4c4e4e9d1..0c7448c58 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2166,7 +2166,7 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() {
   }
   for (auto type : func->vars) {
     for (const auto& t : type) {
-      numLocalsByType[t]++;
+      noteLocalType(t);
     }
   }
   countScratchLocals();
@@ -2176,24 +2176,31 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() {
     for (const auto& type : func->getLocalType(i)) {
       auto fullIndex = std::make_pair(i, j++);
       Index index = func->getVarIndexBase();
-      for (auto& typeCount : numLocalsByType) {
-        if (type == typeCount.first) {
-          mappedLocals[fullIndex] = index + currLocalsByType[typeCount.first];
+      for (auto& localType : localTypes) {
+        if (type == localType) {
+          mappedLocals[fullIndex] = index + currLocalsByType[localType];
           currLocalsByType[type]++;
           break;
         }
-        index += typeCount.second;
+        index += numLocalsByType.at(localType);
       }
     }
   }
   setScratchLocals();
   o << U32LEB(numLocalsByType.size());
-  for (auto& typeCount : numLocalsByType) {
-    o << U32LEB(typeCount.second);
-    parent.writeType(typeCount.first);
+  for (auto& localType : localTypes) {
+    o << U32LEB(numLocalsByType.at(localType));
+    parent.writeType(localType);
   }
 }
 
+void BinaryInstWriter::noteLocalType(Type type) {
+  if (!numLocalsByType.count(type)) {
+    localTypes.push_back(type);
+  }
+  numLocalsByType[type]++;
+}
+
 void BinaryInstWriter::countScratchLocals() {
   // Add a scratch register in `numLocalsByType` for each type of
   // tuple.extract with nonzero index present.
@@ -2204,16 +2211,16 @@ void BinaryInstWriter::countScratchLocals() {
     }
   }
   for (auto t : scratchLocals) {
-    numLocalsByType[t.first]++;
+    noteLocalType(t.first);
   }
 }
 
 void BinaryInstWriter::setScratchLocals() {
   Index index = func->getVarIndexBase();
-  for (auto& typeCount : numLocalsByType) {
-    index += typeCount.second;
-    if (scratchLocals.find(typeCount.first) != scratchLocals.end()) {
-      scratchLocals[typeCount.first] = index - 1;
+  for (auto& localType : localTypes) {
+    index += numLocalsByType[localType];
+    if (scratchLocals.find(localType) != scratchLocals.end()) {
+      scratchLocals[localType] = index - 1;
     }
   }
 }