Add a closure function to return a graph in explicit bfs order

^KT-65540

libraries/tools/kotlin-tooling-core/src/main/kotlin/org/jetbrains/kotlin/tooling/core/closure.kt

@@ -140,6 +140,36 @@ inline fun <reified T : Any> T.withLinearClosure(next: (T) -> T?): Set<T> {
     return results
 }
 
+/**
+ * Similar to [closure], but the result is returned in explicit BFS depth order starting with the receiver
+ * @see closure
+ */
+inline fun <reified T> Iterable<T>.withClosureGroupingByDistance(edges: (T) -> Iterable<T>): List<Set<T>> {
+    val dequeue = if (this is Collection) {
+        if (this.isEmpty()) return emptyList()
+        createDequeue(this)
+    } else createDequeueFromIterable(this)
+
+    val allElements = createResultSet<T>(dequeue.size)
+    val results = mutableListOf<MutableSet<T>>()
+
+    while (dequeue.isNotEmpty()) {
+        val levelElements = createResultSet<T>(dequeue.size)
+        var levelSize = dequeue.size
+        while (levelSize != 0) {
+            levelSize -= 1
+            val element = dequeue.removeAt(0)
+            if (allElements.add(element) && levelElements.add(element)) {
+                dequeue.addAll(edges(element))
+            }
+        }
+        if (levelElements.isNotEmpty()) {
+            results.add(levelElements)
+        }
+    }
+    return results
+}
+
 @PublishedApi
 internal inline fun <reified T> createDequeue(initialSize: Int = 16): MutableList<T> {
     return if (KotlinVersion.CURRENT.isAtLeast(1, 4)) ArrayDeque(initialSize)

libraries/tools/kotlin-tooling-core/src/test/kotlin/org/jetbrains/kotlin/tooling/core/ClosureTest.kt

@@ -346,4 +346,150 @@ class ClosureTest {
             listOf("a"), Node("a").withLinearClosure { it.parent }.map { it.value },
         )
     }
+
+    @Test
+    fun `withClosureGroupingByDistance on List`() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+        val d = Node("d")
+        val e = Node("e")
+        val f = Node("f")
+        val g = Node("g")
+
+        // a -> (b, c)
+        // c -> (d)
+        a.children.add(b)
+        a.children.add(c)
+        c.children.add(d)
+
+        // e -> (f, g, a)
+        e.children.add(f)
+        e.children.add(g)
+        e.children.add(a) // <- cycle back to a!
+
+        assertEquals(
+            listOf(
+                listOf("a", "e"),
+                listOf("b", "c", "f", "g"),
+                listOf("d"),
+            ),
+            listOf(a, e).withClosureGroupingByDistance<Node> { it.children }.map { it.map { it.value } }
+        )
+    }
+
+
+    @Test
+    fun `withClosureGroupingByDistance on emptyList`() {
+        assertSame(
+            emptyList(), listOf<Node>().withClosureGroupingByDistance<Node> { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
+    @Test
+    fun `withClosureGroupingByDistance with no further nodes`() {
+        assertEquals(
+            listOf(listOf("a", "b")),
+            listOf(Node("a"), Node("b")).withClosureGroupingByDistance<Node> { it.children }.map { it.map { it.value } }
+        )
+    }
+
+    @Test
+    fun `withClosureGroupingByDistance handles loop and self references`() {
+        val nodeA = Node("a")
+        val nodeB = Node("b")
+        val nodeC = Node("c")
+        val nodeD = Node("d")
+
+        // a -> b -> c -> d
+        nodeA.children.add(nodeB)
+        nodeB.children.add(nodeC)
+        nodeC.children.add(nodeD)
+
+        // add self reference to b
+        nodeB.children.add(nodeB)
+
+        // add loop from c -> a
+        nodeC.children.add(nodeA)
+
+        val closure = listOf(nodeA).withClosureGroupingByDistance { it.children }.map { it.map { it } }
+        assertEquals(
+            listOf(
+                listOf(nodeA),
+                listOf(nodeB),
+                listOf(nodeC),
+                listOf(nodeD)
+            ), closure,
+            "Expected withClosureGroupingByDistance to be robust against loops and self references"
+        )
+    }
+
+    @Test
+    fun `withClosureGroupingByDistance handles a simple self reference`() {
+        val nodeA = Node("a")
+        nodeA.children.add(nodeA)
+
+        assertEquals(
+            listOf(
+                listOf("a")
+            ),
+            listOf(nodeA).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
+        )
+    }
+
+    @Test
+    fun `withClosureGroupingByDistance handles a simple cyclic reference`() {
+        val nodeA = Node("a")
+        val nodeB = Node("b")
+        nodeA.children.add(nodeB)
+        nodeB.children.add(nodeA)
+
+        assertEquals(
+            listOf(
+                listOf("a"),
+                listOf("b"),
+            ),
+            listOf(nodeA).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
+        )
+        assertEquals(
+            listOf(
+                listOf("a", "b"),
+            ),
+            listOf(nodeA, nodeB).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
+        )
+    }
+
+    @Test
+    fun `withClosureGroupingByDistance handles a loop`() {
+        val nodeA = Node("a")
+        val nodeB = Node("b")
+        val nodeC = Node("c")
+        nodeA.children.add(nodeB)
+        nodeB.children.add(nodeC)
+        nodeC.children.add(nodeA)
+
+        assertEquals(
+            listOf(
+                listOf("a"),
+                listOf("b"),
+                listOf("c"),
+            ),
+            listOf(nodeA).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
+        )
+        assertEquals(
+            listOf(
+                listOf("a", "b"),
+                listOf("c"),
+            ),
+            listOf(nodeA, nodeB).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
+        )
+        assertEquals(
+            listOf(
+                listOf("a", "b", "c"),
+            ),
+            listOf(nodeA, nodeB, nodeC).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
+        )
+    }
+
 }