FIR: support copy functions in data classes

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/DataClassUtils.kt

@@ -9,21 +9,24 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirFunctionTarget
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.impl.FirClassImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirMemberFunctionImpl
+import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
+import org.jetbrains.kotlin.fir.expressions.impl.FirFunctionCallImpl
 import org.jetbrains.kotlin.fir.expressions.impl.FirQualifiedAccessExpressionImpl
 import org.jetbrains.kotlin.fir.expressions.impl.FirReturnExpressionImpl
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedCallableReferenceImpl
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImpl
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtTypeReference
 
 internal fun KtClassOrObject.generateComponentFunctions(
     session: FirSession, firClass: FirClassImpl, packageFqName: FqName, classFqName: FqName
@@ -67,4 +70,64 @@ internal fun KtClassOrObject.generateComponentFunctions(
             }
         )
     }
+}
+
+private val copyName = Name.identifier("copy")
+
+internal fun KtClassOrObject.generateCopyFunction(
+    session: FirSession, firClass: FirClassImpl, packageFqName: FqName, classFqName: FqName,
+    firPrimaryConstructor: FirConstructor,
+    toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef
+) {
+    val symbol = FirFunctionSymbol(CallableId(packageFqName, classFqName, copyName))
+    firClass.addDeclaration(
+        FirMemberFunctionImpl(
+            session, this, symbol, copyName,
+            Visibilities.PUBLIC, Modality.FINAL,
+            isExpect = false, isActual = false,
+            isOverride = false, isOperator = false,
+            isInfix = false, isInline = false,
+            isTailRec = false, isExternal = false,
+            isSuspend = false, receiverTypeRef = null,
+            returnTypeRef = FirImplicitTypeRefImpl(session, this)
+        ).apply {
+            val copyFunction = this
+            val zippedParameters =
+                primaryConstructorParameters.zip(firClass.declarations.filterIsInstance<FirProperty>())
+            for ((ktParameter, firProperty) in zippedParameters) {
+                val name = ktParameter.nameAsSafeName
+                valueParameters += FirValueParameterImpl(
+                    session, ktParameter, name,
+                    ktParameter.typeReference.toFirOrErrorTypeRef(),
+                    FirQualifiedAccessExpressionImpl(session, ktParameter).apply {
+                        calleeReference = FirResolvedCallableReferenceImpl(session, ktParameter, name, firProperty.symbol)
+                    },
+                    isCrossinline = false, isNoinline = false, isVararg = false
+                )
+            }
+            body = FirSingleExpressionBlock(
+                session,
+                FirReturnExpressionImpl(
+                    session, this@generateCopyFunction,
+                    FirFunctionCallImpl(session, this@generateCopyFunction).apply {
+                        calleeReference = FirResolvedCallableReferenceImpl(
+                            session, this@generateCopyFunction, firClass.name,
+                            firPrimaryConstructor.symbol
+                        )
+                    }.apply {
+                        for ((ktParameter, firParameter) in primaryConstructorParameters.zip(valueParameters)) {
+                            this.arguments += FirQualifiedAccessExpressionImpl(session, ktParameter).apply {
+                                calleeReference = FirResolvedCallableReferenceImpl(
+                                    session, ktParameter, firParameter.name, firParameter.symbol
+                                )
+                            }
+                        }
+                    }
+                ).apply {
+                    target = FirFunctionTarget(null)
+                    target.bind(copyFunction)
+                }
+            )
+        }
+    )
 }

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

@@ -547,6 +547,10 @@ class RawFirBuilder(val session: FirSession, val stubMode: Boolean) {
 
                 if (classOrObject.hasModifier(DATA_KEYWORD) && firPrimaryConstructor != null) {
                     classOrObject.generateComponentFunctions(session, firClass, packageFqName, className)
+                    classOrObject.generateCopyFunction(session, firClass, packageFqName, className, firPrimaryConstructor) {
+                        toFirOrErrorType()
+                    }
+                    // TODO: equals, hashCode, toString
                 }
 
                 firClass

compiler/fir/psi2fir/testData/rawBuilder/expressions/annotated.txt

@@ -29,6 +29,10 @@ FILE: annotated.kt
             super<kotlin/Any>()
         }
 
+        public final fun copy(): <implicit> {
+            ^copy R|/Two.Two|()
+        }
+
     }
     public? final? fun bar(two: Two): kotlin/Unit {
         lval <destruct>: <implicit> = two#

compiler/fir/psi2fir/testData/rawBuilder/expressions/destructuring.txt

@@ -25,6 +25,10 @@ FILE: destructuring.kt
             ^component3 R|/Some.third|
         }
 
+        public final fun copy(first: Int = R|/Some.first|, second: Double = R|/Some.second|, third: String = R|/Some.third|): <implicit> {
+            ^copy R|/Some.Some|(R|<local>/first|, R|<local>/second|, R|<local>/third|)
+        }
+
     }
     public? final? fun foo(some: Some): kotlin/Unit {
         lval <destruct>: <implicit> = some#

compiler/fir/psi2fir/testData/rawBuilder/expressions/for.txt

@@ -43,6 +43,10 @@ FILE: for.kt
             ^component2 R|/Some.y|
         }
 
+        public final fun copy(x: Int = R|/Some.x|, y: Int = R|/Some.y|): <implicit> {
+            ^copy R|/Some.Some|(R|<local>/x|, R|<local>/y|)
+        }
+
     }
     public? final? fun baz(set: Set<Some>): kotlin/Unit {
         lval <range>: <implicit> = set#

compiler/fir/psi2fir/testData/rawBuilder/expressions/lambda.txt

@@ -18,6 +18,10 @@ FILE: lambda.kt
             ^component2 R|/Tuple.y|
         }
 
+        public final fun copy(x: Int = R|/Tuple.x|, y: Int = R|/Tuple.y|): <implicit> {
+            ^copy R|/Tuple.Tuple|(R|<local>/x|, R|<local>/y|)
+        }
+
     }
     public? final? inline fun use(f: ( (Tuple) -> Int )): <implicit> {
         ^use f#(Tuple#(Int(1), Int(2)))