FIR IDE: expected type of lambdas and block expression

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

@@ -8,13 +8,12 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirLabel
 import org.jetbrains.kotlin.fir.FirPackageDirective
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirFunction
-import org.jetbrains.kotlin.fir.declarations.FirImport
 import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirNamedReference
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.constructFunctionalType
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -59,28 +58,55 @@ internal class KtFirExpressionTypeProvider(
     }
 
     override fun getExpectedType(expression: PsiElement): KtType? {
-        val expectedType = getExpectedTypeByReturnExpression(expression)
+        val unwrapped = expression.unwrap()
-            ?: getExpressionTypeByIfOrBooleanCondition(expression)
+        val expectedType = getExpectedTypeByReturnExpression(unwrapped)
-            ?: getExpectedTypeOfFunctionParameter(expression)
+            ?: getExpressionTypeByIfOrBooleanCondition(unwrapped)
-            ?: getExpectedTypeOfInfixFunctionParameter(expression)
+            ?: getExpectedTypeByTypeCast(unwrapped)
-            ?: getExpectedTypeByVariableAssignment(expression)
+            ?: getExpectedTypeOfFunctionParameter(unwrapped)
-            ?: getExpectedTypeByPropertyDeclaration(expression)
+            ?: getExpectedTypeOfInfixFunctionParameter(unwrapped)
-            ?: getExpectedTypeByFunctionExpressionBody(expression)
+            ?: getExpectedTypeByVariableAssignment(unwrapped)
+            ?: getExpectedTypeByPropertyDeclaration(unwrapped)
+            ?: getExpectedTypeByFunctionExpressionBody(unwrapped)
         return expectedType.takeIf { it !is KtClassErrorType }
     }
 
+    private fun getExpectedTypeByTypeCast(expression: PsiElement): KtType? {
+        val typeCastExpression =
+            expression.unwrapQualified<KtBinaryExpressionWithTypeRHS> { castExpr, expr -> castExpr.left == expr } ?: return null
+        with(analysisSession) {
+            return typeCastExpression.right?.getKtType()
+        }
+    }
+
     private fun getExpectedTypeOfFunctionParameter(expression: PsiElement): KtType? {
-        val (ktCallExpression, ktArgument) = expression.getFunctionCallAsWithThisAsParameter() ?: return null
+        val (ktCallExpression, argumentExpression) = expression.getFunctionCallAsWithThisAsParameter() ?: return null
         val firCall = ktCallExpression.getOrBuildFirSafe<FirFunctionCall>(firResolveState) ?: return null
+
+        val callee = (firCall.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol
+        if (callee?.fir?.origin == FirDeclarationOrigin.SamConstructor) {
+            return (callee.fir as FirSimpleFunction).returnTypeRef.coneType.asKtType()
+        }
+
         val arguments = firCall.argumentMapping ?: return null
-        val firParameterForExpression = arguments.entries.firstOrNull { (arg, _) -> arg.psi == ktArgument }?.value ?: return null
+        val firParameterForExpression =
+            arguments.entries.firstOrNull { (arg, _) ->
+                when (arg) {
+                    // TODO: better to utilize. See `createArgumentMapping` in [KtFirCallResolver]
+                    is FirLambdaArgumentExpression, is FirNamedArgumentExpression, is FirSpreadArgumentExpression ->
+                        arg.psi == argumentExpression.parent
+                    else ->
+                        arg.psi == argumentExpression
+                }
+            }?.value ?: return null
         return firParameterForExpression.returnTypeRef.coneType.asKtType()
     }
 
     private fun PsiElement.getFunctionCallAsWithThisAsParameter(): KtCallWithArgument? {
         val valueArgument = unwrapQualified<KtValueArgument> { valueArg, expr -> valueArg.getArgumentExpression() == expr } ?: return null
-        val argumentsList = valueArgument.parent as? KtValueArgumentList ?: return null
+        val callExpression =
-        val callExpression = argumentsList.parent as? KtCallExpression ?: return null
+            (valueArgument.parent as? KtValueArgumentList)?.parent as? KtCallExpression
+                ?: valueArgument.parent as? KtCallExpression // KtLambdaArgument
+                ?: return null
         val argumentExpression = valueArgument.getArgumentExpression() ?: return null
         return KtCallWithArgument(callExpression, argumentExpression)
     }
@@ -131,6 +157,11 @@ internal class KtFirExpressionTypeProvider(
         // Given: `fun f(): T = expression`
         // Expected type of `expression` is `T`
         val function = expression.unwrapQualified<KtFunction> { function, expr -> function.bodyExpression == expr } ?: return null
+        if (function.bodyBlockExpression != null) {
+            // Given `fun f(...): R { blockExpression }`, `{ blockExpression }` is mapped to the enclosing anonymous function,
+            // which may raise an exception if we attempt to retrieve, e.g., callable declaration from it.
+            return null
+        }
         return getReturnTypeForKtDeclaration(function)
     }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/firUtils.kt

@@ -4,6 +4,7 @@
  */
 package org.jetbrains.kotlin.idea.frontend.api.fir.utils
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
@@ -26,17 +27,22 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSimpleConstantVa
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtUnsupportedConstantValue
 import org.jetbrains.kotlin.idea.frontend.api.types.KtTypeNullability
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
-import org.jetbrains.kotlin.psi.KtAnnotatedExpression
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtExpression
+
-import org.jetbrains.kotlin.psi.KtLabeledExpression
+internal fun PsiElement.unwrap(): PsiElement {
-import org.jetbrains.kotlin.psi.KtObjectLiteralExpression
+    return when (this) {
+        is KtExpression -> this.unwrap()
+        else -> this
+    }
+}
 
 internal fun KtExpression.unwrap(): KtExpression {
     return when (this) {
         is KtLabeledExpression -> baseExpression?.unwrap()
         is KtAnnotatedExpression -> baseExpression?.unwrap()
         is KtObjectLiteralExpression -> objectDeclaration
-        else -> null
+        is KtFunctionLiteral -> (parent as? KtLambdaExpression)?.unwrap()
+        else -> this
     } ?: this
 }
 

idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBodyBlockExpression.kt

@@ -0,0 +1,8 @@
+val lam4 = fun(a: Int): String <caret>{
+    if (a < 5) return "5"
+
+    if (a > 0)
+        return "1"
+    else
+        return "2"
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBodyBlockExpression.txt

@@ -0,0 +1,9 @@
+expression: {
+    if (a < 5) return "5"
+
+    if (a > 0)
+        return "1"
+    else
+        return "2"
+}
+expected type: null

idea/idea-frontend-fir/testData/components/expectedExpressionType/functionNamedlParam.kt

@@ -1,5 +1,3 @@
-// FIX_ME: should work on non fully resolved calls
-
 fun x() {
     toCall(1, z = a<caret>v)
 }

idea/idea-frontend-fir/testData/components/expectedExpressionType/functionNamedlParam.txt

@@ -1,2 +1,2 @@
 expression: av
-expected type: null
+expected type: kotlin/Boolean

idea/idea-frontend-fir/testData/components/expectedExpressionType/lambdaWithExplicitTypeFromVariable.kt

@@ -0,0 +1,3 @@
+val x: () -> Unit = {<caret>
+    val (a, b) = listOf(1, 2)
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/lambdaWithExplicitTypeFromVariable.txt

@@ -0,0 +1,4 @@
+expression: {
+    val (a, b) = listOf(1, 2)
+}
+expected type: kotlin/Function0<kotlin/Unit>

idea/idea-frontend-fir/testData/components/expectedExpressionType/lambdaWithoutReturnNorExplicitType.kt

@@ -0,0 +1,4 @@
+val lam1 = <caret>{ a: Int ->
+    val b = 1
+    a + b
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/lambdaWithoutReturnNorExplicitType.txt

@@ -0,0 +1,5 @@
+expression: { a: Int ->
+    val b = 1
+    a + b
+}
+expected type: kotlin/Function1<kotlin/Int, kotlin/Int>

idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithSafeCast.kt

@@ -0,0 +1,3 @@
+fun foo(p0 : Any) {
+    val s = p<caret>0 as? String
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithSafeCast.txt

@@ -0,0 +1,2 @@
+expression: p0
+expected type: kotlin/String?

idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithTypeCast.kt

@@ -0,0 +1,3 @@
+fun foo(p0 : Any) {
+    val s = p<caret>0 as String
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithTypeCast.txt

@@ -0,0 +1,2 @@
+expression: p0
+expected type: kotlin/String

idea/idea-frontend-fir/testData/components/expectedExpressionType/sam.kt

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

idea/idea-frontend-fir/testData/components/expectedExpressionType/sam.txt

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

idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsArgument.kt

@@ -0,0 +1,5 @@
+fun runRunnable(r: java.lang.Runnable) = r()
+
+fun foo() {
+    runRunnable {<caret> /* Argument */ }
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsArgument.txt

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

idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsConstructorArgument.kt

@@ -0,0 +1,3 @@
+fun test1() {
+    val thread1 = Thread({<caret> println("hello1") })
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsConstructorArgument.txt

@@ -0,0 +1,2 @@
+expression: { println("hello1") }
+expected type: ft<java/lang/Runnable, java/lang/Runnable?>

idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsReturn.kt

@@ -0,0 +1,3 @@
+fun bar(): java.lang.Runnable {
+    return {<caret> /* Return */ }
+}

idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsReturn.txt

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

idea/idea-frontend-fir/testData/components/expectedExpressionType/samWithExplicitTypeFromProperty.kt

@@ -0,0 +1 @@
+val foo : java.lang.Runnable = {<caret> /* Variable */ }

idea/idea-frontend-fir/testData/components/expectedExpressionType/samWithExplicitTypeFromProperty.txt

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

idea/idea-frontend-fir/testData/components/expectedExpressionType/samWithTypeCast.kt

@@ -0,0 +1 @@
+val a = {<caret> /* Type Cast */ } as java.lang.Runnable

idea/idea-frontend-fir/testData/components/expectedExpressionType/samWithTypeCast.txt

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

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/ExpectedExpressionTypeTestGenerated.java

@@ -30,6 +30,12 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBody.kt");
     }
 
+    @Test
+    @TestMetadata("functionExpressionBodyBlockExpression.kt")
+    public void testFunctionExpressionBodyBlockExpression() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionExpressionBodyBlockExpression.kt");
+    }
+
     @Test
     @TestMetadata("functionExpressionBodyQualified.kt")
     public void testFunctionExpressionBodyQualified() throws Exception {
@@ -108,6 +114,18 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/infixFunctionParamQualified.kt");
     }
 
+    @Test
+    @TestMetadata("lambdaWithExplicitTypeFromVariable.kt")
+    public void testLambdaWithExplicitTypeFromVariable() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/lambdaWithExplicitTypeFromVariable.kt");
+    }
+
+    @Test
+    @TestMetadata("lambdaWithoutReturnNorExplicitType.kt")
+    public void testLambdaWithoutReturnNorExplicitType() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/lambdaWithoutReturnNorExplicitType.kt");
+    }
+
     @Test
     @TestMetadata("propertyDeclaration.kt")
     public void testPropertyDeclaration() throws Exception {
@@ -120,6 +138,18 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationQualified.kt");
     }
 
+    @Test
+    @TestMetadata("propertyDeclarationWithSafeCast.kt")
+    public void testPropertyDeclarationWithSafeCast() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithSafeCast.kt");
+    }
+
+    @Test
+    @TestMetadata("propertyDeclarationWithTypeCast.kt")
+    public void testPropertyDeclarationWithTypeCast() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/propertyDeclarationWithTypeCast.kt");
+    }
+
     @Test
     @TestMetadata("propertyDeclarationWithTypeFromRHS.kt")
     public void testPropertyDeclarationWithTypeFromRHS() throws Exception {
@@ -156,6 +186,42 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
         runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromLambda.kt");
     }
 
+    @Test
+    @TestMetadata("sam.kt")
+    public void testSam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/sam.kt");
+    }
+
+    @Test
+    @TestMetadata("samAsArgument.kt")
+    public void testSamAsArgument() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsArgument.kt");
+    }
+
+    @Test
+    @TestMetadata("samAsConstructorArgument.kt")
+    public void testSamAsConstructorArgument() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsConstructorArgument.kt");
+    }
+
+    @Test
+    @TestMetadata("samAsReturn.kt")
+    public void testSamAsReturn() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/samAsReturn.kt");
+    }
+
+    @Test
+    @TestMetadata("samWithExplicitTypeFromProperty.kt")
+    public void testSamWithExplicitTypeFromProperty() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/samWithExplicitTypeFromProperty.kt");
+    }
+
+    @Test
+    @TestMetadata("samWithTypeCast.kt")
+    public void testSamWithTypeCast() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/samWithTypeCast.kt");
+    }
+
     @Test
     @TestMetadata("variableAssignment.kt")
     public void testVariableAssignment() throws Exception {