1 files changed, 35 insertions, 57 deletions

diff --git a/src/passes/ReorderGlobals.cpp b/src/passes/ReorderGlobals.cpp
index dab4625fe..9f934bb17 100644
--- a/src/passes/ReorderGlobals.cpp
+++ b/src/passes/ReorderGlobals.cpp
@@ -77,7 +77,7 @@ struct ReorderGlobals : public Pass {
   // use the index of the global in the original ordering to identify each
   // global. A different ordering is then a vector of old indices, saying where
   // each element comes from, which is logically a mapping between indices.
-  using IndexIndexMap = std::vector<Index>;
+  using IndexIndexMap = std::vector<size_t>;
 
   // We will also track counts of uses for each global. We use floating-point
   // values here since while the initial counts are integers, we will be
@@ -93,8 +93,8 @@ struct ReorderGlobals : public Pass {
   // To do so we construct a map from each global to those it depends on. We
   // also build the reverse map, of those that it is depended upon by.
   struct Dependencies {
-    std::unordered_map<Index, std::unordered_set<Index>> dependsOn;
-    std::unordered_map<Index, std::unordered_set<Index>> dependedUpon;
+    TopologicalSort::Graph dependsOn;
+    TopologicalSort::Graph dependedUpon;
   };
 
   void run(Module* module) override {
@@ -133,17 +133,26 @@ struct ReorderGlobals : public Pass {
     }
 
     // Compute dependencies.
-    Dependencies deps;
+    std::vector<std::unordered_set<size_t>> dependsOn(globals.size()),
+      dependedUpon(globals.size());
     for (Index i = 0; i < globals.size(); i++) {
       auto& global = globals[i];
       if (!global->imported()) {
         for (auto* get : FindAll<GlobalGet>(global->init).list) {
           auto getIndex = originalIndices[get->name];
-          deps.dependsOn[i].insert(getIndex);
-          deps.dependedUpon[getIndex].insert(i);
+          dependsOn[i].insert(getIndex);
+          dependedUpon[getIndex].insert(i);
         }
       }
     }
+    Dependencies deps;
+    deps.dependsOn.reserve(globals.size());
+    deps.dependedUpon.reserve(globals.size());
+    for (size_t i = 0; i < globals.size(); ++i) {
+      deps.dependsOn.emplace_back(dependsOn[i].begin(), dependsOn[i].end());
+      deps.dependedUpon.emplace_back(dependedUpon[i].begin(),
+                                     dependedUpon[i].end());
+    }
 
     // Compute various sorting options. All the options use a variation of the
     // algorithm in doSort() below, see there for more details; the only
@@ -190,16 +199,7 @@ struct ReorderGlobals : public Pass {
     double const EXPONENTIAL_FACTOR = 0.095;
     IndexCountMap sumCounts(globals.size()), exponentialCounts(globals.size());
 
-    std::vector<std::vector<size_t>> dependenceGraph(globals.size());
-    for (size_t i = 0; i < globals.size(); ++i) {
-      if (auto it = deps.dependsOn.find(i); it != deps.dependsOn.end()) {
-        for (auto dep : it->second) {
-          dependenceGraph[i].push_back(dep);
-        }
-      }
-    }
-
-    for (auto global : TopologicalSort::sort(dependenceGraph)) {
+    for (auto global : TopologicalSort::sort(deps.dependsOn)) {
       // We can compute this global's count as in the sorted order all the
       // values it cares about are resolved. Start with the self-count, then
       // add the deps.
@@ -236,8 +236,6 @@ struct ReorderGlobals : public Pass {
   IndexIndexMap doSort(const IndexCountMap& counts,
                        const Dependencies& deps,
                        Module* module) {
-    auto& globals = module->globals;
-
     // To sort the globals we do a simple greedy approach of always picking the
     // global with the highest count at every point in time, subject to the
     // constraint that we can only emit globals that have all of their
@@ -249,51 +247,31 @@ struct ReorderGlobals : public Pass {
     // global $c that depends on $b, and $c might have a much higher use count
     // than $a. For that reason we try several variations of this with different
     // counts, see earlier.
-    //
-    // Sort the globals into the optimal order based on the counts, ignoring
-    // dependencies for now.
-    std::vector<Index> sortedGlobals;
-    sortedGlobals.resize(globals.size());
-    for (Index i = 0; i < globals.size(); ++i) {
-      sortedGlobals[i] = i;
-    }
-    std::sort(
-      sortedGlobals.begin(), sortedGlobals.end(), [&](Index a, Index b) {
-        // Imports always go first. The binary writer takes care of this itself
-        // anyhow, but it is better to do it here in the IR so we can actually
-        // see what the final layout will be.
-        auto aImported = globals[a]->imported();
-        auto bImported = globals[b]->imported();
-        if (aImported != bImported) {
-          return aImported;
-        }
-
-        // Sort by the counts. Higher counts come first.
-        auto aCount = counts[a];
-        auto bCount = counts[b];
-        if (aCount != bCount) {
-          return aCount > bCount;
-        }
-
-        // Break ties using the original order, which means just using the
-        // indices we have.
-        return a < b;
-      });
 
     // Now use that optimal order to create an ordered graph that includes the
     // dependencies. The final order will be the minimum topological sort of
     // this graph.
-    std::vector<std::pair<Index, std::vector<Index>>> graph;
-    graph.reserve(globals.size());
-    for (auto i : sortedGlobals) {
-      std::vector<Index> children;
-      if (auto it = deps.dependedUpon.find(i); it != deps.dependedUpon.end()) {
-        children = std::vector<Index>(it->second.begin(), it->second.end());
+    return TopologicalSort::minSort(deps.dependedUpon, [&](size_t a, size_t b) {
+      // Imports always go first. The binary writer takes care of this itself
+      // anyhow, but it is better to do it here in the IR so we can actually
+      // see what the final layout will be.
+      auto aImported = module->globals[a]->imported();
+      auto bImported = module->globals[b]->imported();
+      if (aImported != bImported) {
+        return aImported;
+      }
+
+      // Sort by the counts. Higher counts come first.
+      auto aCount = counts[a];
+      auto bCount = counts[b];
+      if (aCount != bCount) {
+        return aCount > bCount;
       }
-      graph.emplace_back(i, std::move(children));
-    }
 
-    return TopologicalSort::minSortOf(graph.begin(), graph.end());
+      // Break ties using the original order, which means just using the
+      // indices.
+      return a < b;
+    });
   }
 
   // Given an indexing of the globals and the counts of how many times each is