FIR IDE: find FIR element by PSI one only if FIR have a real source

Also, add a check that there is only one element with the same real source

idea/idea-fir/build.gradle.kts

@@ -18,6 +18,7 @@ dependencies {
     testCompile(projectTests(":idea"))
     testCompile(projectTests(":compiler:tests-common"))
     testCompile(projectTests(":idea:idea-test-framework"))
+    testCompile(projectTests(":idea:idea-frontend-fir"))
     testCompile(project(":kotlin-test:kotlin-test-junit"))
     testCompile(commonDep("junit:junit"))
 

idea/idea-fir/tests/org/jetbrains/kotlin/checkers/AbstractFirPsiCheckerTest.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.checkers
 
 import com.intellij.rt.execution.junit.FileComparisonFailure
 import org.jetbrains.kotlin.idea.test.withCustomCompilerOptions
+import org.jetbrains.kotlin.idea.withPossiblyDisabledDuplicatedFirSourceElementsException
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import java.io.File
 
@@ -30,10 +31,13 @@ abstract class AbstractFirPsiCheckerTest : AbstractPsiCheckerTest() {
         checkWeakWarnings: Boolean
     ): Long {
         val file = file
-        return withCustomCompilerOptions(file.text, project, module) {
+        val fileText = file.text
+        return withCustomCompilerOptions(fileText, project, module) {
             val doComparison = InTextDirectivesUtils.isDirectiveDefined(myFixture.file.text, "FIR_COMPARISON")
             try {
-                myFixture.checkHighlighting(checkWarnings, checkInfos, checkWeakWarnings)
+                withPossiblyDisabledDuplicatedFirSourceElementsException(fileText) {
+                    myFixture.checkHighlighting(checkWarnings, checkInfos, checkWeakWarnings)
+                }
             } catch (e: FileComparisonFailure) {
                 if (doComparison) {
                     // Even this is very partial check (only error compatibility, no warnings / infos)

idea/idea-fir/tests/org/jetbrains/kotlin/idea/highlighter/AbstractFirHighlightingTest.kt

@@ -5,7 +5,9 @@
 
 package org.jetbrains.kotlin.idea.highlighter
 
+import org.jetbrains.kotlin.idea.shouldBeRethrown
 import org.jetbrains.kotlin.idea.test.ProjectDescriptorWithStdlibSources
+import org.jetbrains.kotlin.idea.withPossiblyDisabledDuplicatedFirSourceElementsException
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 
 abstract class AbstractFirHighlightingTest : AbstractHighlightingTest() {
@@ -21,9 +23,11 @@ abstract class AbstractFirHighlightingTest : AbstractHighlightingTest() {
 
         try {
             // warnings are not supported yet
-            myFixture.checkHighlighting(/* checkWarnings= */ false, checkInfos, /* checkWeakWarnings= */ false)
+            withPossiblyDisabledDuplicatedFirSourceElementsException(fileText) {
+                myFixture.checkHighlighting(/* checkWarnings= */ false, checkInfos, /* checkWeakWarnings= */ false)
+            }
         } catch (e: Throwable) {
-            if (doComparison) throw e
+            if (doComparison || e.shouldBeRethrown()) throw e
             return
         }
         if (!doComparison) {

idea/idea-fir/tests/org/jetbrains/kotlin/idea/ideaFirTestUtils.kt

@@ -0,0 +1,13 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.idea.fir.low.level.api.DuplicatedFirSourceElementsException
+
+fun Throwable.shouldBeRethrown(): Boolean = when (this) {
+    is DuplicatedFirSourceElementsException -> true
+    else -> false
+}

idea/idea-fir/tests/org/jetbrains/kotlin/idea/resolve/AbstractFirReferenceResolveTest.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.resolve
 
 import org.jetbrains.kotlin.idea.completion.test.configureWithExtraFile
+import org.jetbrains.kotlin.idea.shouldBeRethrown
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
@@ -27,6 +28,7 @@ abstract class AbstractFirReferenceResolveTest : AbstractReferenceResolveTest()
             try {
                 performChecks()
             } catch (t: Throwable) {
+                if (t.shouldBeRethrown()) throw t
                 return
             }
             throw AssertionError("Looks like test is passing, please remove IGNORE_FIR")

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveState.kt

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
+import com.intellij.psi.impl.source.tree.LeafPsiElement
+import org.jetbrains.kotlin.cfg.pseudocode.containingDeclarationForPseudocode
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -18,12 +20,16 @@ import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.LibrarySourceInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtTypeReference
 
 internal interface FirModuleResolveState {
     val sessionProvider: FirProjectSessionProvider
@@ -96,7 +102,36 @@ internal class FirModuleResolveStateImpl(override val sessionProvider: FirProjec
     }
 
     override fun record(psi: KtElement, fir: FirElement) {
-        cache.putIfAbsent(psi, fir)
+        val existingFir = cache[psi]
+        if (existingFir != null && existingFir !== fir) {
+            when {
+                existingFir is FirTypeRef && fir is FirTypeRef && psi is KtTypeReference -> {
+                    // FirTypeRefs are often created during resolve
+                    // a lot of them with have the same source
+                    // we want to take the most "resolved one" here
+                    if (fir is FirResolvedTypeRefImpl && existingFir !is FirResolvedTypeRefImpl) {
+                        cache[psi] = fir
+                    }
+                }
+                existingFir.isErrorElement && !fir.isErrorElement -> {
+                    // TODO better handle error elements
+                    // but for now just take first non-error one if such exist
+                    cache[psi] = fir
+                }
+                existingFir.isErrorElement || fir.isErrorElement -> {
+                    // do nothing and maybe upgrade to a non-error element in the branch above in the future
+                }
+                else -> {
+
+                    if (DuplicatedFirSourceElementsException.IS_ENABLED) {
+                        throw DuplicatedFirSourceElementsException(existingFir, fir, psi)
+                    }
+                }
+            }
+        }
+        if (existingFir == null) {
+            cache[psi] = fir
+        }
     }
 
     override fun record(psi: KtElement, diagnostic: Diagnostic) {
@@ -115,5 +150,52 @@ internal class FirModuleResolveStateImpl(override val sessionProvider: FirProjec
     }
 }
 
+class DuplicatedFirSourceElementsException(
+    private val existingFir: FirElement,
+    private val newFir: FirElement,
+    private val psi: KtElement
+) : IllegalStateException() {
+    override val message: String?
+        get() =
+            """|The PSI element should be used only once as a real PSI source of FirElement,
+               |the elements ${if (existingFir.source === newFir.source) "HAVE" else "DON'T HAVE"} the same instances of source elements 
+               |
+               |existing FIR element is $existingFir with text:
+               |${existingFir.render().trim()}
+               |
+               |new FIR element is $newFir with text:
+               | ${newFir.render().trim()}
+               |
+               |PSI element is $psi with text in context:
+               |${getElementInContext(psi)}""".trimMargin()
+
+    private fun getElementInContext(neededElement: KtElement): String {
+        val context = neededElement.containingDeclarationForPseudocode ?: neededElement.containingKtFile
+        val builder = StringBuilder()
+        context.accept(object : PsiElementVisitor() {
+            override fun visitElement(element: PsiElement) {
+                if (element === neededElement) builder.append("<$ELEMENT_TAG>")
+                if (element is LeafPsiElement) {
+                    builder.append(element.text)
+                } else {
+                    element.acceptChildren(this)
+                }
+                if (element === neededElement) builder.append("</$ELEMENT_TAG>")
+            }
+        })
+        return builder.toString().trimIndent().trim()
+    }
+
+    companion object {
+        private const val ELEMENT_TAG = "ELEMENT"
+
+        // The are some cases which are still generates FIR elements with duplicated source elements
+        // Then such case is met, it's better to be fixed
+        // but exception reporting can be easily disabled by setting this to false
+        var IS_ENABLED = true
+    }
+}
+
 internal fun KtElement.firResolveState(): FirModuleResolveState =
-    FirIdeResolveStateService.getInstance(project).getResolveState(getModuleInfo())
+    FirIdeResolveStateService.getInstance(project).getResolveState(getModuleInfo())
+

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirResolutionApi.kt

@@ -7,12 +7,9 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
-import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredCallableSymbols
@@ -23,6 +20,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.FirUserTypeRef
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
@@ -260,7 +258,7 @@ internal fun KtElement.getOrBuildFir(
     return state[this] ?: run {
         containerFir.accept(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {
-                (element.psi as? KtElement)?.let {
+                (element.realPsi as? KtElement)?.let {
                     state.record(it, element)
                 }
                 element.acceptChildren(this)

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

@@ -0,0 +1,12 @@
+/*
+ * 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.fir.low.level.api
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
+
+internal val FirElement.isErrorElement
+    get() = this is FirDiagnosticHolder

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/testUtils.kt

@@ -0,0 +1,30 @@
+/*
+ * 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
+
+import org.jetbrains.kotlin.idea.fir.low.level.api.DuplicatedFirSourceElementsException
+import org.jetbrains.kotlin.test.InTextDirectivesUtils
+
+/**
+ * Temporary
+ * @see org.jetbrains.kotlin.idea.fir.low.level.api.DuplicatedFirSourceElementsException.IS_ENABLED
+ */
+inline fun <T> withPossiblyDisabledDuplicatedFirSourceElementsException(fileText: String, action: () -> T): T {
+    val isDisabled = InTextDirectivesUtils.isDirectiveDefined(fileText, "IGNORE_DUPLICATED_FIR_SOURCE_EXCEPTION")
+
+    @Suppress("LiftReturnOrAssignment")
+    if (isDisabled) {
+        val wasEnabled = DuplicatedFirSourceElementsException.IS_ENABLED
+        DuplicatedFirSourceElementsException.IS_ENABLED = false
+        try {
+            return action()
+        } finally {
+            DuplicatedFirSourceElementsException.IS_ENABLED = wasEnabled
+        }
+    } else {
+        return action()
+    }
+}

idea/testData/checker/ResolveTypeInAnnotationArgumentRuntime.kt

@@ -1,3 +1,4 @@
+// IGNORE_DUPLICATED_FIR_SOURCE_EXCEPTION
 import kotlin.reflect.KClass
 
 annotation class Ann(val value: KClass<*>)

idea/testData/highlighter/deprecated/Setter.kt

@@ -1,3 +1,4 @@
+// IGNORE_DUPLICATED_FIR_SOURCE_EXCEPTION
 fun test() {
     MyClass().test1
     MyClass().<warning descr="[DEPRECATION] 'setter for test1: Int' is deprecated. Use A instead">test1</warning> = 0