FIR: refactor member generations for data class

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

@@ -642,37 +642,30 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
     }
 
-    fun List<Pair<T, FirProperty>>.generateComponentFunctions(
+    inner class DataClassMemberGenerator(
-        session: FirSession, firClassBuilder: AbstractFirRegularClassBuilder, packageFqName: FqName, classFqName: FqName,
+        private val session: FirSession,
-        firPrimaryConstructor: FirConstructor,
+        private val source: T,
+        private val classBuilder: AbstractFirRegularClassBuilder,
+        private val classTypeRef: FirTypeRef,
+        private val zippedParameters: List<Pair<T, FirProperty>>,
+        private val packageFqName: FqName,
+        private val classFqName: FqName,
     ) {
-        var componentIndex = 1
-        for ((sourceNode, firProperty) in this) {
-            if (!firProperty.isVal && !firProperty.isVar) continue
-            val name = Name.identifier("component$componentIndex")
-            componentIndex++
-            val parameterSource = sourceNode?.toFirSourceElement()
-            val target = FirFunctionTarget(labelName = null, isLambda = false)
-            val componentFunction = buildSimpleFunction {
-                source = parameterSource
-                this.session = session
-                returnTypeRef = buildImplicitTypeRef {
-                    source = parameterSource
-                }
-                receiverTypeRef = null
-                this.name = name
-                this.status = FirDeclarationStatusImpl(Visibilities.PUBLIC, Modality.FINAL)
-                this.symbol = FirNamedFunctionSymbol(CallableId(packageFqName, classFqName, name))
 
-                val returnExpression = buildReturnExpression {
+        fun generateMembers() {
-                    source = parameterSource
+            generateComponentFunctions()
-                    result = buildQualifiedAccessExpression {
+            generateCopyFunction()
+            // TODO: equals, hashCode, toString
+        }
+
+        private inline fun generateComponentAccess(parameterSource: FirSourceElement?, firProperty: FirProperty) =
+            buildQualifiedAccessExpression {
                 source = parameterSource
                 dispatchReceiver = buildThisReceiverExpression {
                     calleeReference = buildImplicitThisReference {
-                                boundSymbol = firClassBuilder.symbol
+                        boundSymbol = classBuilder.symbol
                     }
-                            typeRef = firPrimaryConstructor.returnTypeRef
+                    typeRef = classTypeRef
                 }
                 calleeReference = buildResolvedNamedReference {
                     source = parameterSource
@@ -680,57 +673,60 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     resolvedSymbol = firProperty.symbol
                 }
             }
+
+        private fun generateComponentFunctions() {
+            var componentIndex = 1
+            for ((sourceNode, firProperty) in zippedParameters) {
+                if (!firProperty.isVal && !firProperty.isVar) continue
+                val name = Name.identifier("component$componentIndex")
+                componentIndex++
+                val parameterSource = sourceNode?.toFirSourceElement()
+                val target = FirFunctionTarget(labelName = null, isLambda = false)
+                val componentFunction = buildSimpleFunction {
+                    source = parameterSource
+                    session = this@DataClassMemberGenerator.session
+                    returnTypeRef = buildImplicitTypeRef {
+                        source = parameterSource
+                    }
+                    receiverTypeRef = null
+                    this.name = name
+                    status = FirDeclarationStatusImpl(Visibilities.PUBLIC, Modality.FINAL)
+                    symbol = FirNamedFunctionSymbol(CallableId(packageFqName, classFqName, name))
+
+                    val returnExpression = buildReturnExpression {
+                        source = parameterSource
+                        result = generateComponentAccess(parameterSource, firProperty)
                         this.target = target
                     }
                     body = buildSingleExpressionBlock(returnExpression)
                 }.also {
                     target.bind(it)
                 }
-            firClassBuilder.addDeclaration(componentFunction)
+                classBuilder.addDeclaration(componentFunction)
             }
         }
 
         private val copyName = Name.identifier("copy")
 
-    fun List<Pair<T, FirProperty>>.generateCopyFunction(
+        private fun generateCopyFunction() {
-        session: FirSession,
-        classOrObject: T,
-        classBuilder: AbstractFirRegularClassBuilder,
-        packageFqName: FqName,
-        classFqName: FqName,
-        firPrimaryConstructor: FirConstructor,
-    ) {
             classBuilder.addDeclaration(
                 buildSimpleFunction {
-                source = classOrObject?.toFirSourceElement()
+                    source = this@DataClassMemberGenerator.source.toFirSourceElement()
-                this.session = session
+                    session = this@DataClassMemberGenerator.session
-                returnTypeRef = firPrimaryConstructor.returnTypeRef
+                    returnTypeRef = classTypeRef
                     name = copyName
                     status = FirDeclarationStatusImpl(Visibilities.PUBLIC, Modality.FINAL)
                     symbol = FirNamedFunctionSymbol(CallableId(packageFqName, classFqName, copyName))
-                for ((ktParameter, firProperty) in this@generateCopyFunction) {
+                    for ((ktParameter, firProperty) in zippedParameters) {
                         val propertyName = firProperty.name
                         val parameterSource = ktParameter?.toFirSourceElement()
                         valueParameters += buildValueParameter {
                             source = parameterSource
-                        this.session = session
+                            session = this@DataClassMemberGenerator.session
                             returnTypeRef = firProperty.returnTypeRef
                             name = propertyName
                             symbol = FirVariableSymbol(propertyName)
-                        defaultValue = buildQualifiedAccessExpression {
+                            defaultValue = generateComponentAccess(parameterSource, firProperty)
-                            source = parameterSource
-                            dispatchReceiver = buildThisReceiverExpression {
-                                calleeReference = buildImplicitThisReference {
-                                    boundSymbol = classBuilder.symbol
-                                }
-                                typeRef = firPrimaryConstructor.returnTypeRef
-                            }
-                            calleeReference = buildResolvedNamedReference {
-                                source = parameterSource
-                                this.name = propertyName
-                                resolvedSymbol = firProperty.symbol
-                            }
-                        }
                             isCrossinline = false
                             isNoinline = false
                             isVararg = false
@@ -741,6 +737,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 },
             )
         }
+    }
 
     private fun FirVariable<*>.toQualifiedAccess(): FirQualifiedAccessExpression = buildQualifiedAccessExpression {
         calleeReference = buildResolvedNamedReference {

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

@@ -467,13 +467,15 @@ class DeclarationsConverter(
                     //parse data class
                     if (modifiers.isDataClass() && firPrimaryConstructor != null) {
                         val zippedParameters = properties.map { it.source?.lightNode!! to it }
-                        zippedParameters.generateComponentFunctions(
+                        DataClassMemberGenerator(
-                            baseSession, this, context.packageFqName, context.className, firPrimaryConstructor
+                            baseSession,
-                        )
+                            classNode,
-                        zippedParameters.generateCopyFunction(
+                            this,
-                            baseSession, classNode, this, context.packageFqName, context.className, firPrimaryConstructor
+                            firPrimaryConstructor.returnTypeRef,
-                        )
+                            zippedParameters,
-                        // TODO: equals, hashCode, toString
+                            context.packageFqName,
+                            context.className
+                        ).generateMembers()
                     }
 
                     if (modifiers.isEnum()) {

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

@@ -693,13 +693,15 @@ class RawFirBuilder(
                             val zippedParameters = classOrObject.primaryConstructorParameters.zip(
                                 declarations.filterIsInstance<FirProperty>(),
                             )
-                            zippedParameters.generateComponentFunctions(
+                            DataClassMemberGenerator(
-                                baseSession, this, context.packageFqName, context.className, firPrimaryConstructor,
+                                baseSession,
-                            )
+                                classOrObject,
-                            zippedParameters.generateCopyFunction(
+                                this,
-                                baseSession, classOrObject, this, context.packageFqName, context.className, firPrimaryConstructor,
+                                firPrimaryConstructor.returnTypeRef,
-                            )
+                                zippedParameters,
-                            // TODO: equals, hashCode, toString
+                                context.packageFqName,
+                                context.className
+                            ).generateMembers()
                         }
 
                         if (classOrObject.hasModifier(ENUM_KEYWORD)) {