FIR IDE: add expected type calculation for function calls

does not work in case of incomplete call expression :(
2020-12-13 21:17:11 +01:00
parent 7be8d69870
commit 2101816f03
9 changed files with 94 additions and 1 deletions
@@ -32,6 +32,7 @@ dependencies {
    testCompile(projectTests(":idea:idea-test-framework"))
    testCompile(project(":kotlin-test:kotlin-test-junit"))
    testCompile(commonDep("junit:junit"))
+    testCompile(projectTests(":idea:idea-frontend-independent"))

    compile(intellijPluginDep("java"))
 }
@@ -6,12 +6,19 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components

 import com.intellij.psi.PsiElement
+import com.jetbrains.rd.util.firstOrNull
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.expressions.FirCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.argumentMapping
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.types.ConeTypeCheckerContext
 import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
 import org.jetbrains.kotlin.idea.frontend.api.components.KtBuiltinTypes
@@ -23,6 +30,7 @@ import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import kotlin.reflect.KClass

 internal class KtFirTypeProvider(
@@ -41,6 +49,23 @@ internal class KtFirTypeProvider(
    override fun getExpectedType(expression: PsiElement): KtType? =
        getExpectedTypeByReturnExpression(expression)
            ?: getExpressionTypeByIfOrBooleanCondition(expression)
+            ?: getExpectedTypeOfFunctionParameter(expression)
+
+    private fun getExpectedTypeOfFunctionParameter(expression: PsiElement): KtType? {
+        val (ktCallExpression, ktArgument) = expression.getFunctionCallAsWithThisAsParameter() ?: return null
+        val firCall = ktCallExpression.getOrBuildFirSafe<FirFunctionCall>(firResolveState) ?: return null
+        val arguments = firCall.argumentMapping ?: return null
+        val firParameterForExpression = arguments.entries.firstOrNull { (arg, _) -> arg.psi == ktArgument }?.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 argumentExpression = valueArgument.getArgumentExpression() ?: return null
+        return KtCallWithArgument(callExpression, argumentExpression)
+    }

    private fun getExpectedTypeByReturnExpression(expression: PsiElement): KtType? {
        val returnParent = expression.getReturnExpressionWithThisType() ?: return null
@@ -99,6 +124,8 @@ internal class KtFirTypeProvider(
    )
 }

+private data class KtCallWithArgument(val call: KtCallExpression, val argument: KtExpression)
+
 private inline fun <reified R : Any> PsiElement.unwrapQualified(check: (R, PsiElement) -> Boolean): R? {
    val parent = nonContainerParent
    return when {
@@ -0,0 +1,7 @@
+fun x() {
+    toCall(1, "", a<caret>v)
+}
+
+fun toCall(x: Int, y: String, lambda: (Int) -> String): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/Function1<kotlin/Int, kotlin/String>
@@ -0,0 +1,9 @@
+// FIX_ME: should work on non fully resolved calls
+
+fun x() {
+    toCall(1, z = a<caret>v)
+}
+
+fun toCall(x: Int, y: String, z: Boolean): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/Boolean
@@ -0,0 +1,9 @@
+// FIX_ME: should work on non fully resolved calls
+
+fun x() {
+    toCall(a<caret>b, 12)
+}
+
+fun <T> toCall(x: T, y: T): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/Int
@@ -0,0 +1,7 @@
+fun x() {
+    toCall(1, a<caret>v, true)
+}
+
+fun toCall(x: Int, y: String, z: Boolean): Char = 'a'
+
+// EXPECTED_TYPE: kotlin/String
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
+import org.jetbrains.kotlin.test.uitls.IgnoreTests
 import java.io.File

 abstract class AbstractExpectedExpressionTypeTest : KotlinLightCodeInsightFixtureTestCase() {
@@ -37,7 +38,9 @@ abstract class AbstractExpectedExpressionTypeTest : KotlinLightCodeInsightFixtur
            }
        }

-        KotlinTestUtils.assertEqualsToFile(File(path), testDataFile.getTextWithActualType(actualExpectedTypeText))
+        IgnoreTests.runTestWithFixMeSupport(testDataFile.toPath()) {
+            KotlinTestUtils.assertEqualsToFile(File(path), testDataFile.getTextWithActualType(actualExpectedTypeText))
+        }
    }

    private fun File.getTextWithActualType(actualType: String?) : String {
@@ -28,6 +28,26 @@ public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpress
        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/expectedExpressionType"), Pattern.compile("^(.+)\\.kt$"), null, true);
    }

+    @TestMetadata("functionLambdaParam.kt")
+    public void testFunctionLambdaParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionLambdaParam.kt");
+    }
+
+    @TestMetadata("functionNamedlParam.kt")
+    public void testFunctionNamedlParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionNamedlParam.kt");
+    }
+
+    @TestMetadata("functionParamWithTypeParam.kt")
+    public void testFunctionParamWithTypeParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionParamWithTypeParam.kt");
+    }
+
+    @TestMetadata("functionPositionalParam.kt")
+    public void testFunctionPositionalParam() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/functionPositionalParam.kt");
+    }
+
    @TestMetadata("ifCondition.kt")
    public void testIfCondition() throws Exception {
        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/ifCondition.kt");
@@ -27,6 +27,15 @@ object IgnoreTests {
        )
    }

+    fun runTestWithFixMeSupport(testFile: Path, test: () -> Unit) {
+        runTestIfEnabledByDirective(
+            testFile,
+            EnableOrDisableTestDirective.Disable(DIRECTIVES.FIX_ME),
+            additionalFilesExtensions = emptyList(),
+            test = test
+        )
+    }
+
    fun runTestIfNotDisabledByFileDirective(
        testFile: Path,
        disableTestDirective: String,
@@ -128,5 +137,6 @@ object IgnoreTests {
    object DIRECTIVES {
        const val FIR_COMPARISON = "// FIR_COMPARISON"
        const val IGNORE_FIR = "// IGNORE_FIR"
+        const val FIX_ME = "// FIX_ME: "
    }
 }