Add a utility for finding minimal topological sorts (#6884)

Reuse the code implementing Kahn's topological sort algorithm with a new
configuration that uses a min-heap to always choose the best available
element.

Also add wrapper utilities that can find topological sorts of graphs
with arbitrary element types, not just indices.

2 files changed, 115 insertions, 9 deletions

diff --git a/src/support/topological_orders.cpp b/src/support/topological_orders.cpp
index 145ba3004..0536d4ed9 100644
--- a/src/support/topological_orders.cpp
+++ b/src/support/topological_orders.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <algorithm>
 #include <cassert>
 
 #include "topological_orders.h"
@@ -21,9 +22,13 @@
 namespace wasm {
 
 TopologicalOrders::Selector
-TopologicalOrders::Selector::select(TopologicalOrders& ctx) {
+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.
@@ -33,7 +38,12 @@ TopologicalOrders::Selector::select(TopologicalOrders& ctx) {
   for (auto child : ctx.graph[selection]) {
     assert(ctx.indegrees[child] > 0);
     if (--ctx.indegrees[child] == 0) {
-      ctx.buf[next.start + next.count++] = child;
+      if (method == MinHeap) {
+        ctx.pushChoice(child);
+      } else {
+        ctx.buf[next.start + next.count] = child;
+      }
+      ++next.count;
     }
   }
   return next;
@@ -69,7 +79,7 @@ TopologicalOrders::Selector::advance(TopologicalOrders& ctx) {
 }
 
 TopologicalOrders::TopologicalOrders(
-  const std::vector<std::vector<size_t>>& graph)
+  const std::vector<std::vector<size_t>>& graph, SelectionMethod method)
   : graph(graph), indegrees(graph.size()), buf(graph.size()) {
   if (graph.size() == 0) {
     return;
@@ -86,13 +96,19 @@ TopologicalOrders::TopologicalOrders(
   auto& first = selectors.back();
   for (size_t i = 0; i < graph.size(); ++i) {
     if (indegrees[i] == 0) {
-      buf[first.count++] = i;
+      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));
+    selectors.push_back(selectors.back().select(*this, method));
   }
+  selectors.back().select(*this, method);
 }
 
 TopologicalOrders& TopologicalOrders::operator++() {
@@ -117,4 +133,16 @@ TopologicalOrders& TopologicalOrders::operator++() {
   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 wasm
diff --git a/src/support/topological_orders.h b/src/support/topological_orders.h
index 48941c021..4332f7a91 100644
--- a/src/support/topological_orders.h
+++ b/src/support/topological_orders.h
@@ -17,8 +17,10 @@
 #ifndef wasm_support_topological_orders_h
 #define wasm_support_topological_orders_h
 
+#include <cassert>
 #include <cstddef>
 #include <optional>
+#include <unordered_map>
 #include <vector>
 
 namespace wasm {
@@ -36,7 +38,8 @@ 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 std::vector<std::vector<size_t>>& graph);
+  TopologicalOrders(const std::vector<std::vector<size_t>>& graph)
+    : TopologicalOrders(graph, InPlace) {}
 
   TopologicalOrders begin() { return TopologicalOrders(graph); }
   TopologicalOrders end() { return TopologicalOrders({}); }
@@ -47,11 +50,16 @@ struct TopologicalOrders {
   bool operator!=(const TopologicalOrders& other) const {
     return !(*this == other);
   }
-  const std::vector<size_t>& operator*() { return buf; }
-  const std::vector<size_t>* operator->() { return &buf; }
+  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); }
 
+protected:
+  enum SelectionMethod { InPlace, MinHeap };
+  TopologicalOrders(const std::vector<std::vector<size_t>>& graph,
+                    SelectionMethod method);
+
 private:
   // The input graph given as an adjacency list with edges from vertices to
   // their dependent children.
@@ -64,6 +72,9 @@ 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
+  // choices in this separate min-heap instead of directly in `buf`.
+  std::vector<size_t> choiceHeap;
 
   // The state for tracking the possible choices for a single vertex in the
   // output order.
@@ -78,7 +89,7 @@ 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);
+    Selector select(TopologicalOrders& ctx, SelectionMethod method);
 
     // Undo the current selection, move the next selection into the first
     // position and return the new selector for the next position. Returns
@@ -86,11 +97,78 @@ private:
     std::optional<Selector> advance(TopologicalOrders& ctx);
   };
 
+  void pushChoice(size_t);
+  size_t popChoice();
+
   // 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;
 };
 
+// A utility for finding a single topological order of a graph.
+struct TopologicalSort : private TopologicalOrders {
+  TopologicalSort(const std::vector<std::vector<size_t>>& graph)
+    : TopologicalOrders(graph) {}
+
+  const std::vector<size_t>& operator*() const {
+    return TopologicalOrders::operator*();
+  }
+};
+
+// A utility for finding the topological order that is as close as possible to
+// the original order of elements. Internally uses a min-heap to choose the best
+// available next element.
+struct MinTopologicalSort : private TopologicalOrders {
+  MinTopologicalSort(const std::vector<std::vector<size_t>>& graph)
+    : TopologicalOrders(graph, MinHeap) {}
+
+  const std::vector<size_t>& operator*() const {
+    return TopologicalOrders::operator*();
+  }
+};
+
+// A utility that finds a topological sort of a graph with arbitrary element
+// types.
+template<typename T, typename TopoSort = TopologicalSort>
+struct TopologicalSortOf {
+  std::vector<T> order;
+
+  // The value of the iterators must be a pair of an element and an iterable of
+  // its children.
+  template<typename It> TopologicalSortOf(It begin, It end) {
+    std::unordered_map<T, size_t> indices;
+    std::vector<T> elements;
+    // Assign indices to each element.
+    for (auto it = begin; it != end; ++it) {
+      auto inserted = indices.insert({it->first, elements.size()});
+      assert(inserted.second && "unexpected repeat element");
+      elements.push_back(inserted.first->first);
+    }
+    // Collect the graph in terms of indices.
+    std::vector<std::vector<size_t>> indexGraph;
+    indexGraph.reserve(elements.size());
+    for (auto it = begin; it != end; ++it) {
+      indexGraph.emplace_back();
+      for (const auto& child : it->second) {
+        indexGraph.back().push_back(indices.at(child));
+      }
+    }
+    // Compute the topological order and convert back to original elements.
+    order.reserve(elements.size());
+    auto indexOrder = *TopoSort(indexGraph);
+    for (auto i : indexOrder) {
+      order.emplace_back(std::move(elements[i]));
+    }
+  }
+
+  const std::vector<T>& operator*() const { return order; }
+};
+
+// A utility that finds the minimum topological sort of a graph with arbitrary
+// element types.
+template<typename T>
+using MinTopologicalSortOf = TopologicalSortOf<T, MinTopologicalSort>;
+
 } // namespace wasm
 
 #endif // wasm_support_topological_orders_h