commit 5fe8ac584bfb35813cfabe6db902b481397962d8
author Jim Sproch <jsproch@google.com> Fri Mar 18 16:28:28 2022 -0700
committer Jim Sproch <jsproch@google.com> Fri Mar 18 18:22:25 2022 -0700
tree 97659be2724c2b830d37fb46aea57897819dcf21
parent dc3325f34bb59fc7b9eee76e193f962f2cb15b1f

ComposeCompiler: Handle KtLabeledExpression when finding lambdas

Fixes: b/222979253
Change-Id: I8215d0e7c00f3f91ccf678503c9549d41cf35997

compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/SanityCheckCodegenTests.kt

diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/SanityCheckCodegenTests.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/SanityCheckCodegenTests.kt
index 237d15f..38c974d 100644
--- a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/SanityCheckCodegenTests.kt
+++ b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/SanityCheckCodegenTests.kt
@@ -57,4 +57,21 @@
         """
         )
     }
+
+    // Regression test for b/222979253
+    fun testLabeledLambda() = ensureSetup {
+        testCompile(
+            """
+                import androidx.compose.runtime.Composable
+
+                @Composable
+                fun test(): Unit {
+                    Box box@{}
+                }
+
+                @Composable
+                fun Box(content: @Composable () -> Unit) {}
+        """
+        )
+    }
 }

compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposableTargetChecker.kt

diff --git a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposableTargetChecker.kt b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposableTargetChecker.kt
index 34b23fb..3973744 100644
--- a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposableTargetChecker.kt
+++ b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposableTargetChecker.kt
@@ -50,6 +50,7 @@
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtLabeledExpression
 import org.jetbrains.kotlin.psi.KtLambdaArgument
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtProperty
@@ -349,9 +350,18 @@
         return PsiElementNode(element, bindingContext)
     }
 
-    private fun lambdaOrNull(element: PsiElement): KtFunctionLiteral? =
-        (element as? KtLambdaArgument)?.children?.firstOrNull()?.children
-            ?.firstOrNull() as KtFunctionLiteral?
+    private fun lambdaOrNull(element: PsiElement): KtFunctionLiteral? {
+        var container = (element as? KtLambdaArgument)?.children?.singleOrNull()
+        while (true) {
+            container = when (container) {
+                null -> return null
+                is KtLabeledExpression -> container.lastChild
+                is KtFunctionLiteral -> return container
+                is KtLambdaExpression -> container.children.single()
+                else -> throw Error("Unknown type: ${container.javaClass}")
+            }
+        }
+    }
 
     private fun descriptorToInferenceNode(
         descriptor: CallableDescriptor,