IR KT-47840 fix-point solution for closure conversion

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeProjection
 import org.jetbrains.kotlin.ir.util.isLocal
+import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 
 data class Closure(val capturedValues: List<IrValueSymbol>, val capturedTypeParameters: List<IrTypeParameter>)
@@ -43,7 +44,6 @@ class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
     fun getClassClosure(declaration: IrClass) = getClosure(declaration)
 
     private fun getClosure(declaration: IrDeclaration): Closure {
-        closureBuilders.values.forEach { it.processed = false }
         return closureBuilders
             .getOrElse(declaration) { throw AssertionError("No closure builder for passed declaration ${ir2string(declaration)}.") }
             .buildClosure()
@@ -57,24 +57,63 @@ class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
         private val potentiallyCapturedTypeParameters = mutableSetOf<IrTypeParameter>()
         private val capturedTypeParameters = mutableSetOf<IrTypeParameter>()
 
-        var processed = false
+        private var closure: Closure? = null
 
         /*
-         * Node's closure = variables captured by the node +
-         *                  closure of all included nodes -
-         *                  variables declared in the node.
+         * This will solve a system of equations for each dependent closure:
+         *
+         *  closure[V] = captured(V) + { closure[U] | U <- included(V) } - declared(V)
+         *
          */
         fun buildClosure(): Closure {
-            includes.forEach { builder ->
-                if (!builder.processed) {
-                    builder.processed = true
-                    val subClosure = builder.buildClosure()
-                    subClosure.capturedValues.filterTo(capturedValues) { isExternal(it.owner) }
-                    subClosure.capturedTypeParameters.filterTo(capturedTypeParameters) { isExternal(it) }
+            closure?.let { return it }
+
+            val work = collectConnectedClosures()
+
+            do {
+                var changes = false
+                for (c in work) {
+                    if (c.updateFromIncluded()) {
+                        changes = true
+                    }
+                }
+            } while (changes)
+
+            for (c in work) {
+                c.closure = Closure(c.capturedValues.toList(), c.capturedTypeParameters.toList())
+            }
+
+            return closure
+                ?: throw AssertionError("Closure should have been built for ${owner.render()}")
+        }
+
+        private fun collectConnectedClosures(): List<ClosureBuilder> {
+            val connected = LinkedHashSet<ClosureBuilder>()
+            fun collectRec(current: ClosureBuilder) {
+                for (included in current.includes) {
+                    if (included.closure == null && connected.add(included)) {
+                        collectRec(included)
+                    }
                 }
             }
-            // TODO: We can save the closure and reuse it.
-            return Closure(capturedValues.toList(), capturedTypeParameters.toList())
+            connected.add(this)
+            collectRec(this)
+            return connected.toList().asReversed()
+        }
+
+        private fun updateFromIncluded(): Boolean {
+            if (closure != null)
+                throw AssertionError("Closure has already been built for ${owner.render()}")
+
+            val capturedValuesBefore = capturedValues.size
+            val capturedTypeParametersBefore = capturedTypeParameters.size
+            for (subClosure in includes) {
+                subClosure.capturedValues.filterTo(capturedValues) { isExternal(it.owner) }
+                subClosure.capturedTypeParameters.filterTo(capturedTypeParameters) { isExternal(it) }
+            }
+
+            return capturedValues.size != capturedValuesBefore ||
+                    capturedTypeParameters.size != capturedTypeParametersBefore
         }