[NFC] Take custom comparator in TopologicalSort::minSort (#6890)

Rather than finding the minimum sort with respect to the original order
of vertices, find the minimum sort with respect to an arbitrary
user-provided comparator. Users of the minSort utility previously had to
sort their input graphs according to their desired ordering, but now
they can simply provide their comparator instead.

Take advantage of the new functionality in ReorderGlobals and also
standardize on a single data type for representing dependence graphs to
avoid unnecessary conversions. Together, these changes slightly speed up
ReorderGlobals.

Move the topological sort code previously in a .cpp file into the header
so the comparator can be provided as a lambda template parameter instead
of as a `std::function`. This makes ReorderGlobals about 5% faster.

4 files changed, 200 insertions, 248 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
diff --git a/src/support/CMakeLists.txt b/src/support/CMakeLists.txt
index 20470a3cd..54ee7b206 100644
--- a/src/support/CMakeLists.txt
+++ b/src/support/CMakeLists.txt
@@ -14,7 +14,6 @@ set(support_SOURCES
   safe_integer.cpp
   string.cpp
   threads.cpp
-  topological_orders.cpp
   utilities.cpp
   ${support_HEADERS}
 )
diff --git a/src/support/topological_orders.cpp b/src/support/topological_orders.cpp
deleted file mode 100644
index aeec95d18..000000000
--- a/src/support/topological_orders.cpp
+++ /dev/null
@@ -1,159 +0,0 @@
-/*
- * Copyright 2024 WebAssembly Community Group participants
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <algorithm>
-#include <cassert>
-
-#include "topological_orders.h"
-
-namespace wasm {
-
-TopologicalOrders::Selector
-TopologicalOrders::Selector::select(TopologicalOrders& ctx,
-                                    SelectionMethod method = InPlace) {
-  assert(count >= 1);
-  assert(start + count <= ctx.buf.size());
-  if (method == MinHeap) {
-    ctx.buf[start] = ctx.popChoice();
-  }
-  auto selection = ctx.buf[start];
-  // The next selector will select the next index and will not be able to choose
-  // the vertex we just selected.
-  Selector next = {start + 1, count - 1, 0};
-  // Append any child that this selection makes available to the choices for the
-  // next selector.
-  for (auto child : ctx.graph[selection]) {
-    assert(ctx.indegrees[child] > 0);
-    if (--ctx.indegrees[child] == 0) {
-      if (method == MinHeap) {
-        ctx.pushChoice(child);
-      } else {
-        ctx.buf[next.start + next.count] = child;
-      }
-      ++next.count;
-    }
-  }
-  return next;
-}
-
-std::optional<TopologicalOrders::Selector>
-TopologicalOrders::Selector::advance(TopologicalOrders& ctx) {
-  assert(count >= 1);
-  // Undo the current selection. Backtrack by incrementing the in-degree for
-  // each child of the vertex we are unselecting. No need to remove the newly
-  // unavailable children from the buffer; they will be overwritten with valid
-  // selections.
-  auto unselected = ctx.buf[start];
-  for (auto child : ctx.graph[unselected]) {
-    ++ctx.indegrees[child];
-  }
-  if (index == count - 1) {
-    // We are wrapping back to the original configuration. The current selection
-    // element needs to go back on the end and everything else needs to be
-    // shifted back to its original location. This ensures that we leave
-    // everything how we found it so the previous selector can make its next
-    // selection without observing anything having changed in the meantime.
-    for (size_t i = 1; i < count; ++i) {
-      ctx.buf[start + i - 1] = ctx.buf[start + i];
-    }
-    ctx.buf[start + count - 1] = unselected;
-    return std::nullopt;
-  }
-  // Otherwise, just swap the next selection into the first position and
-  // finalize the selection.
-  std::swap(ctx.buf[start], ctx.buf[start + ++index]);
-  return select(ctx);
-}
-
-TopologicalOrders::TopologicalOrders(const Graph& graph, SelectionMethod method)
-  : graph(graph), indegrees(graph.size()), buf(graph.size()) {
-  if (graph.size() == 0) {
-    return;
-  }
-  // Find the in-degree of each vertex.
-  for (const auto& vertex : graph) {
-    for (auto child : vertex) {
-      ++indegrees[child];
-    }
-  }
-  // Set up the first selector with its possible selections.
-  selectors.reserve(graph.size());
-  selectors.push_back({0, 0, 0});
-  auto& first = selectors.back();
-  for (size_t i = 0; i < graph.size(); ++i) {
-    if (indegrees[i] == 0) {
-      if (method == MinHeap) {
-        pushChoice(i);
-      } else {
-        buf[first.count] = i;
-      }
-      ++first.count;
-    }
-  }
-  // Initialize the full stack of selectors.
-  while (selectors.size() < graph.size()) {
-    selectors.push_back(selectors.back().select(*this, method));
-  }
-  selectors.back().select(*this, method);
-}
-
-TopologicalOrders& TopologicalOrders::operator++() {
-  // Find the last selector that can be advanced, popping any that cannot.
-  std::optional<Selector> next;
-  while (!selectors.empty() && !(next = selectors.back().advance(*this))) {
-    selectors.pop_back();
-  }
-  if (!next) {
-    // No selector could be advanced, so we've seen every possible ordering.
-    assert(selectors.empty());
-    return *this;
-  }
-  // We've advanced the last selector on the stack, so initialize the
-  // subsequent selectors.
-  assert(selectors.size() < graph.size());
-  selectors.push_back(*next);
-  while (selectors.size() < graph.size()) {
-    selectors.push_back(selectors.back().select(*this));
-  }
-
-  return *this;
-}
-
-void TopologicalOrders::pushChoice(size_t choice) {
-  choiceHeap.push_back(choice);
-  std::push_heap(choiceHeap.begin(), choiceHeap.end(), std::greater<size_t>{});
-}
-
-size_t TopologicalOrders::popChoice() {
-  std::pop_heap(choiceHeap.begin(), choiceHeap.end(), std::greater<size_t>{});
-  auto choice = choiceHeap.back();
-  choiceHeap.pop_back();
-  return choice;
-}
-
-namespace TopologicalSort {
-
-std::vector<size_t> sort(const Graph& graph) {
-  return *TopologicalOrders(graph, TopologicalOrders::InPlace);
-}
-
-std::vector<size_t> minSort(const Graph& graph) {
-  return *TopologicalOrders(graph, TopologicalOrders::MinHeap);
-}
-
-} // namespace TopologicalSort
-
-} // namespace wasm
diff --git a/src/support/topological_orders.h b/src/support/topological_orders.h
index 1a5828247..925ad1117 100644
--- a/src/support/topological_orders.h
+++ b/src/support/topological_orders.h
@@ -17,11 +17,14 @@
 #ifndef wasm_support_topological_orders_h
 #define wasm_support_topological_orders_h
 
+#include <algorithm>
 #include <cassert>
 #include <cstddef>
+#include <functional>
 #include <optional>
 #include <type_traits>
 #include <unordered_map>
+#include <variant>
 #include <vector>
 
 namespace wasm {
@@ -32,17 +35,12 @@ namespace TopologicalSort {
 using Graph = std::vector<std::vector<size_t>>;
 
 // Return a topological sort of the vertices in the given adjacency graph.
-std::vector<size_t> sort(const Graph& graph);
-
-// Return the topological sort of the vertices in the given adjacency graph that
-// is closest to their original order. Implemented using a min-heap internally.
-std::vector<size_t> minSort(const Graph& graph);
+inline std::vector<size_t> sort(const Graph& graph);
 
 // A utility that finds a topological sort of a graph with arbitrary element
 // types. The provided iterators must be to pairs of elements and collections of
 // their children.
-template<typename It, std::vector<size_t> (*TopoSort)(const Graph&) = sort>
-decltype(auto) sortOf(It begin, It end) {
+template<typename It> decltype(auto) sortOf(It begin, It end) {
   using T = std::remove_cv_t<typename It::value_type::first_type>;
   std::unordered_map<T, size_t> indices;
   std::vector<T> elements;
@@ -64,26 +62,31 @@ decltype(auto) sortOf(It begin, It end) {
   // Compute the topological order and convert back to original elements.
   std::vector<T> order;
   order.reserve(elements.size());
-  auto indexOrder = TopoSort(indexGraph);
-  for (auto i : indexOrder) {
+  for (auto i : sort(indexGraph)) {
     order.emplace_back(std::move(elements[i]));
   }
   return order;
 }
 
-// Find the topological sort of a graph with arbitrary element types that is
-// closest to their original order.
-template<typename It> decltype(auto) minSortOf(It begin, It end) {
-  return sortOf<It, minSort>(begin, end);
-}
+// Return the topological sort of the vertices in the given adjacency graph that
+// is lexicographically minimal with respect to the provided comparator on
+// vertex indices. Implemented using a min-heap internally.
+template<typename F = std::less<size_t>>
+std::vector<size_t> minSort(const Graph& graph, F cmp = std::less<size_t>{});
 
 } // namespace TopologicalSort
 
+template<typename F = std::monostate> struct TopologicalOrdersImpl;
+
 // A utility for iterating through all possible topological orders in a graph
 // using an extension of Kahn's algorithm (see
 // https://en.wikipedia.org/wiki/Topological_sorting) that iteratively makes all
 // possible choices for each position of the output order.
-struct TopologicalOrders {
+using TopologicalOrders = TopologicalOrdersImpl<std::monostate>;
+
+// The template parameter `Cmp` is only used as in internal implementation
+// detail of `minSort`.
+template<typename Cmp> struct TopologicalOrdersImpl {
   using Graph = TopologicalSort::Graph;
   using value_type = const std::vector<size_t>;
   using difference_type = std::ptrdiff_t;
@@ -93,25 +96,75 @@ struct TopologicalOrders {
 
   // Takes an adjacency list, where the list for each vertex is a sorted list of
   // the indices of its children, which will appear after it in the order.
-  TopologicalOrders(const Graph& graph) : TopologicalOrders(graph, InPlace) {}
+  TopologicalOrdersImpl(const Graph& graph)
+    : TopologicalOrdersImpl(graph, {}) {}
 
-  TopologicalOrders begin() { return TopologicalOrders(graph); }
-  TopologicalOrders end() { return TopologicalOrders({}); }
+  TopologicalOrdersImpl begin() { return TopologicalOrdersImpl(graph); }
+  TopologicalOrdersImpl end() { return TopologicalOrdersImpl({}); }
 
-  bool operator==(const TopologicalOrders& other) const {
+  bool operator==(const TopologicalOrdersImpl& other) const {
     return selectors.empty() == other.selectors.empty();
   }
-  bool operator!=(const TopologicalOrders& other) const {
+  bool operator!=(const TopologicalOrdersImpl& other) const {
     return !(*this == other);
   }
   const std::vector<size_t>& operator*() const { return buf; }
   const std::vector<size_t>* operator->() const { return &buf; }
-  TopologicalOrders& operator++();
-  TopologicalOrders operator++(int) { return ++(*this); }
+  TopologicalOrdersImpl& operator++() {
+    // Find the last selector that can be advanced, popping any that cannot.
+    std::optional<Selector> next;
+    while (!selectors.empty() && !(next = selectors.back().advance(*this))) {
+      selectors.pop_back();
+    }
+    if (!next) {
+      // No selector could be advanced, so we've seen every possible ordering.
+      assert(selectors.empty());
+      return *this;
+    }
+    // We've advanced the last selector on the stack, so initialize the
+    // subsequent selectors.
+    assert(selectors.size() < graph.size());
+    selectors.push_back(*next);
+    while (selectors.size() < graph.size()) {
+      selectors.push_back(selectors.back().select(*this));
+    }
+
+    return *this;
+  }
+  TopologicalOrdersImpl operator++(int) { return ++(*this); }
 
 private:
-  enum SelectionMethod { InPlace, MinHeap };
-  TopologicalOrders(const Graph& graph, SelectionMethod method);
+  TopologicalOrdersImpl(const Graph& graph, Cmp cmp)
+    : graph(graph), indegrees(graph.size()), buf(graph.size()), cmp(cmp) {
+    if (graph.size() == 0) {
+      return;
+    }
+    // Find the in-degree of each vertex.
+    for (const auto& vertex : graph) {
+      for (auto child : vertex) {
+        ++indegrees[child];
+      }
+    }
+    // Set up the first selector with its possible selections.
+    selectors.reserve(graph.size());
+    selectors.push_back({0, 0, 0});
+    auto& first = selectors.back();
+    for (size_t i = 0; i < graph.size(); ++i) {
+      if (indegrees[i] == 0) {
+        if constexpr (useMinHeap) {
+          pushChoice(i);
+        } else {
+          buf[first.count] = i;
+        }
+        ++first.count;
+      }
+    }
+    // Initialize the full stack of selectors.
+    while (selectors.size() < graph.size()) {
+      selectors.push_back(selectors.back().select(*this));
+    }
+    selectors.back().select(*this);
+  }
 
   // The input graph given as an adjacency list with edges from vertices to
   // their dependent children.
@@ -124,9 +177,15 @@ private:
   // sequence of selected vertices followed by a sequence of possible choices
   // for the next vertex.
   std::vector<size_t> buf;
-  // When we are finding the minimal topological order, store the possible
+  // When we are finding a minimal topological order, store the possible
   // choices in this separate min-heap instead of directly in `buf`.
   std::vector<size_t> choiceHeap;
+  // When we are finding a minimal topological order, use this function to
+  // compare possible choices. Empty iff we are not finding a minimal
+  // topological order.
+  Cmp cmp;
+
+  static constexpr bool useMinHeap = !std::is_same_v<Cmp, std::monostate>;
 
   // The state for tracking the possible choices for a single vertex in the
   // output order.
@@ -141,25 +200,100 @@ private:
 
     // Select the next available vertex, decrement in-degrees, and update the
     // sequence of available vertices. Return the Selector for the next vertex.
-    Selector select(TopologicalOrders& ctx, SelectionMethod method);
+    Selector select(TopologicalOrdersImpl& ctx) {
+      assert(count >= 1);
+      assert(start + count <= ctx.buf.size());
+      if constexpr (TopologicalOrdersImpl::useMinHeap) {
+        ctx.buf[start] = ctx.popChoice();
+      }
+      auto selection = ctx.buf[start];
+      // The next selector will select the next index and will not be able to
+      // choose the vertex we just selected.
+      Selector next = {start + 1, count - 1, 0};
+      // Append any child that this selection makes available to the choices for
+      // the next selector.
+      for (auto child : ctx.graph[selection]) {
+        assert(ctx.indegrees[child] > 0);
+        if (--ctx.indegrees[child] == 0) {
+          if constexpr (TopologicalOrdersImpl::useMinHeap) {
+            ctx.pushChoice(child);
+          } else {
+            ctx.buf[next.start + next.count] = child;
+          }
+          ++next.count;
+        }
+      }
+      return next;
+    }
 
     // Undo the current selection, move the next selection into the first
     // position and return the new selector for the next position. Returns
     // nullopt if advancing wraps back around to the original configuration.
-    std::optional<Selector> advance(TopologicalOrders& ctx);
+    std::optional<Selector> advance(TopologicalOrdersImpl& ctx) {
+      assert(count >= 1);
+      // Undo the current selection. Backtrack by incrementing the in-degree for
+      // each child of the vertex we are unselecting. No need to remove the
+      // newly unavailable children from the buffer; they will be overwritten
+      // with valid selections.
+      auto unselected = ctx.buf[start];
+      for (auto child : ctx.graph[unselected]) {
+        ++ctx.indegrees[child];
+      }
+      if (index == count - 1) {
+        // We are wrapping back to the original configuration. The current
+        // selection element needs to go back on the end and everything else
+        // needs to be shifted back to its original location. This ensures that
+        // we leave everything how we found it so the previous selector can make
+        // its next selection without observing anything having changed in the
+        // meantime.
+        for (size_t i = 1; i < count; ++i) {
+          ctx.buf[start + i - 1] = ctx.buf[start + i];
+        }
+        ctx.buf[start + count - 1] = unselected;
+        return std::nullopt;
+      }
+      // Otherwise, just swap the next selection into the first position and
+      // finalize the selection.
+      std::swap(ctx.buf[start], ctx.buf[start + ++index]);
+      return select(ctx);
+    }
   };
 
-  void pushChoice(size_t);
-  size_t popChoice();
+  void pushChoice(size_t choice) {
+    choiceHeap.push_back(choice);
+    std::push_heap(choiceHeap.begin(),
+                   choiceHeap.end(),
+                   [&](size_t a, size_t b) { return cmp(b, a); });
+  }
+  size_t popChoice() {
+    std::pop_heap(choiceHeap.begin(),
+                  choiceHeap.end(),
+                  [&](size_t a, size_t b) { return cmp(b, a); });
+    auto choice = choiceHeap.back();
+    choiceHeap.pop_back();
+    return choice;
+  }
 
   // A stack of selectors, one for each vertex in a complete topological order.
   // Empty if we've already seen every possible ordering.
   std::vector<Selector> selectors;
 
-  friend std::vector<size_t> TopologicalSort::sort(const Graph&);
-  friend std::vector<size_t> TopologicalSort::minSort(const Graph&);
+  friend std::vector<size_t> TopologicalSort::minSort<Cmp>(const Graph&, Cmp);
 };
 
+namespace TopologicalSort {
+
+std::vector<size_t> sort(const Graph& graph) {
+  return *TopologicalOrders(graph);
+}
+
+template<typename Cmp>
+std::vector<size_t> minSort(const Graph& graph, Cmp cmp) {
+  return *TopologicalOrdersImpl<Cmp>(graph, cmp);
+}
+
+} // namespace TopologicalSort
+
 } // namespace wasm
 
 #endif // wasm_support_topological_orders_h