Find more intuitive orderings when there are multiple possible toposorts

loader/src/main/java/net/neoforged/fml/loading/toposort/TopologicalSort.java

@@ -7,15 +7,10 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.graph.Graph;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Comparator;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.PriorityQueue;
-import java.util.Queue;
 import java.util.Set;
 import org.jetbrains.annotations.Nullable;
 
@@ -26,7 +21,8 @@
  * utilized in other fashions, e.g. topology-based registry loading, prioritization
  * for renderers, and even mod module loading.
  */
-public final class TopologicalSort {
+public final class TopologicalSort<T> {
+    // TODO: this comment here needs updating, in particular the complexity is now O(VE) I believe
     /**
      * A breath-first-search based topological sort.
      *
@@ -61,42 +57,64 @@ public static <T> List<T> topologicalSort(Graph<T> graph, @Nullable Comparator<?
         Preconditions.checkArgument(graph.isDirected(), "Cannot topologically sort an undirected graph!");
         Preconditions.checkArgument(!graph.allowsSelfLoops(), "Cannot topologically sort a graph with self loops!");
 
-        final Queue<T> queue = comparator == null ? new ArrayDeque<>() : new PriorityQueue<>(comparator);
-        final Map<T, Integer> degrees = new HashMap<>();
-        final List<T> results = new ArrayList<>();
-
-        for (final T node : graph.nodes()) {
-            final int degree = graph.inDegree(node);
-            if (degree == 0) {
-                queue.add(node);
-            } else {
-                degrees.put(node, degree);
-            }
+        // Check for cycles
+        Set<Set<T>> components = new StronglyConnectedComponentDetector<>(graph).getComponents();
+        components.removeIf(set -> set.size() < 2);
+        if (!components.isEmpty()) {
+            //noinspection unchecked
+            throw new CyclePresentException((Set<Set<?>>) (Set<?>) components);
         }
 
-        while (!queue.isEmpty()) {
-            final T current = queue.remove();
-            results.add(current);
-            for (final T successor : graph.successors(current)) {
-                final int updated = degrees.compute(successor, (node, degree) -> Objects.requireNonNull(degree, () -> "Invalid degree present for " + node) - 1);
-                if (updated == 0) {
-                    queue.add(successor);
-                    degrees.remove(successor);
-                }
-            }
+        var toposort = new TopologicalSort<>(graph, comparator);
+        toposort.resolveSubgraph(new ArrayList<>(graph.nodes()));
+        return toposort.result;
+    }
+
+    private final Graph<T> graph;
+    @Nullable
+    private final Comparator<? super T> comparator;
+    private final Set<T> taken = new HashSet<>(); // nodes that have been added to `result` already
+    private final List<T> result = new ArrayList<>();
+
+    private TopologicalSort(Graph<T> graph, @Nullable Comparator<? super T> comparator) {
+        this.graph = graph;
+        this.comparator = comparator;
+    }
+
+    private void resolveSubgraph(List<T> nodes) {
+        while (!nodes.isEmpty()) {
+            var min = removeMin(nodes);
+            // Process all dependencies of `min`
+            var subGraph = new ArrayList<T>();
+            collectChildren(min, subGraph);
+            resolveSubgraph(subGraph);
+            nodes.removeIf(taken::contains);
+            // Add `min`
+            result.add(min);
+            taken.add(min);
         }
+    }
 
-        if (!degrees.isEmpty()) {
-            Set<Set<T>> components = new StronglyConnectedComponentDetector<>(graph).getComponents();
-            components.removeIf(set -> set.size() < 2);
-            throwCyclePresentException(components);
+    private T removeMin(List<T> nodes) {
+        if (this.comparator == null) {
+            return nodes.removeFirst();
         }
 
-        return results;
+        int minIndex = 0;
+        for (int i = 1; i < nodes.size(); ++i) {
+            if (this.comparator.compare(nodes.get(i), nodes.get(minIndex)) < 0) {
+                minIndex = i;
+            }
+        }
+        return nodes.remove(minIndex);
     }
 
-    @SuppressWarnings("unchecked") // for unchecked annotation
-    private static <T> void throwCyclePresentException(Set<Set<T>> components) {
-        throw new CyclePresentException((Set<Set<?>>) (Set<?>) components);
+    private void collectChildren(T node, List<T> out) {
+        for (var child : this.graph.predecessors(node)) {
+            if (!this.taken.contains(child)) {
+                out.add(child);
+                collectChildren(child, out);
+            }
+        }
     }
 }

loader/src/test/java/net/neoforged/fml/TopologicalSortTest.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.fml;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import com.google.common.graph.ElementOrder;
+import com.google.common.graph.GraphBuilder;
+import com.google.common.graph.MutableGraph;
+import java.util.Comparator;
+import net.neoforged.fml.loading.toposort.CyclePresentException;
+import net.neoforged.fml.loading.toposort.TopologicalSort;
+import org.junit.jupiter.api.Test;
+
+@SuppressWarnings("UnstableApiUsage") // guava graph is @Beta
+public class TopologicalSortTest {
+    @Test
+    public void testDefaultOrdering() {
+        MutableGraph<Integer> g = GraphBuilder.directed().nodeOrder(ElementOrder.insertion()).build();
+        g.putEdge(5, 0);
+        g.putEdge(5, 4);
+        g.putEdge(0, 1);
+        g.putEdge(1, 2);
+        g.putEdge(0, 4);
+        g.addNode(3);
+
+        var sorted = TopologicalSort.topologicalSort(g, Comparator.naturalOrder());
+        assertThat(sorted)
+                .containsExactly(5, 0, 1, 2, 3, 4);
+    }
+
+    @Test
+    public void testCycle() {
+        MutableGraph<Integer> g = GraphBuilder.directed().nodeOrder(ElementOrder.insertion()).build();
+        g.putEdge(0, 1);
+        g.putEdge(1, 3);
+        g.putEdge(3, 0);
+        g.addNode(2);
+
+        assertThatThrownBy(() -> TopologicalSort.topologicalSort(g, Comparator.naturalOrder()))
+                .isInstanceOf(CyclePresentException.class);
+    }
+}