[kotlin-tooling-core] Implement 'closure' sequence APIs

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", "DeprecatedCallableAddReplaceWith")
+@file:Suppress("FunctionName", "DeprecatedCallableAddReplaceWith", "DuplicatedCode")
 
 package org.jetbrains.kotlin.tooling.core
 
@@ -34,6 +34,38 @@ inline fun <reified T> T.closure(edges: (T) -> Iterable<T>): Set<T> {
     return results
 }
 
+/**
+ * Same as [closure], but returns a lazy sequence instead
+ */
+inline fun <reified T> T.closureSequence(crossinline edges: (T) -> Iterable<T>): Sequence<T> {
+    val initialEdges = edges(this)
+
+    val resolveDequeue = if (initialEdges is Collection && initialEdges.isEmpty()) return emptySequence()
+    else createDequeue<T>()
+
+    return sequence {
+        val results = HashSet<T>()
+        results.add(this@closureSequence)
+
+        initialEdges.forEach { initialEdge ->
+            if (results.add(initialEdge)) {
+                yield(initialEdge)
+                resolveDequeue.add(initialEdge)
+            }
+        }
+
+        while (resolveDequeue.isNotEmpty()) {
+            edges(resolveDequeue.removeAt(0)).forEach { edge ->
+                if (results.add(edge)) {
+                    yield(edge)
+                    resolveDequeue.add(edge)
+                }
+            }
+        }
+    }
+}
+
+
 /**
  * Similar to [closure], but will also include the receiver(seed) of this function into the final set
  * @see closure
@@ -58,6 +90,39 @@ inline fun <reified T> T.withClosure(edges: (T) -> Iterable<T>): Set<T> {
     return results
 }
 
+/**
+ * Similar to [closureSequence], but will also include the receiver(seed) of this function into the final set
+ * @see closure
+ */
+inline fun <reified T> T.withClosureSequence(crossinline edges: (T) -> Iterable<T>): Sequence<T> {
+    val initialEdges = edges(this)
+
+    val resolveDequeue = if (initialEdges is Collection && initialEdges.isEmpty()) return sequenceOf(this)
+    else createDequeue<T>()
+
+    return sequence {
+        val results = HashSet<T>()
+        yield(this@withClosureSequence)
+        results.add(this@withClosureSequence)
+
+        initialEdges.forEach { initialEdge ->
+            if (results.add(initialEdge)) {
+                yield(initialEdge)
+                resolveDequeue.add(initialEdge)
+            }
+        }
+
+        while (resolveDequeue.isNotEmpty()) {
+            edges(resolveDequeue.removeAt(0)).forEach { edge ->
+                if (results.add(edge)) {
+                    yield(edge)
+                    resolveDequeue.add(edge)
+                }
+            }
+        }
+    }
+}
+
 /**
  * @see closure
  * @receiver: Will not be included in the return set
@@ -82,6 +147,37 @@ inline fun <reified T> Iterable<T>.closure(edges: (T) -> Iterable<T>): Set<T> {
     return results
 }
 
+/**
+ * @see closure
+ * @receiver: Will not be included in the return set
+ */
+inline fun <reified T> Iterable<T>.closureSequence(crossinline edges: (T) -> Iterable<T>): Sequence<T> {
+    if (this is Collection && this.isEmpty()) return emptySequence()
+    val thisSet = if (this is Set<T>) this else this.toSet()
+
+    val results = HashSet<T>(thisSet)
+    val resolveQueue = createDequeue<T>(thisSet)
+
+    return sequence {
+        while (resolveQueue.isNotEmpty()) {
+            edges(resolveQueue.removeAt(0)).forEach { edge ->
+                if (results.add(edge)) {
+                    yield(edge)
+                    resolveQueue.add(edge)
+                }
+            }
+        }
+    }
+}
+
+/**
+ * @see closure
+ * @receiver Initial sequence of elements will not be included in the return squence
+ */
+inline fun <reified T> Sequence<T>.closureSequence(crossinline edges: (T) -> Iterable<T>): Sequence<T> {
+    return toSet().closureSequence(edges)
+}
+
 /**
  * @see closure
  * @receiver: Will be included in the return set
@@ -104,6 +200,42 @@ inline fun <reified T> Iterable<T>.withClosure(edges: (T) -> Iterable<T>): Set<T
     return results
 }
 
+/**
+ * @see withClosure
+ * @receiver: Will be included in the return set
+ */
+inline fun <reified T> Iterable<T>.withClosureSequence(crossinline edges: (T) -> Iterable<T>): Sequence<T> {
+    return asSequence().withClosureSequence(edges)
+}
+
+/**
+ * @see closure
+ * @receiver: Will be included in the return set
+ */
+inline fun <reified T> Sequence<T>.withClosureSequence(crossinline edges: (T) -> Iterable<T>): Sequence<T> {
+    val results = HashSet<T>()
+    val resolveQueue = createDequeue<T>()
+
+    return sequence {
+        this@withClosureSequence.forEach { element ->
+            if (results.add(element)) {
+                yield(element)
+                resolveQueue.add(element)
+            }
+        }
+
+        while (resolveQueue.isNotEmpty()) {
+            val seed = resolveQueue.removeAt(0)
+            edges(seed).forEach { resolvedEdge ->
+                if (results.add(resolvedEdge)) {
+                    yield(resolvedEdge)
+                    resolveQueue.add(resolvedEdge)
+                }
+            }
+        }
+    }
+}
+
 /**
  * @see closure
  * @receiver is not included in the return set
@@ -121,6 +253,25 @@ inline fun <reified T : Any> T.linearClosure(next: (T) -> T?): Set<T> {
     return results
 }
 
+/**
+ * @see closure
+ * @receiver is not included in the return set
+ */
+inline fun <reified T : Any> T.linearClosureSequence(crossinline next: (T) -> T?): Sequence<T> {
+    val initial = next(this) ?: return emptySequence()
+    val results = HashSet<T>()
+
+    return sequence {
+        var enqueued: T? = initial
+        while (enqueued != null) {
+            if (enqueued != this@linearClosureSequence && results.add(enqueued)) {
+                yield(enqueued)
+                enqueued = next(enqueued)
+            } else break
+        }
+    }
+}
+
 /**
  * @see closure
  * @receiver is included in the return set
@@ -140,6 +291,27 @@ inline fun <reified T : Any> T.withLinearClosure(next: (T) -> T?): Set<T> {
     return results
 }
 
+/**
+ * @see closure
+ * @receiver is included in the return set
+ */
+inline fun <reified T : Any> T.withLinearClosureSequence(crossinline next: (T) -> T?): Sequence<T> {
+    val initial = next(this) ?: return sequenceOf(this)
+    val results = HashSet<T>()
+    return sequence {
+        results.add(this@withLinearClosureSequence)
+        yield(this@withLinearClosureSequence)
+
+        var enqueued: T? = initial
+        while (enqueued != null) {
+            if (results.add(enqueued)) {
+                yield(enqueued)
+                enqueued = next(enqueued)
+            } else break
+        }
+    }
+}
+
 /**
  * Similar to [closure], but the result is returned in explicit BFS depth order starting with the receiver
  * @see closure
@@ -219,7 +391,7 @@ internal inline fun <T> closureTo(
     exclude: Set<T>,
     dequeue: MutableList<T>,
     seed: T,
-    enqueueNextElements: MutableList<T>.(value: T) -> Unit
+    enqueueNextElements: MutableList<T>.(value: T) -> Unit,
 ): Set<T> {
     dequeue.enqueueNextElements(seed)
     while (dequeue.isNotEmpty()) {

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

@@ -36,6 +36,15 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence does not include root node`() {
+        val closure = Node("a", children = mutableListOf(Node("b"), Node("c"))).closureSequence { it.children }
+        assertEquals(
+            listOf("b", "c"), closure.map { it.value }.toList(),
+            "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 }
@@ -45,6 +54,15 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosureSequence does include root node`() {
+        val closure = Node("a", children = mutableListOf(Node("b"), Node("c"))).withClosureSequence { it.children }
+        assertEquals(
+            listOf("a", "b", "c"), closure.map { it.value }.toList(),
+            "Expected 'withClosure' to include root node"
+        )
+    }
+
     @Test
     fun `closure handles loop and self references`() {
         val nodeA = Node("a")
@@ -70,6 +88,31 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence 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.closureSequence { it.children }
+        assertEquals(
+            setOf(nodeB, nodeC, nodeD), closure.toSet(),
+            "Expected transitiveClosure to be robust against loops and self references"
+        )
+    }
+
     @Test
     fun `closure handles loop and self references - iterable node`() {
         val nodeA = IterableNode("a")
@@ -95,6 +138,30 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence 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.closureSequence { it.children }
+        assertEquals(
+            setOf(nodeB, nodeC, nodeD), closure.toSet(),
+            "Expected transitiveClosure to be robust against loops and self references"
+        )
+    }
 
     @Test
     fun `withClosure handles loop and self references`() {
@@ -121,6 +188,31 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosureSequence 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.withClosureSequence { it.children }
+        assertEquals(
+            setOf(nodeA, nodeB, nodeC, nodeD), closure.toSet(),
+            "Expected transitiveClosure to be robust against loops and self references"
+        )
+    }
+
     @Test
     fun `withClosure handles loop and self references - iterable node`() {
         val nodeA = IterableNode("a")
@@ -146,6 +238,30 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosureSequence 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.withClosureSequence { it.children }
+        assertEquals(
+            setOf(nodeA, nodeB, nodeC, nodeD), closure.toSet(),
+            "Expected transitiveClosure to be robust against loops and self references"
+        )
+    }
 
     @Test
     fun `closure with empty nodes`() {
@@ -155,6 +271,14 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence with empty nodes`() {
+        assertSame(
+            emptySequence(), Node("").closureSequence { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
     @Test
     fun `closure with only self reference`() {
         val node = Node("a")
@@ -162,6 +286,13 @@ class ClosureTest {
         assertEquals(emptySet(), node.closure { it.children })
     }
 
+    @Test
+    fun `closureSequence with only self reference`() {
+        val node = Node("a")
+        node.children.add(node)
+        assertEquals(emptyList(), node.closureSequence { it.children }.toList())
+    }
+
     @Test
     fun `closure on List`() {
         val a = Node("a")
@@ -189,6 +320,33 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence 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("b", "c", "f", "g", "d"), // <- a *is not* listed in closure!!!
+            listOf(a, e).closureSequence<Node> { it.children }.map { it.value }.toList()
+        )
+    }
+
     @Test
     fun `closure on List - iterable node`() {
         val a = IterableNode("a")
@@ -216,6 +374,33 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence 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).closureSequence<IterableNode> { it.children }.map { it.value }.toList()
+        )
+    }
+
     @Test
     fun `closure on empty list`() {
         assertSame(
@@ -224,6 +409,15 @@ class ClosureTest {
         )
     }
 
+
+    @Test
+    fun `closureSequence on empty list`() {
+        assertSame(
+            emptySequence(), listOf<Node>().closureSequence<Node> { it.children },
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
     @Test
     fun `closure - on list - no edges`() {
         assertSame(
@@ -232,6 +426,14 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `closureSequence - on list - no edges`() {
+        assertEquals(
+            emptyList(), listOf(Node("a"), Node("b")).closureSequence<Node> { it.children }.toList(),
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
     @Test
     fun `withClosure on List`() {
         val a = Node("a")
@@ -259,6 +461,33 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosureSequence 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("a", "e", "b", "c", "f", "g", "d"),
+            listOf(a, e).withClosureSequence<Node> { it.children }.map { it.value }.toList()
+        )
+    }
+
     @Test
     fun `withClosure on emptyList`() {
         assertSame(
@@ -267,6 +496,14 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosureSequence on emptyList`() {
+        assertSame(
+            emptySet(), listOf<Node>().withClosureSequence<Node> { it.children }.toSet(),
+            "Expected no Set being allocated on empty closure"
+        )
+    }
+
     @Test
     fun `withClosure with no further nodes`() {
         assertEquals(
@@ -274,6 +511,13 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withClosureSequence with no further nodes`() {
+        assertEquals(
+            listOf("a", "b"), listOf(Node("a"), Node("b")).withClosureSequence<Node> { it.children }.map { it.value }.toList()
+        )
+    }
+
     @Test
     fun linearClosure() {
         val a = Node("a")
@@ -288,6 +532,20 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun linearClosureSequence() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+
+        c.parent = b
+        b.parent = a
+
+        assertEquals(
+            listOf("b", "a"), c.linearClosureSequence { it.parent }.map { it.value }.toList(),
+        )
+    }
+
     @Test
     fun `linearClosure - loop`() {
         val a = Node("a")
@@ -303,6 +561,21 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `linearClosureSequence - loop`() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+
+        c.parent = b
+        b.parent = a
+        a.parent = c
+
+        assertEquals(
+            listOf("b", "a"), c.linearClosureSequence { it.parent }.map { it.value }.toList(),
+        )
+    }
+
     @Test
     fun `linearClosure on empty`() {
         assertSame(
@@ -311,6 +584,14 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `linearClosureSequence on empty`() {
+        assertSame(
+            emptySet(), Node("").linearClosureSequence { it.parent }.toSet(),
+            "Expected no Set being allocated on empty linearClosure"
+        )
+    }
+
     @Test
     fun withLinearClosure() {
         val a = Node("a")
@@ -325,6 +606,20 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun withLinearClosureSequence() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+
+        c.parent = b
+        b.parent = a
+
+        assertEquals(
+            listOf("c", "b", "a"), c.withLinearClosureSequence { it.parent }.map { it.value }.toList(),
+        )
+    }
+
     @Test
     fun `withLinearClosure - loop`() {
         val a = Node("a")
@@ -340,6 +635,21 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withLinearClosureSequence - loop`() {
+        val a = Node("a")
+        val b = Node("b")
+        val c = Node("c")
+
+        c.parent = b
+        b.parent = a
+        a.parent = c
+
+        assertEquals(
+            listOf("c", "b", "a"), c.withLinearClosureSequence { it.parent }.map { it.value }.toList(),
+        )
+    }
+
     @Test
     fun `withLinearClosure on empty`() {
         assertEquals(
@@ -347,6 +657,13 @@ class ClosureTest {
         )
     }
 
+    @Test
+    fun `withLinearClosureSequence on empty`() {
+        assertEquals(
+            listOf("a"), Node("a").withLinearClosureSequence { it.parent }.map { it.value }.toList(),
+        )
+    }
+
     @Test
     fun `withClosureGroupingByDistance on List`() {
         val a = Node("a")
@@ -491,5 +808,4 @@ class ClosureTest {
             listOf(nodeA, nodeB, nodeC).withClosureGroupingByDistance { it.children }.map { it.map { it.value } },
         )
     }
-
 }