FIR IDE: introduce expected type provider for return expressions

2020-12-12 17:44:45 +01:00
parent 299f36183c
commit 835577383b
10 changed files with 152 additions and 5 deletions
@@ -89,6 +89,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.AbstractFileSt
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.AbstractSessionsInvalidationTest
 import org.jetbrains.kotlin.idea.fir.low.level.api.trackers.AbstractProjectWideOutOfBlockKotlinModificationTrackerTest
 import org.jetbrains.kotlin.idea.folding.AbstractKotlinFoldingTest
+import org.jetbrains.kotlin.idea.frontend.api.components.AbstractExpectedExpressionTypeTest
 import org.jetbrains.kotlin.idea.frontend.api.components.AbstractReturnExpressionTargetTest
 import org.jetbrains.kotlin.idea.frontend.api.fir.AbstractResolveCallTest
 import org.jetbrains.kotlin.idea.frontend.api.scopes.AbstractMemberScopeByFqNameTest
@@ -1027,6 +1028,10 @@ fun main(args: Array<String>) {
            testClass<AbstractReturnExpressionTargetTest> {
                model("components/returnExpressionTarget")
            }
+
+            testClass<AbstractExpectedExpressionTypeTest> {
+                model("components/expectedExpressionType")
+            }
        }

        testGroup("idea/idea-frontend-fir/idea-fir-low-level-api/tests", "idea/testData") {
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.idea.frontend.api

+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.frontend.api.calls.KtCall
 import org.jetbrains.kotlin.idea.frontend.api.components.*
@@ -66,7 +67,7 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali

    infix fun KtType.isSubTypeOf(superType: KtType): Boolean = typeProvider.isSubTypeOf(this, superType)

-    fun KtExpression.getExpectedType(): KtType? = typeProvider.getExpectedType(this)
+    fun PsiElement.getExpectedType(): KtType? = typeProvider.getExpectedType(this)

    fun KtType.isBuiltInFunctionalType(): Boolean = typeProvider.isBuiltinFunctionalType(this)

@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.idea.frontend.api.components

+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
@@ -19,5 +20,5 @@ abstract class KtTypeProvider : KtAnalysisSessionComponent() {
    //TODO get rid of
    abstract fun isBuiltinFunctionalType(type: KtType): Boolean

-    abstract fun getExpectedType(expression: KtExpression): KtType?
+    abstract fun getExpectedType(expression: PsiElement): KtType?
 }
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.idea.frontend.api.fir.components

+import com.intellij.psi.PsiElement
 import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirExpression
@@ -21,6 +22,7 @@ import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtQualifiedExpression
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
@@ -38,15 +40,24 @@ internal class KtFirTypeProvider(
        expression.getOrBuildFirOfType<FirExpression>(firResolveState).typeRef.coneType.asKtType()
    }

-    override fun getExpectedType(expression: KtExpression): KtType? =
+    override fun getExpectedType(expression: PsiElement): KtType? =
        getExpectedTypeByReturnExpression(expression)

-    private fun getExpectedTypeByReturnExpression(expression: KtExpression): KtType? {
-        val returnParent = expression.parentOfType<KtReturnExpression>() ?: return null
+    private fun getExpectedTypeByReturnExpression(expression: PsiElement): KtType? {
+        val returnParent = expression.getReturnExpressionWithThisType() ?: return null
        val targetSymbol = with(analysisSession) { returnParent.getReturnTargetSymbol() } ?: return null
        return targetSymbol.type
    }

+    private fun PsiElement.getReturnExpressionWithThisType(): KtReturnExpression? {
+        val parent = parent
+        return when {
+            parent is KtReturnExpression && parent.returnedExpression == this -> parent
+            parent is KtQualifiedExpression && parent.selectorExpression == this -> parent.getReturnExpressionWithThisType()
+            else -> null
+        }
+    }
+
    override fun isEqualTo(first: KtType, second: KtType): Boolean = withValidityAssertion {
        second.assertIsValid()
        check(first is KtFirType)
@@ -0,0 +1,5 @@
+fun foo(): Int {
+    return <caret>a
+}
+
+// EXPECTED_TYPE: kotlin/Int
@@ -0,0 +1,5 @@
+fun foo(): Int {
+    return x<caret>fd.a
+}
+
+// EXPECTED_TYPE: null
@@ -0,0 +1,5 @@
+fun foo(): Int {
+    return x.<caret>a
+}
+
+// EXPECTED_TYPE: kotlin/Int
@@ -0,0 +1,10 @@
+fun x(): Int {
+    receiveLambda {
+        return@receiveLambda <caret>fd
+    }
+    return 2
+}
+
+fun receiveLambda(x: () -> Int){}
+
+// EXPECTED_TYPE: kotlin/Int
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api.components
+
+import com.intellij.openapi.util.io.FileUtil
+import com.intellij.psi.PsiComment
+import com.intellij.psi.util.parentOfType
+import junit.framework.Assert
+import org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
+import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFile
+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 java.io.File
+
+abstract class AbstractExpectedExpressionTypeTest : KotlinLightCodeInsightFixtureTestCase() {
+    override fun isFirPlugin() = true
+
+    protected fun doTest(path: String) {
+        val testDataFile = File(path)
+        val ktFile = myFixture.configureByText(testDataFile.name, FileUtil.loadFile(testDataFile)) as KtFile
+
+        val expressionAtCaret = ktFile.findElementAt(myFixture.caretOffset)?.parentOfType<KtExpression>()
+            ?: error("No element was found at caret or no <caret> is present in the test file")
+
+        val actualExpectedTypeText: String? = executeOnPooledThreadInReadAction {
+            analyze(ktFile) {
+                expressionAtCaret.getExpectedType()?.asStringForDebugging()
+            }
+        }
+
+        KotlinTestUtils.assertEqualsToFile(File(path), testDataFile.getTextWithActualType(actualExpectedTypeText))
+    }
+
+    private fun File.getTextWithActualType(actualType: String?) : String {
+        val text = FileUtil.loadFile(this)
+        val textWithoutTypeDirective = text.split('\n')
+            .filterNot { it.startsWith(EXPECTED_TYPE_TEXT_DIRECTIVE) }
+            .joinToString(separator = "\n")
+        return "$textWithoutTypeDirective\n$EXPECTED_TYPE_TEXT_DIRECTIVE $actualType"
+    }
+
+    companion object {
+        private const val EXPECTED_TYPE_TEXT_DIRECTIVE = "// EXPECTED_TYPE:"
+    }
+}
@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api.components;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("idea/idea-frontend-fir/testData/components/expectedExpressionType")
+@TestDataPath("$PROJECT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+public class ExpectedExpressionTypeTestGenerated extends AbstractExpectedExpressionTypeTest {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    public void testAllFilesPresentInExpectedExpressionType() throws Exception {
+        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/expectedExpressionType"), Pattern.compile("^(.+)\\.kt$"), null, true);
+    }
+
+    @TestMetadata("returnFromFunction.kt")
+    public void testReturnFromFunction() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromFunction.kt");
+    }
+
+    @TestMetadata("returnFromFunctionQualifiedReceiver.kt")
+    public void testReturnFromFunctionQualifiedReceiver() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromFunctionQualifiedReceiver.kt");
+    }
+
+    @TestMetadata("returnFromFunctionQualifiedSelector.kt")
+    public void testReturnFromFunctionQualifiedSelector() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromFunctionQualifiedSelector.kt");
+    }
+
+    @TestMetadata("returnFromLambda.kt")
+    public void testReturnFromLambda() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expectedExpressionType/returnFromLambda.kt");
+    }
+}