FIR IDE: fix CCE in getKtExpressionType on assignment target

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt

@@ -90,6 +90,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.AbstractSessionsInva
 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.AbstractHLExpressionTypeTest
 import org.jetbrains.kotlin.idea.frontend.api.components.AbstractOverriddenDeclarationProviderTest
 import org.jetbrains.kotlin.idea.frontend.api.components.AbstractReturnExpressionTargetTest
 import org.jetbrains.kotlin.idea.frontend.api.fir.AbstractResolveCallTest
@@ -1046,6 +1047,10 @@ fun main(args: Array<String>) {
             testClass<AbstractOverriddenDeclarationProviderTest> {
                 model("components/overridenDeclarations")
             }
+
+            testClass<AbstractHLExpressionTypeTest> {
+                model("components/expressionType")
+            }
         }
 
         testGroup("idea/idea-frontend-fir/idea-fir-low-level-api/tests", "idea/testData") {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FirElementsRecorder.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
@@ -54,12 +55,16 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
     }
 
     override fun visitElement(element: FirElement, data: MutableMap<KtElement, FirElement>) {
-        (element.realPsi as? KtElement)?.let { psi ->
+        cacheElement(element, data)
-            cache(psi, element, data)
-        }
         element.acceptChildren(this, data)
     }
 
+    override fun visitVariableAssignment(variableAssignment: FirVariableAssignment, data: MutableMap<KtElement, FirElement>) {
+        cacheElement(variableAssignment.lValue, data) // FirReference is not cached by default
+        visitElement(variableAssignment, data)
+    }
+
+
     //@formatter:off
     override fun visitReference(reference: FirReference, data: MutableMap<KtElement, FirElement>) {}
     override fun visitControlFlowGraphReference(controlFlowGraphReference: FirControlFlowGraphReference, data: MutableMap<KtElement, FirElement>) {}
@@ -85,6 +90,12 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
         userTypeRef.acceptChildren(this, data)
     }
 
+    private fun cacheElement(element: FirElement, cache: MutableMap<KtElement, FirElement>) {
+        (element.realPsi as? KtElement)?.let { psi ->
+            cache(psi, element, cache)
+        }
+    }
+
     companion object {
         @OptIn(ExperimentalStdlibApi::class)
         fun recordElementsFrom(firDeclaration: FirDeclaration, recorder: FirElementsRecorder): Map<KtElement, FirElement> =

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

@@ -11,7 +11,14 @@ 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.references.FirErrorNamedReference
+import org.jetbrains.kotlin.fir.references.FirNamedReference
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedNameError
+import org.jetbrains.kotlin.fir.types.ConeClassErrorType
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 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
@@ -20,6 +27,7 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtTypedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
 import org.jetbrains.kotlin.psi.*
 
 internal class KtFirExpressionTypeProvider(
@@ -32,7 +40,23 @@ internal class KtFirExpressionTypeProvider(
     }
 
     override fun getKtExpressionType(expression: KtExpression): KtType = withValidityAssertion {
-        expression.getOrBuildFirOfType<FirExpression>(firResolveState).typeRef.coneType.asKtType()
+        val coneType = when (val fir = expression.getOrBuildFir(firResolveState)) {
+            is FirExpression -> fir.typeRef.coneType
+            is FirNamedReference -> fir.getReferencedElementType()
+            else -> error("Unexpected ${fir::class}")
+        }
+        coneType.asKtType()
+    }
+
+    private fun FirNamedReference.getReferencedElementType(): ConeKotlinType {
+        val symbols = when (this) {
+            is FirResolvedNamedReference -> listOf(resolvedSymbol)
+            is FirErrorNamedReference -> FirReferenceResolveHelper.getFirSymbolsByErrorNamedReference(this)
+            else -> error("Unexpected ${this::class}")
+        }
+        val firCallableDeclaration = symbols.singleOrNull()?.fir as? FirCallableDeclaration<*>
+        return firCallableDeclaration?.returnTypeRef?.coneType
+            ?: ConeClassErrorType(ConeUnresolvedNameError(name))
     }
 
     override fun getExpectedType(expression: PsiElement): KtType? =

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeWrongNumberOfTypeArgumen
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.getCandidateSymbols
@@ -174,12 +175,19 @@ internal object FirReferenceResolveHelper {
             is FirReturnExpression -> getSymbolsByReturnExpression(expression, fir, symbolBuilder)
             is FirErrorNamedReference -> getSymbolsByErrorNamedReference(fir, symbolBuilder)
             is FirVariableAssignment -> getSymbolsByVariableAssignment(fir, session, symbolBuilder)
+            is FirResolvedNamedReference -> getSymbolByResolvedNameReference(fir, session, symbolBuilder)
             is FirResolvable -> getSymbolsByResolvable(fir, expression, session, symbolBuilder)
             is FirNamedArgumentExpression -> getSymbolsByNameArgumentExpression(expression, analysisSession, symbolBuilder)
             else -> handleUnknownFirElement(expression, analysisSession, session, symbolBuilder)
         }
     }
 
+    private fun getSymbolByResolvedNameReference(
+        fir: FirResolvedNamedReference,
+        session: FirSession,
+        symbolBuilder: KtSymbolByFirBuilder
+    ): Collection<KtSymbol> = fir.toTargetSymbol(session, symbolBuilder)
+
     private fun KtSimpleNameExpression.isSyntheticOperatorReference() = when (this) {
         is KtOperationReferenceExpression -> operationSignTokenType in syntheticTokenTypes
         else -> false
@@ -265,19 +273,24 @@ internal object FirReferenceResolveHelper {
         return calleeReference.toTargetSymbol(session, symbolBuilder)
     }
 
+
     private fun getSymbolsByErrorNamedReference(
         fir: FirErrorNamedReference,
         symbolBuilder: KtSymbolByFirBuilder
-    ): List<KtSymbol> {
+    ): List<KtSymbol> =
-        val candidates = when (val diagnostic = fir.diagnostic) {
+        getFirSymbolsByErrorNamedReference(fir).mapNotNull { it.fir.buildSymbol(symbolBuilder) }
-            is ConeAmbiguityError -> diagnostic.candidates
+
-            is ConeOperatorAmbiguityError -> diagnostic.candidates
+
-            is ConeInapplicableCandidateError -> listOf(diagnostic.candidate.symbol)
+    fun getFirSymbolsByErrorNamedReference(
-            else -> emptyList()
+        errorNamedReference: FirErrorNamedReference,
-        }
+    ): Collection<FirBasedSymbol<*>> = when (val diagnostic = errorNamedReference.diagnostic) {
-        return candidates.mapNotNull { it.fir.buildSymbol(symbolBuilder) }
+        is ConeAmbiguityError -> diagnostic.candidates
+        is ConeOperatorAmbiguityError -> diagnostic.candidates
+        is ConeInapplicableCandidateError -> listOf(diagnostic.candidate.symbol)
+        else -> emptyList()
     }
 
+
     private fun getSymbolsByReturnExpression(
         expression: KtSimpleNameExpression,
         fir: FirReturnExpression,

idea/idea-frontend-fir/testData/components/expressionType/assignmentExpressionTarget.kt

@@ -0,0 +1,9 @@
+fun test(s: String) {
+    var i: Int = 0
+    <expr>i</expr> = s.length
+}
+// RESULT
+
+// expression: i
+// type: kotlin.Int
+

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/components/AbstractHLExpressionTypeTest.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2021 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 org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
+import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.test.framework.AbstractKtIdeaTest
+import org.jetbrains.kotlin.idea.test.framework.TestFileStructure
+import org.jetbrains.kotlin.idea.test.framework.TestStructureExpectedDataBlock
+import org.jetbrains.kotlin.idea.test.framework.TestStructureRenderer
+import org.jetbrains.kotlin.idea.util.application.executeOnPooledThread
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.test.KotlinTestUtils
+
+abstract class AbstractHLExpressionTypeTest : AbstractKtIdeaTest() {
+    override fun doTestByFileStructure(fileStructure: TestFileStructure) {
+        val expression = fileStructure.mainFile.selectedExpression as KtExpression?
+            ?: error("Selected expression was not provided")
+        val type = executeOnPooledThreadInReadAction {
+            analyze(expression) { expression.getKtType().render() }
+        }
+        val actual = TestStructureRenderer.render(
+            fileStructure,
+            TestStructureExpectedDataBlock(
+                "expression: ${expression.text}",
+                "type: $type"
+            )
+        )
+        KotlinTestUtils.assertEqualsToFile(fileStructure.filePath.toFile(), actual)
+    }
+}

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2021 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.util.KtTestUtil;
+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/expressionType")
+@TestDataPath("$PROJECT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+public class HLExpressionTypeTestGenerated extends AbstractHLExpressionTypeTest {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    public void testAllFilesPresentInExpressionType() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/expressionType"), Pattern.compile("^(.+)\\.kt$"), null, true);
+    }
+
+    @TestMetadata("assignmentExpressionTarget.kt")
+    public void testAssignmentExpressionTarget() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expressionType/assignmentExpressionTarget.kt");
+    }
+}

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/test/framework/AbstractKtIdeaTest.kt

@@ -6,7 +6,10 @@
 package org.jetbrains.kotlin.idea.test.framework
 
 import com.intellij.openapi.util.io.FileUtil
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 abstract class AbstractKtIdeaTest : KotlinLightCodeInsightFixtureTestCase() {
@@ -30,7 +33,7 @@ abstract class AbstractKtIdeaTest : KotlinLightCodeInsightFixtureTestCase() {
             filePath = testDataFile.toPath(),
             caretPosition = getCaretPosition(text),
             directives = directives,
-            mainFile = mainFile as TestFile.KtTestFile,
+            mainFile = mainFile as TestFile.KtTestRootFile,
             otherFiles = testFiles
         )
     }
@@ -51,12 +54,14 @@ abstract class AbstractKtIdeaTest : KotlinLightCodeInsightFixtureTestCase() {
     }
 
     private fun createTestFile(file: FileSplitter.FileNameWithText, isMainFile: Boolean): TestFile {
-        val psiFile = if (isMainFile) {
+        return if (isMainFile) {
-            myFixture.configureByText(file.name, file.text)
+            val (ktFile, expression) = SelectedExpressionProvider.getFileWithSelectedExpressions(file.text) {
+                myFixture.configureByText(file.name, it) as KtFile
+            }
+            TestFile.KtTestRootFile(ktFile, expression)
         } else {
-            myFixture.addFileToProject(file.name, file.text)
+            TestFile.createByPsiFile(myFixture.addFileToProject(file.name, file.text))
         }
-        return TestFile.createByPsiFile(psiFile)
     }
 
     abstract fun doTestByFileStructure(fileStructure: TestFileStructure)
@@ -68,7 +73,7 @@ private object FileSplitter {
     fun splitIntoFiles(text: String, defaultName: String): List<FileNameWithText> {
         val result = mutableListOf<FileNameWithText>()
         val stopAt = text.indexOfOrNull(KtTest.RESULT_DIRECTIVE) ?: text.length
-        var index = text.indexOfOrNull(KtTest.FILE_DIRECTIVE) ?: return listOf(FileNameWithText(defaultName, text))
+        var index = text.indexOfOrNull(KtTest.FILE_DIRECTIVE) ?: return listOf(FileNameWithText(defaultName, text.substring(0, stopAt).trim()))
         while (index < stopAt) {
             val eolIndex = text.indexOfOrNull("\n", index) ?: text.length
             val fileName = text.substring(index + KtTest.FILE_DIRECTIVE.length, eolIndex).trim()

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/test/framework/SelectedExpressionProvider.kt

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2021 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.test.framework
+
+import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
+
+object SelectedExpressionProvider {
+    fun getFileWithSelectedExpressions(fileText: String, createKtFile: (text: String) -> KtFile): KtFileWithSelectedExpression {
+        val fileTextWithoutTags = fileText.replace(TAGS.OPENING_EXPRESSION_TAG, "").replace(TAGS.CLOSING_EXPRESSION_TAG, "")
+        val ktFile = createKtFile(fileTextWithoutTags)
+        val selectedExpression = run {
+            val startCaretPosition = fileText.indexOf(TAGS.OPENING_EXPRESSION_TAG)
+            if (startCaretPosition < 0) {
+                return KtFileWithSelectedExpression(ktFile, selectedExpression = null)
+            }
+            val endCaretPosition = fileText.indexOf(TAGS.CLOSING_EXPRESSION_TAG)
+            if (endCaretPosition < 0) {
+                error("${TAGS.CLOSING_EXPRESSION_TAG} was not found in the file")
+            }
+            val elements = ktFile.elementsInRange(TextRange(startCaretPosition, endCaretPosition - TAGS.OPENING_EXPRESSION_TAG.length))
+            if (elements.size != 1) {
+                error("Expected one element at rage but found ${elements.size} [${elements.joinToString { it.text }}]")
+            }
+            elements.single() as KtElement
+        }
+        return KtFileWithSelectedExpression(ktFile, selectedExpression)
+    }
+
+    data class KtFileWithSelectedExpression(val file: KtFile, val selectedExpression: KtElement?)
+
+    object TAGS {
+        const val OPENING_EXPRESSION_TAG = "<expr>"
+        const val CLOSING_EXPRESSION_TAG = "</expr>"
+    }
+}

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/test/framework/TestFileStructure.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.test.framework
 
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiJavaFile
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import java.nio.file.Path
 
@@ -14,7 +15,7 @@ class TestFileStructure(
     val filePath: Path,
     val caretPosition: Int?,
     val directives: TestFileDirectives,
-    val mainFile: TestFile.KtTestFile,
+    val mainFile: TestFile.KtTestRootFile,
     val otherFiles: List<TestFile>,
 ) {
     val mainKtFile: KtFile
@@ -23,7 +24,9 @@ class TestFileStructure(
     val allFiles: List<TestFile> = listOf(mainFile) + otherFiles
 }
 
-data class TestStructureExpectedDataBlock(val name: String, val values: List<String>)
+data class TestStructureExpectedDataBlock(val name: String? = null, val values: List<String>) {
+    constructor(vararg values: String) : this(name = null, values.toList())
+}
 
 class TestFileDirectives(
     private val directives: Map<String, Any>
@@ -44,6 +47,7 @@ sealed class TestFile {
     abstract val psiFile: PsiFile
 
     data class KtTestFile(override val psiFile: KtFile) : TestFile()
+    data class KtTestRootFile(override val psiFile: KtFile, val selectedExpression: KtElement?) : TestFile()
     data class JavaTestFile(override val psiFile: PsiJavaFile) : TestFile()
 
     companion object {

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/test/framework/TestStructureRenderer.kt

@@ -5,7 +5,16 @@
 
 package org.jetbrains.kotlin.idea.test.framework
 
+import org.jetbrains.kotlin.psi.psiUtil.getStartOffsetIn
+
 object TestStructureRenderer {
+    fun render(testStructure: TestFileStructure, vararg expectedData: TestStructureExpectedDataBlock): String = buildString {
+        renderFiles(testStructure)
+        renderExpectedData(expectedData.toList())
+        renderCaretSymbol(testStructure)
+        renderExpressionTag(testStructure)
+    }
+
     fun render(testStructure: TestFileStructure, expectedData: List<TestStructureExpectedDataBlock>): String = buildString {
         renderFiles(testStructure)
         renderExpectedData(expectedData)
@@ -18,6 +27,14 @@ object TestStructureRenderer {
         }
     }
 
+    private fun StringBuilder.renderExpressionTag(testStructure: TestFileStructure) {
+        testStructure.mainFile.selectedExpression?.let { expression ->
+            val offset = expression.getStartOffsetIn(testStructure.mainKtFile)
+            insert(offset + expression.textLength, SelectedExpressionProvider.TAGS.CLOSING_EXPRESSION_TAG)
+            insert(offset, SelectedExpressionProvider.TAGS.OPENING_EXPRESSION_TAG)
+        }
+    }
+
     private fun StringBuilder.renderFiles(testStructure: TestFileStructure) {
         if (testStructure.otherFiles.isEmpty()) {
             appendLine(testStructure.mainFile.psiFile.text)
@@ -40,7 +57,7 @@ object TestStructureRenderer {
     }
 
     private fun StringBuilder.renderExpectedDataBlock(block: TestStructureExpectedDataBlock) {
-        appendLine("// ${block.name}")
+        block.name?.let { name -> appendLine("// $name") }
         block.values.forEach { value ->
             appendLine("// $value")
         }

idea/testData/fir/lazyResolve/inInit/results.txt

@@ -1 +1 @@
-this@R|/Foo|.R|/Foo.x| = Int(42)
+R|/Foo.x|

idea/testData/fir/lazyResolve/inSecondary/results.txt

@@ -1 +1 @@
-this@R|/Foo|.R|/Foo.x| = Int(42)
+R|/Foo.x|