AA: fix expected type for lambda with explicit label

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10ExpressionTypeProvider.kt

@@ -161,7 +161,12 @@ class KtFe10ExpressionTypeProvider(
 
     override fun getExpectedType(expression: PsiElement): KtType? {
         val ktExpression = expression.getParentOfType<KtExpression>(false) ?: return null
-        val parentExpression = ktExpression.parent
+        val parentExpression = if (ktExpression.parent is KtLabeledExpression) {
+            // lambda -> labeled expression -> lambda argument (value argument)
+            ktExpression.parent.parent
+        } else {
+            ktExpression.parent
+        }
 
         // Unwrap specific expressions
         when (ktExpression) {

analysis/analysis-api-fe10/tests-gen/org/jetbrains/kotlin/analysis/api/fe10/test/cases/generated/cases/components/expressionTypeProvider/Fe10IdeNormalAnalysisSourceModuleExpectedExpressionTypeTestGenerated.java

@@ -238,6 +238,18 @@ public class Fe10IdeNormalAnalysisSourceModuleExpectedExpressionTypeTestGenerate
         runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithExplicitTypeFromProperty.kt");
     }
 
+    @Test
+    @TestMetadata("samWithReturnToExplicitLabel.kt")
+    public void testSamWithReturnToExplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToExplicitLabel.kt");
+    }
+
+    @Test
+    @TestMetadata("samWithReturnToImplicitLabel.kt")
+    public void testSamWithReturnToImplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToImplicitLabel.kt");
+    }
+
     @Test
     @TestMetadata("samWithTypeCast.kt")
     public void testSamWithTypeCast() throws Exception {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirExpressionTypeProvider.kt

@@ -155,7 +155,11 @@ internal class KtFirExpressionTypeProvider(
     }
 
     private fun PsiElement.getFunctionCallAsWithThisAsParameter(): KtCallWithArgument? {
-        val valueArgument = unwrapQualified<KtValueArgument> { valueArg, expr -> valueArg.getArgumentExpression() == expr } ?: return null
+        val valueArgument = unwrapQualified<KtValueArgument> { valueArg, expr ->
+            // If `valueArg` is [KtLambdaArgument], its [getArgumentExpression] could be labeled expression (e.g., l@{ ... }).
+            // That is not exactly `expr`, which would be [KtLambdaExpression]. So, we need [unwrap] here.
+            valueArg.getArgumentExpression()?.unwrap() == expr
+        } ?: return null
         val callExpression =
             (valueArgument.parent as? KtValueArgumentList)?.parent as? KtCallExpression
                 ?: valueArgument.parent as? KtCallExpression // KtLambdaArgument
@@ -271,7 +275,8 @@ private inline fun <reified R : Any> PsiElement.unwrapQualified(check: (R, PsiEl
 
 private val PsiElement.nonContainerParent: PsiElement?
     get() = when (val parent = parent) {
-        is KtContainerNode -> parent.parent
+        is KtContainerNode -> parent.nonContainerParent
+        is KtLabeledExpression -> parent.nonContainerParent
         else -> parent
     }
 

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/expressionTypeProvider/FirIdeDependentAnalysisSourceModuleExpectedExpressionTypeTestGenerated.java

@@ -238,6 +238,18 @@ public class FirIdeDependentAnalysisSourceModuleExpectedExpressionTypeTestGenera
         runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithExplicitTypeFromProperty.kt");
     }
 
+    @Test
+    @TestMetadata("samWithReturnToExplicitLabel.kt")
+    public void testSamWithReturnToExplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToExplicitLabel.kt");
+    }
+
+    @Test
+    @TestMetadata("samWithReturnToImplicitLabel.kt")
+    public void testSamWithReturnToImplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToImplicitLabel.kt");
+    }
+
     @Test
     @TestMetadata("samWithTypeCast.kt")
     public void testSamWithTypeCast() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/expressionTypeProvider/FirIdeNormalAnalysisSourceModuleExpectedExpressionTypeTestGenerated.java

@@ -238,6 +238,18 @@ public class FirIdeNormalAnalysisSourceModuleExpectedExpressionTypeTestGenerated
         runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithExplicitTypeFromProperty.kt");
     }
 
+    @Test
+    @TestMetadata("samWithReturnToExplicitLabel.kt")
+    public void testSamWithReturnToExplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToExplicitLabel.kt");
+    }
+
+    @Test
+    @TestMetadata("samWithReturnToImplicitLabel.kt")
+    public void testSamWithReturnToImplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToImplicitLabel.kt");
+    }
+
     @Test
     @TestMetadata("samWithTypeCast.kt")
     public void testSamWithTypeCast() throws Exception {

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/components/expressionTypeProvider/FirStandaloneNormalAnalysisSourceModuleExpectedExpressionTypeTestGenerated.java

@@ -238,6 +238,18 @@ public class FirStandaloneNormalAnalysisSourceModuleExpectedExpressionTypeTestGe
         runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithExplicitTypeFromProperty.kt");
     }
 
+    @Test
+    @TestMetadata("samWithReturnToExplicitLabel.kt")
+    public void testSamWithReturnToExplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToExplicitLabel.kt");
+    }
+
+    @Test
+    @TestMetadata("samWithReturnToImplicitLabel.kt")
+    public void testSamWithReturnToImplicitLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToImplicitLabel.kt");
+    }
+
     @Test
     @TestMetadata("samWithTypeCast.kt")
     public void testSamWithTypeCast() throws Exception {

analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToExplicitLabel.kt

@@ -0,0 +1 @@
+val baz = java.lang.Runnable l@{<caret> return@l }

analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToExplicitLabel.txt

@@ -0,0 +1,2 @@
+expression: { return@l }
+expected type: java/lang/Runnable

analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToImplicitLabel.kt

@@ -0,0 +1 @@
+val baz = java.lang.Runnable {<caret> return@Runnable }

analysis/analysis-api/testData/components/expressionTypeProvider/expectedExpressionType/samWithReturnToImplicitLabel.txt

@@ -0,0 +1,2 @@
+expression: { return@Runnable }
+expected type: java/lang/Runnable