Merge pull request #3359 from square/fast_topo_graph

[swift] Improve performance of topological sort algorithm for partitions

Author: dnkoutso

wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/DirectedAcyclicGraph.kt

@@ -45,20 +45,40 @@ internal class DirectedAcyclicGraph<N>(
   }
 
   fun topologicalOrder(): List<N> {
-    val seen = LinkedHashSet<N>() // Insertion order is important to produce the final list!
-    val queue = ArrayDeque<N>().apply { addAll(seeds) }
+    val incomingEdges = mutableMapOf<N, Int>()
+    for (vertex in nodes) {
+      if (vertex !in incomingEdges) {
+        incomingEdges[vertex] = 0
+      }
+      for (edge in edges(vertex)) {
+        incomingEdges[edge] = (incomingEdges[edge] ?: 0) + 1
+      }
+    }
+
+    val queue = ArrayDeque<N>()
+    for ((vertex, edges) in incomingEdges) {
+      if (edges == 0) queue += vertex
+    }
+
+    val result = mutableListOf<N>()
+
     while (queue.isNotEmpty()) {
-      val currentName = queue.removeFirst()
-      val currentDependencies = edges(currentName).toList()
-      if (seen.containsAll(currentDependencies)) {
-        seen += currentName
-        queue += incomingEdges(currentName)
-      } else {
-        // All dependencies have not been seen so move to the back of the line to try again later.
-        queue += currentName
+      val vertex = queue.removeFirst()
+      result += vertex
+
+      for (edge in edges(vertex)) {
+        incomingEdges[edge] = (incomingEdges[edge] ?: 0) - 1
+        if (incomingEdges[edge] == 0) {
+          queue += edge
+        }
       }
     }
-    return seen.toList()
+
+    check(result.size == incomingEdges.size) {
+      "Graph contains a cycle, topological sort not possible!"
+    }
+
+    return result
   }
 
   fun transitiveNodes(node: N): Set<N> {

wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/PartitionedSchema.kt

@@ -34,12 +34,29 @@ internal class PartitionedSchema(
   )
 }
 
+internal fun computeDepths(modules: Map<String, Module>): Map<String, Int> {
+  val memo = mutableMapOf<String, Int>()
+
+  fun dfs(name: String): Int {
+    memo[name]?.let { return it }
+    val depth = 1 + (modules.getValue(name).dependencies.maxOfOrNull { dfs(it) } ?: 0)
+    memo[name] = depth
+    return depth
+  }
+
+  return modules.keys.associateWith { dfs(it) }
+}
+
 internal fun Schema.partition(modules: Map<String, Module>): PartitionedSchema {
   val moduleGraph = DirectedAcyclicGraph(modules.keys) { modules.getValue(it).dependencies }
 
   val errors = mutableListOf<String>()
   val partitions = mutableMapOf<String, Partition>()
-  for (moduleName in moduleGraph.topologicalOrder()) {
+  val topoGraph = moduleGraph.topologicalOrder()
+  val depths = computeDepths(modules)
+  val ordered = topoGraph.sortedBy { depths.getValue(it) }
+
+  for (moduleName in ordered) {
     val module = modules.getValue(moduleName)
 
     val upstreamTypes = buildMap {

wire-schema/src/jvmTest/kotlin/com/squareup/wire/schema/DirectedAcyclicGraphTest.kt

@@ -0,0 +1,116 @@
+/*
+ * Copyright (C) 2025 Square, Inc.
+ *
+ * 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
+ *
+ *      https://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.
+ */
+package com.squareup.wire.schema
+
+import assertk.assertThat
+import assertk.assertions.containsExactly
+import assertk.assertions.hasMessage
+import assertk.assertions.isEqualTo
+import kotlin.test.Test
+import kotlin.test.fail
+
+class DirectedAcyclicGraphTest {
+
+  @Test
+  fun singleNode() {
+    val nodes = listOf("A")
+    val graph = DirectedAcyclicGraph(nodes) { emptyList() }
+    val order = graph.topologicalOrder()
+    assertThat(order).containsExactly("A")
+  }
+
+  @Test
+  fun simpleLinearDag() {
+    // A -> B -> C
+    val nodes = listOf("A", "B", "C")
+    val edges = mapOf(
+      "A" to listOf("B"),
+      "B" to listOf("C"),
+      "C" to emptyList(),
+    )
+    val graph = DirectedAcyclicGraph(nodes) { edges[it].orEmpty() }
+    val order = graph.topologicalOrder()
+    assertThat(order).isEqualTo(listOf("A", "B", "C"))
+  }
+
+  @Test
+  fun simpleNonLinearDag() {
+    // A
+    // C -> B
+    val nodes = listOf("A", "B", "C")
+    val edges = mapOf(
+      "A" to emptyList(),
+      "B" to emptyList(),
+      "C" to listOf("B"),
+    )
+    val graph = DirectedAcyclicGraph(nodes) { edges[it].orEmpty() }
+    val order = graph.topologicalOrder()
+    assertThat(order).isEqualTo(listOf("A", "C", "B"))
+  }
+
+  @Test
+  fun branchingDag() {
+    //   A
+    //  / \
+    // B   C
+    //  \ /
+    //   D
+    val nodes = listOf("A", "B", "C", "D")
+    val edges = mapOf(
+      "A" to listOf("B", "C"),
+      "B" to listOf("D"),
+      "C" to listOf("D"),
+      "D" to emptyList(),
+    )
+    val graph = DirectedAcyclicGraph(nodes) { edges[it].orEmpty() }
+    val order = graph.topologicalOrder()
+    assertThat(order).isEqualTo(listOf("A", "B", "C", "D"))
+  }
+
+  @Test
+  fun multipleRoots() {
+    // A -> C
+    // B -> C
+    val nodes = listOf("A", "B", "C")
+    val edges = mapOf(
+      "A" to listOf("C"),
+      "B" to listOf("C"),
+      "C" to emptyList(),
+    )
+    val graph = DirectedAcyclicGraph(nodes) { edges[it].orEmpty() }
+    val order = graph.topologicalOrder()
+    assertThat(order).isEqualTo(listOf("A", "B", "C"))
+  }
+
+  @Test
+  fun cycleThrowsError() {
+    // A -> B -> C -> A
+    val nodes = listOf("A", "B", "C")
+    val edges = mapOf(
+      "A" to listOf("B"),
+      "B" to listOf("C"),
+      "C" to listOf("A"),
+    )
+    val graph = DirectedAcyclicGraph(nodes) { edges[it].orEmpty() }
+
+    try {
+      graph.topologicalOrder()
+      fail()
+    } catch (expected: IllegalStateException) {
+      assertThat(expected).hasMessage("Graph contains a cycle, topological sort not possible!")
+    }
+  }
+}