[kotlin-tooling-core] closure.kt: Implement fast paths for empty closures

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

@@ -3,7 +3,7 @@
  * 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")
+@file:Suppress("FunctionName", "DeprecatedCallableAddReplaceWith")
 
 package org.jetbrains.kotlin.tooling.core
 
@@ -15,39 +15,71 @@ package org.jetbrains.kotlin.tooling.core
  * @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> =
+inline fun <reified T> T.closure(edges: (T) -> Iterable<T>): Set<T> {
-    closureTo(
+    val initialEdges = edges(this)
-        destination = createResultSet(),
+
-        exclude = emptySet(),
+    val dequeue = if (initialEdges is Collection) {
-        dequeue = createDequeue(),
+        if (initialEdges.isEmpty()) return emptySet()
-        seed = this,
+        createDequeue(initialEdges)
-        enqueueNextElements = { element -> addAll(edges(element)) }
+    } else createDequeueFromIterable(initialEdges)
-    )
+
+    val results = createResultSet<T>(dequeue.size)
+
+    while (dequeue.isNotEmpty()) {
+        val element = dequeue.removeAt(0)
+        if (element != this && results.add(element)) {
+            dequeue.addAll(edges(element))
+        }
+    }
+    return results
+}
 
 /**
  * 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> =
+inline fun <reified T> T.withClosure(edges: (T) -> Iterable<T>): Set<T> {
-    closureTo(
+    val initialEdges = edges(this)
-        destination = createResultSet(this),
+
-        exclude = emptySet(),
+    val dequeue = if (initialEdges is Collection) {
-        dequeue = createDequeue(),
+        if (initialEdges.isEmpty()) return setOf(this)
-        seed = this,
+        createDequeue(initialEdges)
-        enqueueNextElements = { element -> addAll(edges(element)) }
+    } else createDequeueFromIterable(initialEdges)
-    )
+
+    val results = createResultSet<T>(dequeue.size)
+    results.add(this)
+
+    while (dequeue.isNotEmpty()) {
+        val element = dequeue.removeAt(0)
+        if (results.add(element)) {
+            dequeue.addAll(edges(element))
+        }
+    }
+    return results
+}
 
 /**
  * @see closure
  * @receiver: Will not be included in the return set
  */
 inline fun <reified T> Iterable<T>.closure(edges: (T) -> Iterable<T>): Set<T> {
+    if (this is Collection && this.isEmpty()) return emptySet()
     val thisSet = this.toSet()
-    return closureTo(
+
-        destination = createResultSet(thisSet.size * 2),
+    val dequeue = createDequeue<T>()
-        exclude = thisSet,
+    thisSet.forEach { seed -> dequeue.addAll(edges(seed)) }
-        dequeue = createDequeue(thisSet.size * 2), edges = edges
+    if (dequeue.isEmpty()) return emptySet()
-    )
+
+    val results = createResultSet<T>()
+
+    while (dequeue.isNotEmpty()) {
+        val element = dequeue.removeAt(0)
+        if (element !in thisSet && results.add(element)) {
+            dequeue.addAll(edges(element))
+        }
+    }
+
+    return results
 }
 
 /**
@@ -55,56 +87,109 @@ inline fun <reified T> Iterable<T>.closure(edges: (T) -> Iterable<T>): Set<T> {
  * @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()
+    val dequeue = if (this is Collection) {
-    return closureTo(
+        if (this.isEmpty()) return emptySet()
-        destination = createResultSet(this, size * 2),
+        createDequeue(this)
-        exclude = emptySet(),
+    } else createDequeueFromIterable(this)
-        dequeue = createDequeue(size * 2),
+
-        edges = edges
+    val results = createResultSet<T>()
-    )
+
+    while (dequeue.isNotEmpty()) {
+        val element = dequeue.removeAt(0)
+        if (results.add(element)) {
+            dequeue.addAll(edges(element))
+        }
+    }
+
+    return results
 }
 
 /**
  * @see closure
  * @receiver is not included in the return set
  */
-inline fun <reified T : Any> T.linearClosure(next: (T) -> T?): Set<T> =
+inline fun <reified T : Any> T.linearClosure(next: (T) -> T?): Set<T> {
-    closureTo(createResultSet(), emptySet(), createDequeue(), this) { element -> next(element)?.let { add(it) } }
+    val initial = next(this) ?: return emptySet()
+    val results = createResultSet<T>()
+    var enqueued: T? = initial
+    while (enqueued != null) {
+        if (enqueued != this && results.add(enqueued)) {
+            enqueued = next(enqueued)
+        }
+    }
+
+    return results
+}
 
 /**
  * @see closure
  * @receiver is included in the return set
  */
-inline fun <reified T : Any> T.withLinearClosure(next: (T) -> T?): Set<T> =
+inline fun <reified T : Any> T.withLinearClosure(next: (T) -> T?): Set<T> {
-    closureTo(createResultSet(this), emptySet(), createDequeue(), this) { element -> next(element)?.let { add(it) } }
+    val initial = next(this) ?: return setOf(this)
+    val results = createResultSet<T>()
+    results.add(this)
 
-/* Implementation */
+    var enqueued: T? = initial
-
+    while (enqueued != null) {
-private typealias Results<T> = MutableSet<T>
+        if (results.add(enqueued)) {
-private typealias Dequeue<T> = MutableList<T>
+            enqueued = next(enqueued)
-
+        }
-@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
+
+    return results
 }
 
+@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 inline fun <reified T> createDequeue(elements: Collection<T>): MutableList<T> {
+    return if (KotlinVersion.CURRENT.isAtLeast(1, 4)) ArrayDeque(elements)
+    else ArrayList(elements)
+}
+
+@PublishedApi
+internal inline fun <reified T> createDequeueFromIterable(elements: Iterable<T>): MutableList<T> {
+    return createDequeue<T>().apply {
+        elements.forEach { element -> add(element) }
+    }
+}
+
+
+@PublishedApi
+internal fun <T> createResultSet(initialSize: Int = 16): MutableSet<T> {
+    return LinkedHashSet(initialSize)
+}
+
+//region Deprecations
+
+@Suppress("unused")
+@Deprecated("Scheduled for removal in 1.8", level = DeprecationLevel.ERROR)
+@PublishedApi
+internal fun <T> createResultSet(withValue: T, initialSize: Int = 16): MutableSet<T> {
+    return LinkedHashSet<T>(initialSize).also { it.add(withValue) }
+}
+
+@Suppress("unused")
+@Deprecated("Scheduled for removal in 1.8", level = DeprecationLevel.ERROR)
+@PublishedApi
+internal fun <T> createResultSet(withValues: Iterable<T>, initialSize: Int = 16): MutableSet<T> {
+    return LinkedHashSet<T>(initialSize).also { it.addAll(withValues) }
+}
+
+@Suppress("unused")
+@Deprecated("Scheduled for removal in 1.8", level = DeprecationLevel.ERROR)
 @PublishedApi
 internal inline fun <T> closureTo(
-    destination: Results<T>,
+    destination: MutableSet<T>,
     exclude: Set<T>,
-    dequeue: Dequeue<T>,
+    dequeue: MutableList<T>,
     seed: T,
-    enqueueNextElements: Dequeue<T>.(value: T) -> Unit
+    enqueueNextElements: MutableList<T>.(value: T) -> Unit
 ): Set<T> {
     dequeue.enqueueNextElements(seed)
     while (dequeue.isNotEmpty()) {
@@ -116,23 +201,4 @@ internal inline fun <T> closureTo(
     return destination
 }
 
-@PublishedApi
+//endregion
-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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.tooling.core
 
 import org.junit.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertSame
 
 class ClosureTest {
 
@@ -14,6 +15,18 @@ class ClosureTest {
         override fun toString(): String = value
     }
 
+    /* 'Children' is explicitly not implementing Collection */
+    private class IterableNode(val value: String, val children: Children = Children(mutableListOf())) {
+
+        /* Does not implement Collection */
+        class Children(private val list: MutableList<IterableNode>) : Iterable<IterableNode> {
+            override fun iterator(): Iterator<IterableNode> = list.iterator()
+            fun add(node: IterableNode) = list.add(node)
+        }
+
+        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 }
@@ -57,6 +70,32 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closure handles loop and self references - iterable node`() {
+        val nodeA = IterableNode("a")
+        val nodeB = IterableNode("b")
+        val nodeC = IterableNode("c")
+        val nodeD = IterableNode("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")
@@ -82,6 +121,47 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosure handles loop and self references - iterable node`() {
+        val nodeA = IterableNode("a")
+        val nodeB = IterableNode("b")
+        val nodeC = IterableNode("c")
+        val nodeD = IterableNode("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 with empty nodes`() {
+        assertSame(
+            emptySet(), Node("").closure { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
+    @Test
+    fun `closure with only self reference`() {
+        val node = Node("a")
+        node.children.add(node)
+        assertEquals(emptySet(), node.closure { it.children })
+    }
+
     @Test
     fun `closure on List`() {
         val a = Node("a")
@@ -92,22 +172,66 @@ class ClosureTest {
         val f = Node("f")
         val g = Node("g")
 
-        // a -> b, c -> d
+        // a -> (b, c)
+        // c -> (d)
         a.children.add(b)
         a.children.add(c)
         c.children.add(d)
 
-        // e -> f, g, a
+        // 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("b", "c", "f", "g", "d"), // <- a *is not* listed in closure!!!
             listOf(a, e).closure<Node> { it.children }.map { it.value }
         )
     }
 
+    @Test
+    fun `closure on List - iterable node`() {
+        val a = IterableNode("a")
+        val b = IterableNode("b")
+        val c = IterableNode("c")
+        val d = IterableNode("d")
+        val e = IterableNode("e")
+        val f = IterableNode("f")
+        val g = IterableNode("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("b", "c", "f", "g", "d"), // <- a *is not* listed in closure!!!
+            listOf(a, e).closure<IterableNode> { it.children }.map { it.value }
+        )
+    }
+
+    @Test
+    fun `closure on empty list`() {
+        assertSame(
+            emptySet(), listOf<Node>().closure<Node> { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
+    @Test
+    fun `closure - on list - no edges`() {
+        assertSame(
+            emptySet(), listOf(Node("a"), Node("b")).closure<Node> { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
     @Test
     fun `withClosure on List`() {
         val a = Node("a")
@@ -118,22 +242,38 @@ class ClosureTest {
         val f = Node("f")
         val g = Node("g")
 
-        // a -> b, c -> d
+        // a -> (b, c)
+        // c -> (d)
         a.children.add(b)
         a.children.add(c)
         c.children.add(d)
 
-        // e -> f, g, a
+        // 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", "b", "c", "f", "g", "d"),
             listOf(a, e).withClosure<Node> { it.children }.map { it.value }
         )
     }
 
+    @Test
+    fun `withClosure on emptyList`() {
+        assertSame(
+            emptySet(), listOf<Node>().withClosure<Node> { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
+    @Test
+    fun `withClosure with no further nodes`() {
+        assertEquals(
+            listOf("a", "b"), listOf(Node("a"), Node("b")).withClosure<Node> { it.children }.map { it.value }
+        )
+    }
+
     @Test
     fun linearClosure() {
         val a = Node("a")
@@ -148,6 +288,14 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `linearClosure on empty`() {
+        assertSame(
+            emptySet(), Node("").linearClosure { it.parent },
+            "Expected no Set being allocated on empty linearClosure"
+        )
+    }
+
     @Test
     fun withLinearClosure() {
         val a = Node("a")
@@ -161,4 +309,11 @@ class ClosureTest {
             listOf("c", "b", "a"), c.withLinearClosure { it.parent }.map { it.value },
         )
     }
+
+    @Test
+    fun `withLinearClosure on empty`() {
+        assertEquals(
+            listOf("a"), Node("a").withLinearClosure { it.parent }.map { it.value },
+        )
+    }
 }