From 700432e1114a4a6cdd0525f31897e77a549c4205 Mon Sep 17 00:00:00 2001 From: Timofey Solonin Date: Wed, 14 Feb 2024 15:06:33 +0100 Subject: [PATCH] Add a closure function to return a graph in explicit bfs order ^KT-65540 --- .../jetbrains/kotlin/tooling/core/closure.kt | 30 ++++ .../kotlin/tooling/core/ClosureTest.kt | 146 ++++++++++++++++++ 2 files changed, 176 insertions(+) diff --git a/libraries/tools/kotlin-tooling-core/src/main/kotlin/org/jetbrains/kotlin/tooling/core/closure.kt b/libraries/tools/kotlin-tooling-core/src/main/kotlin/org/jetbrains/kotlin/tooling/core/closure.kt index 23c95c780d4..aadf2216e25 100644 --- a/libraries/tools/kotlin-tooling-core/src/main/kotlin/org/jetbrains/kotlin/tooling/core/closure.kt +++ b/libraries/tools/kotlin-tooling-core/src/main/kotlin/org/jetbrains/kotlin/tooling/core/closure.kt @@ -140,6 +140,36 @@ inline fun T.withLinearClosure(next: (T) -> T?): Set { return results } +/** + * Similar to [closure], but the result is returned in explicit BFS depth order starting with the receiver + * @see closure + */ +inline fun Iterable.withClosureGroupingByDistance(edges: (T) -> Iterable): List> { + val dequeue = if (this is Collection) { + if (this.isEmpty()) return emptyList() + createDequeue(this) + } else createDequeueFromIterable(this) + + val allElements = createResultSet(dequeue.size) + val results = mutableListOf>() + + while (dequeue.isNotEmpty()) { + val levelElements = createResultSet(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 createDequeue(initialSize: Int = 16): MutableList { return if (KotlinVersion.CURRENT.isAtLeast(1, 4)) ArrayDeque(initialSize) diff --git a/libraries/tools/kotlin-tooling-core/src/test/kotlin/org/jetbrains/kotlin/tooling/core/ClosureTest.kt b/libraries/tools/kotlin-tooling-core/src/test/kotlin/org/jetbrains/kotlin/tooling/core/ClosureTest.kt index 9d23dd2201c..5fa0a96d428 100644 --- a/libraries/tools/kotlin-tooling-core/src/test/kotlin/org/jetbrains/kotlin/tooling/core/ClosureTest.kt +++ b/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 { it.children }.map { it.map { it.value } } + ) + } + + + @Test + fun `withClosureGroupingByDistance on emptyList`() { + assertSame( + emptyList(), listOf().withClosureGroupingByDistance { 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 { 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 } }, + ) + } + }