FIR IDE: expected type of lambdas and block expression

2021-08-19 22:16:32 -07:00
parent da3f2c2095
commit b7a99eca55
27 changed files with 192 additions and 22 deletions
@@ -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.FirFunction
-import org.jetbrains.kotlin.fir.declarations.FirImport
+import org.jetbrains.kotlin.fir.declarations.*
 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)
-            ?: getExpressionTypeByIfOrBooleanCondition(expression)
-            ?: getExpectedTypeOfFunctionParameter(expression)
-            ?: getExpectedTypeOfInfixFunctionParameter(expression)
-            ?: getExpectedTypeByVariableAssignment(expression)
-            ?: getExpectedTypeByPropertyDeclaration(expression)
-            ?: getExpectedTypeByFunctionExpressionBody(expression)
+        val unwrapped = expression.unwrap()
+        val expectedType = getExpectedTypeByReturnExpression(unwrapped)
+            ?: getExpressionTypeByIfOrBooleanCondition(unwrapped)
+            ?: getExpectedTypeByTypeCast(unwrapped)
+            ?: getExpectedTypeOfFunctionParameter(unwrapped)
+            ?: getExpectedTypeOfInfixFunctionParameter(unwrapped)
+            ?: 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 = argumentsList.parent as? KtCallExpression ?: return null
+        val callExpression =
+            (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)
    }

@@ -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.KtExpression
-import org.jetbrains.kotlin.psi.KtLabeledExpression
-import org.jetbrains.kotlin.psi.KtObjectLiteralExpression
+import org.jetbrains.kotlin.psi.*
+
+internal fun PsiElement.unwrap(): PsiElement {
+    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
 }

@@ -0,0 +1,8 @@
+val lam4 = fun(a: Int): String <caret>{
+    if (a < 5) return "5"
+
+    if (a > 0)
+        return "1"
+    else
+        return "2"
+}
@@ -0,0 +1,9 @@
+expression: {
+    if (a < 5) return "5"
+
+    if (a > 0)
+        return "1"
+    else
+        return "2"
+}
+expected type: null
@@ -1,5 +1,3 @@
-// FIX_ME: should work on non fully resolved calls
-
 fun x() {
    toCall(1, z = a<caret>v)
 }
@@ -1,2 +1,2 @@
 expression: av
-expected type: null
+expected type: kotlin/Boolean
@@ -0,0 +1,3 @@
+val x: () -> Unit = {<caret>
+    val (a, b) = listOf(1, 2)
+}
@@ -0,0 +1,4 @@
+expression: {
+    val (a, b) = listOf(1, 2)
+}
+expected type: kotlin/Function0<kotlin/Unit>
@@ -0,0 +1,4 @@
+val lam1 = <caret>{ a: Int ->
+    val b = 1
+    a + b
+}
@@ -0,0 +1,5 @@
+expression: { a: Int ->
+    val b = 1
+    a + b
+}
+expected type: kotlin/Function1<kotlin/Int, kotlin/Int>
@@ -0,0 +1,3 @@
+fun foo(p0 : Any) {
+    val s = p<caret>0 as? String
+}
@@ -0,0 +1,2 @@
+expression: p0
+expected type: kotlin/String?
@@ -0,0 +1,3 @@
+fun foo(p0 : Any) {
+    val s = p<caret>0 as String
+}
@@ -0,0 +1,2 @@
+expression: p0
+expected type: kotlin/String
@@ -0,0 +1 @@
+val baz = java.lang.Runnable <caret>{ /* SAM */ }
@@ -0,0 +1,2 @@
+expression: { /* SAM */ }
+expected type: java/lang/Runnable
@@ -0,0 +1,5 @@
+fun runRunnable(r: java.lang.Runnable) = r()
+
+fun foo() {
+    runRunnable {<caret> /* Argument */ }
+}
@@ -0,0 +1,2 @@
+expression: { /* Argument */ }
+expected type: java/lang/Runnable
@@ -0,0 +1,3 @@
+fun test1() {
+    val thread1 = Thread({<caret> println("hello1") })
+}
@@ -0,0 +1,2 @@
+expression: { println("hello1") }
+expected type: ft<java/lang/Runnable, java/lang/Runnable?>
@@ -0,0 +1,3 @@
+fun bar(): java.lang.Runnable {
+    return {<caret> /* Return */ }
+}
@@ -0,0 +1,2 @@
+expression: { /* Return */ }
+expected type: java/lang/Runnable
@@ -0,0 +1 @@
+val foo : java.lang.Runnable = {<caret> /* Variable */ }
@@ -0,0 +1,2 @@
+expression: { /* Variable */ }
+expected type: java/lang/Runnable
@@ -0,0 +1 @@
+val a = {<caret> /* Type Cast */ } as java.lang.Runnable
@@ -0,0 +1,2 @@
+expression: { /* Type Cast */ }
+expected type: java/lang/Runnable
@@ -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 {
				`@@ -0,0 +1 @@`
				`val baz = java.lang.Runnable <caret>{ /* SAM */ }`
				`@@ -0,0 +1 @@`
				`val foo : java.lang.Runnable = {<caret> /* Variable */ }`
				`@@ -0,0 +1 @@`
				`val a = {<caret> /* Type Cast */ } as java.lang.Runnable`