FIR: change source of existing type refs in DataClassMembersGenerator to avoid sources clash

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

@@ -36,7 +36,7 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 //T can be either PsiElement, or LighterASTNode
 abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Context<T> = Context()) {
 
-    abstract fun T.toFirSourceElement(kind: FirSourceElementKind = FirRealSourceElementKind): FirSourceElement
+    abstract fun T.toFirSourceElement(kind: FirFakeSourceElementKind? = null): FirSourceElement
 
     protected val implicitUnitType = baseSession.builtinTypes.unitType
     protected val implicitAnyType = baseSession.builtinTypes.anyType
@@ -697,28 +697,32 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         private val session: FirSession,
         private val source: T,
         private val classBuilder: FirRegularClassBuilder,
-        private val primaryConstructor: FirConstructor,
         private val zippedParameters: List<Pair<T, FirProperty>>,
         private val packageFqName: FqName,
         private val classFqName: FqName,
+        private val createClassTypeRefWithSourceKind: (FirFakeSourceElementKind) -> FirTypeRef,
+        private val createParameterTypeRefWithSourceKind: (FirProperty, FirFakeSourceElementKind) -> FirTypeRef,
     ) {
-        private val classTypeRef = primaryConstructor.returnTypeRef
-
         fun generate() {
             generateComponentFunctions()
             generateCopyFunction()
             // Refer to (IR utils or FIR backend) DataClassMembersGenerator for generating equals, hashCode, and toString
         }
 
-        private fun generateComponentAccess(parameterSource: FirSourceElement?, firProperty: FirProperty) =
+        private fun generateComponentAccess(
+            parameterSource: FirSourceElement?,
+            firProperty: FirProperty,
+            classTypeRefWithCorrectSourceKind: FirTypeRef,
+            firPropertyReturnTypeRefWithCorrectSourceKind: FirTypeRef
+        ) =
             buildQualifiedAccessExpression {
                 source = parameterSource
-                typeRef = firProperty.returnTypeRef
+                typeRef = firPropertyReturnTypeRefWithCorrectSourceKind
                 dispatchReceiver = buildThisReceiverExpression {
                     calleeReference = buildImplicitThisReference {
                         boundSymbol = classBuilder.symbol
                     }
-                    typeRef = classTypeRef
+                    typeRef = classTypeRefWithCorrectSourceKind
                 }
                 calleeReference = buildResolvedNamedReference {
                     source = parameterSource
@@ -755,7 +759,8 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         private fun generateCopyFunction() {
             classBuilder.addDeclaration(
                 buildSimpleFunction {
-                    source = this@DataClassMembersGenerator.source.toFirSourceElement()
+                    val classTypeRef = createClassTypeRefWithSourceKind(FirFakeSourceElementKind.DataClassGeneratedMembers)
+                    source = this@DataClassMembersGenerator.source.toFirSourceElement(FirFakeSourceElementKind.DataClassGeneratedMembers)
                     session = this@DataClassMembersGenerator.session
                     origin = FirDeclarationOrigin.Source
                     returnTypeRef = classTypeRef
@@ -764,7 +769,9 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     symbol = FirNamedFunctionSymbol(CallableId(packageFqName, classFqName, copyName))
                     for ((ktParameter, firProperty) in zippedParameters) {
                         val propertyName = firProperty.name
-                        val parameterSource = ktParameter?.toFirSourceElement()
+                        val parameterSource = ktParameter?.toFirSourceElement(FirFakeSourceElementKind.DataClassGeneratedMembers)
+                        val propertyReturnTypeRef =
+                            createParameterTypeRefWithSourceKind(firProperty, FirFakeSourceElementKind.DataClassGeneratedMembers)
                         valueParameters += buildValueParameter {
                             source = parameterSource
                             session = this@DataClassMembersGenerator.session
@@ -772,7 +779,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                             returnTypeRef = propertyReturnTypeRef
                             name = propertyName
                             symbol = FirVariableSymbol(propertyName)
-                            defaultValue = generateComponentAccess(parameterSource, firProperty)
+                            defaultValue = generateComponentAccess(parameterSource, firProperty, propertyReturnTypeRef, classTypeRef)
                             isCrossinline = false
                             isNoinline = false
                             isVararg = false
@@ -790,4 +797,14 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             resolvedSymbol = this@toQualifiedAccess.symbol
         }
     }
+
+    protected inline fun <R> withDefaultSourceElementKind(newDefault: FirSourceElementKind, action: () -> R): R {
+        val currentForced = context.forcedElementSourceKind
+        context.forcedElementSourceKind = newDefault
+        try {
+            return action()
+        } finally {
+            context.forcedElementSourceKind = currentForced
+        }
+    }
 }

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

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.fir.builder
 
 import kotlinx.collections.immutable.persistentListOf
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
-import org.jetbrains.kotlin.fir.FirFunctionTarget
-import org.jetbrains.kotlin.fir.FirLabel
-import org.jetbrains.kotlin.fir.FirLoopTarget
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.name.ClassId
@@ -28,4 +25,6 @@ class Context<T> {
     val firLoopTargets = mutableListOf<FirLoopTarget>()
     var capturedTypeParameters = persistentListOf<FirTypeParameterSymbol>()
     val arraySetArgument = mutableMapOf<T, FirExpression>()
+
+    var forcedElementSourceKind: FirSourceElementKind? = null
 }

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

@@ -10,13 +10,10 @@ import com.intellij.openapi.util.Ref
 import com.intellij.psi.tree.IElementType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.KtNodeTypes
-import org.jetbrains.kotlin.fir.FirLightSourceElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.FirSourceElementKind
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.BaseFirBuilder
 import org.jetbrains.kotlin.fir.builder.Context
 import org.jetbrains.kotlin.fir.builder.escapedStringToCharacter
-import org.jetbrains.kotlin.fir.toFirLightSourceElement
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens.*
@@ -31,7 +28,7 @@ open class BaseConverter(
 ) : BaseFirBuilder<LighterASTNode>(baseSession, context) {
     protected val implicitType = buildImplicitTypeRef()
 
-    override fun LighterASTNode.toFirSourceElement(kind: FirSourceElementKind): FirLightSourceElement {
+    override fun LighterASTNode.toFirSourceElement(kind: FirFakeSourceElementKind?): FirLightSourceElement {
         val startOffset = offset + tree.getStartOffset(this)
         val endOffset = offset + tree.getEndOffset(this)
         return toFirLightSourceElement(startOffset, endOffset, tree)

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.
  */
 
@@ -474,10 +474,11 @@ class DeclarationsConverter(
                             baseSession,
                             classNode,
                             this,
-                            firPrimaryConstructor,
                             zippedParameters,
                             context.packageFqName,
-                            context.className
+                            context.className,
+                            createClassTypeRefWithSourceKind = { firPrimaryConstructor.returnTypeRef },
+                            createParameterTypeRefWithSourceKind = { property, _ -> property.returnTypeRef },
                         ).generate()
                     }
 

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

@@ -62,8 +62,9 @@ class RawFirBuilder(
         return reference.accept(Visitor(), Unit) as FirTypeRef
     }
 
-    override fun PsiElement.toFirSourceElement(kind: FirSourceElementKind): FirPsiSourceElement<*> {
-        return this.toFirPsiSourceElement(kind)
+    override fun PsiElement.toFirSourceElement(kind: FirFakeSourceElementKind?): FirPsiSourceElement<*> {
+        val actualKind = kind ?: this@RawFirBuilder.context.forcedElementSourceKind ?: FirRealSourceElementKind
+        return this.toFirPsiSourceElement(actualKind)
     }
 
     override val PsiElement.elementType: IElementType
@@ -749,10 +750,17 @@ class RawFirBuilder(
                                 baseSession,
                                 classOrObject,
                                 this,
-                                firPrimaryConstructor,
                                 zippedParameters,
                                 context.packageFqName,
-                                context.className
+                                context.className,
+                                createClassTypeRefWithSourceKind = { firPrimaryConstructor.returnTypeRef.copyWithNewSourceKind(it) },
+                                createParameterTypeRefWithSourceKind = { property, newKind ->
+                                    // just making a shallow copy isn't enough type ref may be a function type ref
+                                    // and contain value parameters inside
+                                    withDefaultSourceElementKind(newKind) {
+                                        (property.returnTypeRef.psi as KtTypeReference).toFirOrImplicitType()
+                                    }
+                                },
                             ).generate()
                         }