IR KT-47840 fix-point solution for closure conversion

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -5606,6 +5606,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/closures/kt47120a.kt");
         }
 
+        @Test
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @Test
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {

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
         }
 
 

compiler/testData/codegen/box/closures/kt47840.kt

@@ -0,0 +1,16 @@
+fun box(): String {
+    var t = ""
+    fun foo(x: String) {
+        fun bar() {
+            fun a() {
+                foo("")
+                bar()
+            }
+            t = x
+        }
+
+        bar()
+    }
+    foo("OK")
+    return t
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -5546,6 +5546,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/closures/kt47120a.kt");
         }
 
+        @Test
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @Test
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -5606,6 +5606,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/closures/kt47120a.kt");
         }
 
+        @Test
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @Test
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -4836,6 +4836,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/closures/kt47120a.kt");
         }
 
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {
             runTest("compiler/testData/codegen/box/closures/kt5589.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -3776,6 +3776,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
             runTest("compiler/testData/codegen/box/closures/kt4137.kt");
         }
 
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {
             runTest("compiler/testData/codegen/box/closures/kt5589.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -3776,6 +3776,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/closures/kt4137.kt");
         }
 
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {
             runTest("compiler/testData/codegen/box/closures/kt5589.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -3776,6 +3776,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/closures/kt4137.kt");
         }
 
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {
             runTest("compiler/testData/codegen/box/closures/kt5589.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -2657,6 +2657,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
             runTest("compiler/testData/codegen/box/closures/kt4137.kt");
         }
 
+        @TestMetadata("kt47840.kt")
+        public void testKt47840() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/kt47840.kt");
+        }
+
         @TestMetadata("kt5589.kt")
         public void testKt5589() throws Exception {
             runTest("compiler/testData/codegen/box/closures/kt5589.kt");