kotlin-tooling-core: Implement closure util

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/util/reportSourceSetCommonizerDependencies.kt

@@ -10,12 +10,12 @@ import org.intellij.lang.annotations.RegExp
 import org.jetbrains.kotlin.commonizer.CommonizerTarget
 import org.jetbrains.kotlin.commonizer.SharedCommonizerTarget
 import org.jetbrains.kotlin.commonizer.parseCommonizerTarget
-import org.jetbrains.kotlin.commonizer.util.transitiveClosure
 import org.jetbrains.kotlin.gradle.BaseGradleIT
 import org.jetbrains.kotlin.gradle.BaseGradleIT.CompiledProject
 import org.jetbrains.kotlin.library.ToolingSingleFileKlibResolveStrategy
 import org.jetbrains.kotlin.library.commonizerTarget
 import org.jetbrains.kotlin.library.resolveSingleFileKlib
+import org.jetbrains.kotlin.tooling.core.singleClosure
 import java.io.File
 import javax.annotation.RegEx
 import kotlin.test.fail
@@ -50,7 +50,7 @@ data class SourceSetCommonizerDependencies(
         if (konanDataDir != null) {
             return file.startsWith(konanDataDir)
         }
-        return file.allParents.any { parentFile -> parentFile.name == ".konan" }
+        return file.parentsClosure.any { parentFile -> parentFile.name == ".konan" }
     }
 
     fun assertTargetOnAllDependencies(target: CommonizerTarget) = apply {
@@ -165,7 +165,7 @@ private fun inferCommonizerTargetOrNull(libraryFile: File): CommonizerTarget? =
     strategy = ToolingSingleFileKlibResolveStrategy
 ).commonizerTarget?.let(::parseCommonizerTarget)
 
-private val File.allParents: Set<File> get() = transitiveClosure(this) { listOfNotNull(parentFile) }
+private val File.parentsClosure: Set<File> get() = this.singleClosure { parentFile }
 
 private const val dollar = "\$"
 

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

@@ -0,0 +1,138 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+@file:Suppress("FunctionName")
+
+package org.jetbrains.kotlin.tooling.core
+
+/**
+ * General purpose implementation of a transitive closure
+ * - Recursion free
+ * - Predictable amount of allocations
+ * - Handles loops and self references gracefully
+ * @param edges: Producer function from one node to all its children. This implementation can handle loops and self references gracefully.
+ * @return Note: The order of the set is guaranteed to be bfs
+ */
+inline fun <reified T> T.closure(edges: (T) -> Iterable<T>): Set<T> =
+    closureTo(
+        destination = createResultSet(),
+        exclude = emptySet(),
+        dequeue = createDequeue(),
+        seed = this,
+        enqueueNextElements = { element -> addAll(edges(element)) }
+    )
+
+/**
+ * Similar to [closure], but will also include the receiver(seed) of this function into the final set
+ * @see closure
+ */
+inline fun <reified T> T.withClosure(edges: (T) -> Iterable<T>): Set<T> =
+    closureTo(
+        destination = createResultSet(this),
+        exclude = emptySet(),
+        dequeue = createDequeue(),
+        seed = this,
+        enqueueNextElements = { element -> addAll(edges(element)) }
+    )
+
+/**
+ * @see closure
+ * @receiver: Will not be included in the return set
+ */
+inline fun <reified T> Iterable<T>.closure(edges: (T) -> Iterable<T>): Set<T> {
+    val thisSet = this.toSet()
+    return closureTo(
+        destination = createResultSet(thisSet.size * 2),
+        exclude = thisSet,
+        dequeue = createDequeue(thisSet.size * 2), edges = edges
+    )
+}
+
+/**
+ * @see closure
+ * @receiver: Will be included in the return set
+ */
+inline fun <reified T> Iterable<T>.withClosure(edges: (T) -> Iterable<T>): Set<T> {
+    val size = this.count()
+    return closureTo(
+        destination = createResultSet(this, size * 2),
+        exclude = emptySet(),
+        dequeue = createDequeue(size * 2),
+        edges = edges
+    )
+}
+
+/**
+ * @see closure
+ * @receiver is not included in the return set
+ */
+inline fun <reified T : Any> T.singleClosure(next: (T) -> T?): Set<T> =
+    closureTo(createResultSet(), emptySet(), createDequeue(), this) { element -> next(element)?.let { add(it) } }
+
+/**
+ * @see closure
+ * @receiver is included in the return set
+ */
+inline fun <reified T : Any> T.withSingleClosure(next: (T) -> T?): Set<T> =
+    closureTo(createResultSet(this), emptySet(), createDequeue(), this) { element -> next(element)?.let { add(it) } }
+
+/* Implementation */
+
+private typealias Results<T> = MutableSet<T>
+private typealias Dequeue<T> = MutableList<T>
+
+@PublishedApi
+internal inline fun <reified T> Iterable<T>.closureTo(
+    destination: Results<T>, exclude: Set<T>, dequeue: Dequeue<T>, edges: (T) -> Iterable<T>
+): Set<T> {
+    forEach { seed ->
+        closureTo(
+            destination = destination,
+            exclude = exclude,
+            dequeue = dequeue,
+            seed = seed,
+            enqueueNextElements = { element -> addAll(edges(element)) })
+    }
+    return destination
+}
+
+@PublishedApi
+internal inline fun <T> closureTo(
+    destination: Results<T>,
+    exclude: Set<T>,
+    dequeue: Dequeue<T>,
+    seed: T,
+    enqueueNextElements: Dequeue<T>.(value: T) -> Unit
+): Set<T> {
+    dequeue.enqueueNextElements(seed)
+    while (dequeue.isNotEmpty()) {
+        val element = dequeue.removeAt(0)
+        if (element != seed && element !in exclude && destination.add(element)) {
+            dequeue.enqueueNextElements(element)
+        }
+    }
+    return destination
+}
+
+@PublishedApi
+internal inline fun <reified T> createDequeue(initialSize: Int = 16): MutableList<T> {
+    return if (KotlinVersion.CURRENT.isAtLeast(1, 4)) ArrayDeque(initialSize)
+    else ArrayList(initialSize)
+}
+
+@PublishedApi
+internal fun <T> createResultSet(initialSize: Int = 16): MutableSet<T> {
+    return LinkedHashSet(initialSize)
+}
+
+@PublishedApi
+internal fun <T> createResultSet(withValue: T, initialSize: Int = 16): MutableSet<T> {
+    return LinkedHashSet<T>(initialSize).also { it.add(withValue) }
+}
+
+@PublishedApi
+internal fun <T> createResultSet(withValues: Iterable<T>, initialSize: Int = 16): MutableSet<T> {
+    return LinkedHashSet<T>(initialSize).also { it.addAll(withValues) }
+}

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

@@ -0,0 +1,164 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.tooling.core
+
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class ClosureTest {
+
+    private class Node(val value: String, var parent: Node? = null, val children: MutableList<Node> = mutableListOf()) {
+        override fun toString(): String = value
+    }
+
+    @Test
+    fun `closure does not include root node`() {
+        val closure = Node("a", children = mutableListOf(Node("b"), Node("c"))).closure { it.children }
+        assertEquals(
+            listOf("b", "c"), closure.map { it.value },
+            "Expected closure to not include root node"
+        )
+    }
+
+    @Test
+    fun `withClosure does include root node`() {
+        val closure = Node("a", children = mutableListOf(Node("b"), Node("c"))).withClosure { it.children }
+        assertEquals(
+            listOf("a", "b", "c"), closure.map { it.value },
+            "Expected 'withClosure' to include root node"
+        )
+    }
+
+    @Test
+    fun `closure 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 = nodeA.closure { it.children }
+        assertEquals(
+            setOf(nodeB, nodeC, nodeD), closure,
+            "Expected transitiveClosure to be robust against loops and self references"
+        )
+    }
+
+    @Test
+    fun `withClosure 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 = nodeA.withClosure { it.children }
+        assertEquals(
+            setOf(nodeA, nodeB, nodeC, nodeD), closure,
+            "Expected transitiveClosure to be robust against loops and self references"
+        )
+    }
+
+    @Test
+    fun `closure 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 -> 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("b", "c", "d", "f", "g"), // <- a *is not* listed in closure!!!
+            listOf(a, e).closure<Node> { it.children }.map { it.value }
+        )
+    }
+
+    @Test
+    fun `withClosure 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 -> 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("a", "e", "b", "c", "d", "f", "g"),
+            listOf(a, e).withClosure<Node> { it.children }.map { it.value }
+        )
+    }
+
+    @Test
+    fun singleClosure() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+
+        c.parent = b
+        b.parent = a
+
+        assertEquals(
+            listOf("b", "a"), c.singleClosure { it.parent }.map { it.value },
+        )
+    }
+
+    @Test
+    fun withSingleClosure() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+
+        c.parent = b
+        b.parent = a
+
+        assertEquals(
+            listOf("c", "b", "a"), c.withSingleClosure { it.parent }.map { it.value },
+        )
+    }
+}

native/commonizer-api/src/org/jetbrains/kotlin/commonizer/util/transitiveClosure.kt

@@ -5,14 +5,7 @@
 
 package org.jetbrains.kotlin.commonizer.util
 
-/**
+
- * General purpose implementation of a transitive closure
- * - Recursion free
- * - Predictable amount of allocations
- * - Handles loops and self references gracefully
- * @param edges: Producer function from one node to all its children. This implementation can handle loops and self references gracefully.
- * @return Note: No guarantees given about the order ot this [Set]
- */
 public inline fun <reified T> transitiveClosure(seed: T, edges: T.() -> Iterable<T>): Set<T> {
     // Fast path when initial edges are empty
     val initialEdges = seed.edges()
@@ -32,7 +25,6 @@ public inline fun <reified T> transitiveClosure(seed: T, edges: T.() -> Iterable
     return results.toSet()
 }
 
-@OptIn(ExperimentalStdlibApi::class)
 @PublishedApi
 internal inline fun <reified T> deque(initialSize: Int): MutableList<T> {
     return if (KotlinVersion.CURRENT.isAtLeast(1, 4)) ArrayDeque(initialSize)