1 files changed, 76 insertions, 11 deletions

diff --git a/src/passes/ReorderLocals.cpp b/src/passes/ReorderLocals.cpp
index 66ea131de..22295dc1a 100644
--- a/src/passes/ReorderLocals.cpp
+++ b/src/passes/ReorderLocals.cpp
@@ -29,29 +29,94 @@ namespace wasm {
 struct ReorderLocals : public WalkerPass<PostWalker<ReorderLocals>> {
   bool isFunctionParallel() { return true; }
 
-  std::map<Name, uint32_t> counts;
+  std::map<Index, uint32_t> counts;
 
   void visitFunction(Function *curr) {
-    auto& vars = curr->vars;
-    sort(vars.begin(), vars.end(), [this](NameType a, NameType b) -> bool {
-      if (this->counts[a.name] == this->counts[b.name]) {
-        return strcmp(a.name.str, b.name.str) > 0;
+    Index num = curr->getNumLocals();
+    std::vector<Index> newToOld;
+    for (size_t i = 0; i < num; i++) {
+      newToOld.push_back(i);
+    }
+    // sort, keeping params in front (where they will not be moved)
+    sort(newToOld.begin(), newToOld.end(), [this, curr, &newToOld](Index a, Index b) -> bool {
+      if (curr->isParam(a) && !curr->isParam(b)) return true;
+      if (curr->isParam(b) && !curr->isParam(a)) return false;
+      if (curr->isParam(b) && curr->isParam(a)) {
+        return a < b;
+      }
+      if (this->counts[a] == this->counts[b]) {
+        return a < b;
       }
-      return this->counts[a.name] > this->counts[b.name];
+      return this->counts[a] > this->counts[b];
     });
-    // drop completely unused vars
-    while (vars.size() > 0 && counts[vars.back().name] == 0) {
-      vars.pop_back();
+    // sorting left params in front, perhaps slightly reordered. verify and fix.
+    for (size_t i = 0; i < curr->params.size(); i++) {
+      assert(newToOld[i] < curr->params.size());
+    }
+    for (size_t i = 0; i < curr->params.size(); i++) {
+      newToOld[i] = i;
+    }
+    // sort vars, and drop unused ones
+    auto oldVars = curr->vars;
+    curr->vars.clear();
+    for (size_t i = curr->getVarIndexBase(); i < newToOld.size(); i++) {
+      Index index = newToOld[i];
+      if (counts[index] > 0) {
+        curr->vars.push_back(oldVars[index - curr->getVarIndexBase()]);
+      } else {
+        newToOld.resize(i);
+        break;
+      }
     }
     counts.clear();
+    std::vector<Index> oldToNew;
+    oldToNew.resize(num);
+    for (size_t i = 0; i < newToOld.size(); i++) {
+      if (curr->isParam(i)) {
+        oldToNew[i] = i;
+      } else {
+        oldToNew[newToOld[i]] = i;
+      }
+    }
+    // apply the renaming to AST nodes
+    struct ReIndexer : public PostWalker<ReIndexer> {
+      Function* func;
+      std::vector<Index>& oldToNew;
+
+      ReIndexer(Function* func, std::vector<Index>& oldToNew) : func(func), oldToNew(oldToNew) {}
+
+      void visitGetLocal(GetLocal *curr) {
+        if (func->isVar(curr->index)) {
+          curr->index = oldToNew[curr->index];
+        }
+      }
+
+      void visitSetLocal(SetLocal *curr) {
+        if (func->isVar(curr->index)) {
+          curr->index = oldToNew[curr->index];
+        }
+      }
+    };
+    ReIndexer reIndexer(curr, oldToNew);
+    reIndexer.walk(curr->body);
+    // apply to the names
+    auto oldLocalNames = curr->localNames;
+    auto oldLocalIndices = curr->localIndices;
+    curr->localNames.clear();
+    curr->localNames.resize(newToOld.size());
+    curr->localIndices.clear();
+    for (size_t i = 0; i < newToOld.size(); i++) {
+      curr->localNames[i] = oldLocalNames[newToOld[i]];
+      curr->localIndices[oldLocalNames[newToOld[i]]] = i;
+    }
   }
 
   void visitGetLocal(GetLocal *curr) {
-    counts[curr->name]++;
+    counts[curr->index]++;
   }
 
   void visitSetLocal(SetLocal *curr) {
-    counts[curr->name]++;
+    counts[curr->index]++;
   }
 };