DurableFunctionKeyTransformer: Fix non-composable functions stopping …

…the 'realizeKeyMetaAnnotations' transformation ^KT-75159 fixed
JetBrains · Feb 12, 2025 · bbae865 · bbae865
1 parent 9058a9e
commit bbae865
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 5 deletions.
diff --git a/...vmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionKeyMetaAnnotationsTests.kt b/...vmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionKeyMetaAnnotationsTests.kt
@@ -104,6 +104,26 @@ class FunctionKeyMetaAnnotationsTests(useFir: Boolean) : AbstractCodegenTest(use
         """.trimIndent().replace("%", "$")
     )
 
+    @Test
+    fun testCapturingComposableLambda_entryPoint(): Unit = verifyGoldenBytecodeRender(
+        """
+         import androidx.compose.runtime.*
+         
+         fun runApplication(child: @Composable () -> Unit) {
+             /* Pretend to be an entry point */
+         }
+         
+         fun Foo() {
+             var state = 255
+             runApplication {
+                 println("%state")
+             }
+         }
+
+        """.trimIndent().replace("%", "$")
+    )
+
+
     private fun verifyGoldenBytecodeRender(@Language("kotlin") source: String) {
         val files = compileToClassFiles(source, "Test.kt")
         val rendered = renderAnnotatedDeclarations(files)

diff --git a/...ctionKeyMetaAnnotationsTests/testCapturingComposableLambda_entryPoint[useFir = false].txt b/...ctionKeyMetaAnnotationsTests/testCapturingComposableLambda_entryPoint[useFir = false].txt
@@ -0,0 +1,32 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+fun runApplication(child: @Composable () -> Unit) {
+    /* Pretend to be an entry point */
+}
+
+fun Foo() {
+    var state = 255
+    runApplication {
+        println("$state")
+    }
+}
+
+
+//
+// Transformed IR
+// ------------------------------------------
+
+TestKt {
+    runApplication (Lkotlin/jvm/functions/Function2;)V null
+    Foo ()V null
+}
+
+TestKt$Foo$1 {
+    <init> (Lkotlin/jvm/internal/Ref$IntRef;)V null
+    invoke (Landroidx/compose/runtime/Composer;I)V [key=2066003522, startOffset=180, endOffset=213]
+    invoke (Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object; null
+}
diff --git a/...nctionKeyMetaAnnotationsTests/testCapturingComposableLambda_entryPoint[useFir = true].txt b/...nctionKeyMetaAnnotationsTests/testCapturingComposableLambda_entryPoint[useFir = true].txt
@@ -0,0 +1,32 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+fun runApplication(child: @Composable () -> Unit) {
+    /* Pretend to be an entry point */
+}
+
+fun Foo() {
+    var state = 255
+    runApplication {
+        println("$state")
+    }
+}
+
+
+//
+// Transformed IR
+// ------------------------------------------
+
+TestKt {
+    runApplication (Lkotlin/jvm/functions/Function2;)V null
+    Foo ()V null
+}
+
+TestKt$Foo$1 {
+    <init> (Lkotlin/jvm/internal/Ref$IntRef;)V null
+    invoke (Landroidx/compose/runtime/Composer;I)V [key=2066003522, startOffset=180, endOffset=213]
+    invoke (Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object; null
+}
diff --git a/...main/java/androidx/compose/compiler/plugins/kotlin/lower/DurableFunctionKeyTransformer.kt b/...main/java/androidx/compose/compiler/plugins/kotlin/lower/DurableFunctionKeyTransformer.kt
@@ -101,17 +101,18 @@ class DurableFunctionKeyTransformer(
     fun realizeKeyMetaAnnotations(moduleFragment: IrModuleFragment) {
         moduleFragment.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitSimpleFunction(declaration: IrSimpleFunction): IrStatement {
-                val functionKey = context.irTrace[DURABLE_FUNCTION_KEY, declaration] ?: return declaration
-                if (!declaration.hasComposableAnnotation()) return declaration
-                if (declaration.hasAnnotation(ComposeClassIds.FunctionKeyMeta)) return declaration
+                run transform@{
+                    val functionKey = context.irTrace[DURABLE_FUNCTION_KEY, declaration] ?: return@transform
+                    if (!declaration.hasComposableAnnotation()) return@transform
+                    if (declaration.hasAnnotation(ComposeClassIds.FunctionKeyMeta)) return@transform
+                    declaration.annotations += irKeyMetaAnnotation(functionKey)
+                }
 
-                declaration.annotations += irKeyMetaAnnotation(functionKey)
                 return super.visitSimpleFunction(declaration)
             }
         })
     }
 
-
     private val keyMetaAnnotation =
         getTopLevelClassOrNull(ComposeClassIds.FunctionKeyMeta)