Add a utility for iterating over all topological orders (#6801)

Use an extension of Kahn's algorithm for finding topological orders that
iteratively makes every possible choice at every step to find all the
topological orders. The order being constructed and the set of possible
choices are managed in-place in the same buffer, so the algorithm takes
linear time and space plus amortized constant time per generated order.

This will be used in an upcoming type optimization.

4 files changed, 225 insertions, 2 deletions

diff --git a/src/support/CMakeLists.txt b/src/support/CMakeLists.txt
index 54ee7b206..20470a3cd 100644
--- a/src/support/CMakeLists.txt
+++ b/src/support/CMakeLists.txt
@@ -14,6 +14,7 @@ set(support_SOURCES
   safe_integer.cpp
   string.cpp
   threads.cpp
+  topological_orders.cpp
   utilities.cpp
   ${support_HEADERS}
 )
diff --git a/src/support/strongly_connected_components.h b/src/support/strongly_connected_components.h
index d54200ad9..7fcbdf14c 100644
--- a/src/support/strongly_connected_components.h
+++ b/src/support/strongly_connected_components.h
@@ -223,8 +223,8 @@ public:
   // Iterate over the strongly connected components of the graph.
   struct Iterator {
     using value_type = SCC;
-    using different_type = std::ptrdiff_t;
-    using reference = SCC;
+    using difference_type = std::ptrdiff_t;
+    using reference = SCC&;
     using pointer = SCC*;
     using iterator_category = std::input_iterator_tag;
 
diff --git a/src/support/topological_orders.cpp b/src/support/topological_orders.cpp
new file mode 100644
index 000000000..939c76c6e
--- /dev/null
+++ b/src/support/topological_orders.cpp
@@ -0,0 +1,118 @@
+/*
+ * 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 <cassert>
+
+#include "topological_orders.h"
+
+namespace wasm {
+
+TopologicalOrders::Selector
+TopologicalOrders::Selector::select(TopologicalOrders& ctx) {
+  assert(count >= 1);
+  assert(start + count <= ctx.buf.size());
+  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) {
+      ctx.buf[next.start + next.count++] = child;
+    }
+  }
+  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 std::vector<std::vector<size_t>>& graph)
+  : 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) {
+      buf[first.count++] = i;
+    }
+  }
+  // Initialize the full stack of selectors.
+  while (selectors.size() < graph.size()) {
+    selectors.push_back(selectors.back().select(*this));
+  }
+}
+
+void TopologicalOrders::advance() {
+  // 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;
+  }
+  // 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));
+  }
+}
+
+} // namespace wasm
diff --git a/src/support/topological_orders.h b/src/support/topological_orders.h
new file mode 100644
index 000000000..acab78ee5
--- /dev/null
+++ b/src/support/topological_orders.h
@@ -0,0 +1,104 @@
+/*
+ * 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.
+ */
+
+#ifndef wasm_support_topological_orders_h
+#define wasm_support_topological_orders_h
+
+#include <cstddef>
+#include <optional>
+#include <vector>
+
+namespace wasm {
+
+// 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 {
+  // 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);
+
+  struct Iterator {
+    using value_type = const std::vector<size_t>;
+    using difference_type = std::ptrdiff_t;
+    using reference = const std::vector<size_t>&;
+    using pointer = const std::vector<size_t>*;
+    using iterator_category = std::input_iterator_tag;
+
+    TopologicalOrders* parent;
+
+    bool isEnd() const { return !parent || parent->selectors.empty(); }
+    bool operator==(const Iterator& other) const {
+      return isEnd() == other.isEnd();
+    }
+    bool operator!=(const Iterator& other) const { return !(*this == other); }
+    const std::vector<size_t>& operator*() { return parent->buf; }
+    const std::vector<size_t>* operator->() { return &parent->buf; }
+    Iterator& operator++() {
+      parent->advance();
+      return *this;
+    }
+    Iterator operator++(int) { return ++(*this); }
+  };
+
+  Iterator begin() { return {this}; }
+  Iterator end() { return {nullptr}; }
+
+private:
+  // The input graph given as an adjacency list with edges from vertices to
+  // their dependent children.
+  const std::vector<std::vector<size_t>>& graph;
+  // The current in-degrees for each vertex. When a vertex is appended to our
+  // permutation, the in-degrees of its children are decremented and those that
+  // go to zero become available for the next selection.
+  std::vector<size_t> indegrees;
+  // The buffer in which we are constructing a permutation. It contains a
+  // sequence of selected vertices followed by a sequence of possible choices
+  // for the next vertex.
+  std::vector<size_t> buf;
+
+  // The state for tracking the possible choices for a single vertex in the
+  // output order.
+  struct Selector {
+    // The start index of the sequence of available choices. Also the index
+    // where we place the current choice.
+    size_t start;
+    // The number of choices we have.
+    size_t count;
+    // The index of the current choice in the original order.
+    size_t index;
+
+    // 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);
+
+    // 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);
+  };
+
+  // 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;
+
+  void advance();
+};
+
+} // namespace wasm
+
+#endif // wasm_support_topological_orders_h