FIR: fix atomic qualified acess sharing the same PSI with callee ref

Consider an atomic reference `i`, it's FIR representation is

- FirQualifiedAccessExpression
  - calleeReference : FirNamedReference

Currently, both of the above FIR elements uses the same PSI element as
the source. Such problems are not present with `this` or `super` because
we manually mark them as `ExplicitThisOrSuperReference`.

This change generalizes the previous `ExplicitThisOrSuperReference` as
`ReferenceInAtomicQualifiedAccess` and fixes it for more cases.

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirUnderscoreHelpers.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.analysis.checkers
 
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import org.jetbrains.kotlin.fir.FirRealSourceElementKind
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
@@ -18,7 +19,7 @@ fun checkUnderscoreDiagnostics(
     reporter: DiagnosticReporter,
     isExpression: Boolean
 ) {
-    if (source != null && source.kind !is FirFakeSourceElementKind) {
+    if (source != null && (source.kind is FirRealSourceElementKind || source.kind is FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)) {
         with(SourceNavigator.forSource(source)) {
             if (source.getRawIdentifier()?.isUnderscore == true) {
                 reporter.reportOn(

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/coneDiagnosticToFirDiagnostic.kt

@@ -77,8 +77,8 @@ private fun ConeDiagnostic.toFirDiagnostic(
         FirErrors.OUTER_CLASS_ARGUMENTS_REQUIRED.createOn(qualifiedAccessSource ?: source, this.symbol)
     is ConeNoTypeArgumentsOnRhsError ->
         FirErrors.NO_TYPE_ARGUMENTS_ON_RHS.createOn(qualifiedAccessSource ?: source, this.desiredCount, this.type)
-    is ConeSimpleDiagnostic -> when (source.kind) {
-        is FirFakeSourceElementKind -> null
+    is ConeSimpleDiagnostic -> when {
+        source.kind is FirFakeSourceElementKind && source.kind != FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess -> null
         else -> this.getFactory(source).createOn(qualifiedAccessSource ?: source)
     }
     is ConeInstanceAccessBeforeSuperCall -> FirErrors.INSTANCE_ACCESS_BEFORE_SUPER_CALL.createOn(source, this.target)

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -468,7 +468,7 @@ class ExpressionsConverter(
         var isReceiver = true
         var hasQuestionMarkAtLHS = false
         var firReceiverExpression: FirExpression? = null
-        lateinit var firCallableReference: FirQualifiedAccess
+        lateinit var namedReference: FirNamedReference
         callableReferenceExpression.forEachChildren {
             when (it.tokenType) {
                 COLONCOLON -> isReceiver = false
@@ -477,7 +477,7 @@ class ExpressionsConverter(
                     if (isReceiver) {
                         firReceiverExpression = getAsFirExpression(it, "Incorrect receiver expression")
                     } else {
-                        firCallableReference = convertSimpleNameExpression(it)
+                        namedReference = createSimpleNamedReference(it.toFirSourceElement(), it)
                     }
                 }
             }
@@ -485,7 +485,7 @@ class ExpressionsConverter(
 
         return buildCallableReferenceAccess {
             source = callableReferenceExpression.toFirSourceElement()
-            calleeReference = firCallableReference.calleeReference as FirNamedReference
+            calleeReference = namedReference
             explicitReceiver = firReceiverExpression
             this.hasQuestionMarkAtLHS = hasQuestionMarkAtLHS
         }
@@ -965,19 +965,29 @@ class ExpressionsConverter(
      * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.visitSimpleNameExpression
      */
     private fun convertSimpleNameExpression(referenceExpression: LighterASTNode): FirQualifiedAccessExpression {
+        val nameSource = referenceExpression.toFirSourceElement()
+        val referenceSourceElement = if (nameSource.kind is FirFakeSourceElementKind) {
+            nameSource
+        } else {
+            nameSource.fakeElement(FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
+        }
         return buildQualifiedAccessExpression {
-            source = referenceExpression.toFirSourceElement()
-
-            val nameSource = this@buildQualifiedAccessExpression.source
             val rawText = referenceExpression.asText
-            if (nameSource != null && rawText.isUnderscore) {
+            if (rawText.isUnderscore) {
                 nonFatalDiagnostics.add(ConeUnderscoreUsageWithoutBackticks(nameSource))
             }
+            source = nameSource
+            calleeReference = createSimpleNamedReference(referenceSourceElement, referenceExpression)
+        }
+    }
 
-            calleeReference = buildSimpleNamedReference {
-                source = nameSource
-                name = rawText.nameAsSafeName()
-            }
+    private fun createSimpleNamedReference(
+        sourceElement: FirSourceElement,
+        referenceExpression: LighterASTNode
+    ): FirNamedReference {
+        return buildSimpleNamedReference {
+            source = sourceElement
+            name = referenceExpression.asText.nameAsSafeName()
         }
     }
 

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -2222,7 +2222,7 @@ open class RawFirBuilder(
                 val sourceElement = expression.toFirSourceElement()
                 source = sourceElement
                 calleeReference = buildExplicitThisReference {
-                    source = sourceElement.fakeElement(FirFakeSourceElementKind.ExplicitThisOrSuperReference)
+                    source = sourceElement.fakeElement(FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
                     labelName = expression.getLabelName()
                 }
             }
@@ -2234,7 +2234,7 @@ open class RawFirBuilder(
             return buildQualifiedAccessExpression {
                 this.source = theSource
                 calleeReference = buildExplicitSuperReference {
-                    source = theSource.fakeElement(FirFakeSourceElementKind.ExplicitThisOrSuperReference)
+                    source = theSource.fakeElement(FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
                     labelName = expression.getLabelName()
                     superTypeRef = superType.toFirOrImplicitType()
                 }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -1183,7 +1183,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
 
     private fun FirVariable.toQualifiedAccess(): FirQualifiedAccessExpression = buildQualifiedAccessExpression {
         calleeReference = buildResolvedNamedReference {
-            source = this@toQualifiedAccess.source
+            source = this@toQualifiedAccess.source?.fakeElement(FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
             name = this@toQualifiedAccess.name
             resolvedSymbol = this@toQualifiedAccess.symbol
         }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -257,7 +257,10 @@ fun generateAccessExpression(
     buildQualifiedAccessExpression {
         this.source = qualifiedSource
         calleeReference = buildSimpleNamedReference {
-            this.source = calleeReferenceSource
+            this.source = if (calleeReferenceSource == qualifiedSource)
+                calleeReferenceSource?.fakeElement(FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess)
+            else
+                calleeReferenceSource
             this.name = name
         }
         if (diagnostic != null) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt

@@ -83,9 +83,10 @@ sealed class FirFakeSourceElementKind : FirSourceElementKind() {
 
     object ImplicitInvokeCall : FirFakeSourceElementKind()
 
-    // this/super expressions have FirThisReference/FirSuperReference
-    // with a fake sources which refers to this this/super expression
-    object ExplicitThisOrSuperReference : FirFakeSourceElementKind()
+    // Consider an atomic qualified access like `i`. In the FIR tree, both the FirQualifiedAccessExpression and its calleeReference uses
+    // `i` as the source. Hence, this fake kind is set on the `calleeReference` to make sure no PSI element is shared by multiple FIR
+    // elements. This also applies to `this` and `super` references.
+    object ReferenceInAtomicQualifiedAccess : FirFakeSourceElementKind()
 
     // for enum classes we have valueOf & values functions generated
     // with a fake sources which refers to this the enum class
@@ -266,13 +267,46 @@ sealed class FirPsiSourceElement(val psi: PsiElement) : FirSourceElement() {
             return element.textRange.endOffset
         }
     }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as FirPsiSourceElement
+
+        if (psi != other.psi) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        return psi.hashCode()
+    }
 }
 
 class FirRealPsiSourceElement(psi: PsiElement) : FirPsiSourceElement(psi) {
     override val kind: FirSourceElementKind get() = FirRealSourceElementKind
 }
 
-class FirFakeSourceElement(psi: PsiElement, override val kind: FirFakeSourceElementKind) : FirPsiSourceElement(psi)
+class FirFakeSourceElement(psi: PsiElement, override val kind: FirFakeSourceElementKind) : FirPsiSourceElement(psi) {
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+        if (!super.equals(other)) return false
+
+        other as FirFakeSourceElement
+
+        if (kind != other.kind) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = super.hashCode()
+        result = 31 * result + kind.hashCode()
+        return result
+    }
+}
 
 fun FirSourceElement.fakeElement(newKind: FirFakeSourceElementKind): FirSourceElement {
     return when (this) {
@@ -310,6 +344,30 @@ class FirLightSourceElement(
         val node = treeStructure.unwrap(lighterASTNode)
         return node.psi?.toFirPsiSourceElement(kind)
     }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as FirLightSourceElement
+
+        if (lighterASTNode != other.lighterASTNode) return false
+        if (startOffset != other.startOffset) return false
+        if (endOffset != other.endOffset) return false
+        if (treeStructure != other.treeStructure) return false
+        if (kind != other.kind) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = lighterASTNode.hashCode()
+        result = 31 * result + startOffset
+        result = 31 * result + endOffset
+        result = 31 * result + treeStructure.hashCode()
+        result = 31 * result + kind.hashCode()
+        return result
+    }
 }
 
 val FirSourceElement?.psi: PsiElement? get() = (this as? FirPsiSourceElement)?.psi

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt

@@ -52,6 +52,7 @@ class FirDiagnosticsHandler(testServices: TestServices) : FirAnalysisHandler(tes
         private val allowedKindsForDebugInfo = setOf(
             FirRealSourceElementKind,
             FirFakeSourceElementKind.DesugaredCompoundAssignment,
+            FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess,
         )
     }
 

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

@@ -41,6 +41,7 @@ internal class KtToFirMapping(firElement: FirElement, recorder: FirElementsRecor
 
     private val userTypeMapping = ConcurrentHashMap<KtUserType, FirElement>()
     fun getElement(ktElement: KtElement, state: FirModuleResolveState): FirElement? {
+
         mapping[ktElement]?.let { return it }
 
         val userType = when (ktElement) {

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

@@ -6,8 +6,10 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import org.jetbrains.kotlin.fir.FirRealPsiSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
-import org.jetbrains.kotlin.fir.realPsi
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
@@ -90,9 +92,12 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
     }
 
     private fun cacheElement(element: FirElement, cache: MutableMap<KtElement, FirElement>) {
-        (element.realPsi as? KtElement)?.let { psi ->
-            cache(psi, element, cache)
-        }
+        val psi = element.source
+            ?.takeIf {
+                it.kind == FirFakeSourceElementKind.ReferenceInAtomicQualifiedAccess || it is FirRealPsiSourceElement
+            }.psi as? KtElement
+            ?: return
+        cache(psi, element, cache)
     }
 
     companion object {