[FIR] Use captured, constructed class type parameters in RawFirBuilder

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.builder
 
 import com.intellij.psi.tree.IElementType
+import kotlinx.collections.immutable.mutate
 import org.jetbrains.kotlin.KtNodeTypes.*
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -68,6 +69,18 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
     }
 
+    inline fun <T> withCapturedTypeParameters(block: () -> T): T {
+        val previous = context.capturedTypeParameters
+        val result = block()
+        context.capturedTypeParameters = previous
+        return result
+    }
+
+    fun addCapturedTypeParameters(typeParameters: List<FirTypeParameterRef>) {
+        context.capturedTypeParameters =
+            context.capturedTypeParameters.addAll(0, typeParameters.map { typeParameter -> typeParameter.symbol })
+    }
+
     fun callableIdForName(name: Name, local: Boolean = false) =
         when {
             local -> CallableId(name)
@@ -146,17 +159,34 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
     fun T?.toDelegatedSelfType(firObject: FirAnonymousObjectBuilder): FirResolvedTypeRef {
         return buildResolvedTypeRef {
             source = this@toDelegatedSelfType?.toFirSourceElement()
-            type = ConeClassLikeTypeImpl(firObject.symbol.toLookupTag(), emptyArray(), false)
+            type = ConeClassLikeTypeImpl(
+                firObject.symbol.toLookupTag(),
+                firObject.typeParameters.map { ConeTypeParameterTypeImpl(it.symbol.toLookupTag(), false) }.toTypedArray(),
+                false
+            )
         }
     }
 
-    fun typeParametersFromSelfType(delegatedSelfTypeRef: FirTypeRef): List<FirTypeParameter> {
+    fun typeParametersFromSelfType(
+        delegatedSelfTypeRef: FirTypeRef
+    ): List<FirTypeParameterRef> {
         return delegatedSelfTypeRef.coneTypeSafe<ConeKotlinType>()
             ?.typeArguments
-            ?.map { ((it as ConeTypeParameterType).lookupTag.symbol as FirTypeParameterSymbol).fir }
+            ?.map {
+                buildConstructedClassTypeParameterRef {
+                    symbol = ((it as ConeTypeParameterType).lookupTag.symbol as FirTypeParameterSymbol)
+                }
+            }
             ?: emptyList()
     }
 
+    fun constructorTypeParametersFromConstructedClass(ownerTypeParameters: List<FirTypeParameterRef>): List<FirTypeParameterRef> {
+        return ownerTypeParameters.mapNotNull {
+            val declaredTypeParameter = (it as? FirTypeParameter) ?: return@mapNotNull null
+            buildConstructedClassTypeParameterRef { symbol = declaredTypeParameter.symbol }
+        }
+    }
+
     fun FirLoopBuilder.configure(generateBlock: () -> FirBlock): FirLoop {
         label = context.firLabels.pop()
         val target = FirLoopTarget(label?.name)

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

@@ -5,11 +5,13 @@
 
 package org.jetbrains.kotlin.fir.builder
 
+import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.FirFunctionTarget
 import org.jetbrains.kotlin.fir.FirLabel
 import org.jetbrains.kotlin.fir.FirLoopTarget
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.builder.FirFunctionCallBuilder
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 
@@ -22,5 +24,6 @@ class Context<T> {
     val firFunctionCalls = mutableListOf<FirFunctionCallBuilder>()
     val firLabels = mutableListOf<FirLabel>()
     val firLoopTargets = mutableListOf<FirLoopTarget>()
+    var capturedTypeParameters = persistentListOf<FirTypeParameterSymbol>()
     val arraySetArgument = mutableMapOf<T, FirExpression>()
 }

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

@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.fir.builder
 
 import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.IElementType
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.mutate
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -156,8 +158,11 @@ class RawFirBuilder(
             convert()
 
         private fun KtDeclaration.toFirDeclaration(
-            delegatedSuperType: FirTypeRef, delegatedSelfType: FirResolvedTypeRef,
+            delegatedSuperType: FirTypeRef,
-            owner: KtClassOrObject, ownerClassBuilder: FirClassBuilder, hasPrimaryConstructor: Boolean,
+            delegatedSelfType: FirResolvedTypeRef,
+            owner: KtClassOrObject,
+           ownerClassBuilder: FirClassBuilder, hasPrimaryConstructor: Boolean,
+            ownerTypeParameters: List<FirTypeParameterRef>
         ): FirDeclaration {
             return when (this) {
                 is KtSecondaryConstructor -> {
@@ -166,6 +171,7 @@ class RawFirBuilder(
                         delegatedSelfType,
                         owner,
                         hasPrimaryConstructor,
+                        ownerTypeParameters
                     )
                 }
                 is KtEnumEntry -> {
@@ -408,6 +414,7 @@ class RawFirBuilder(
             delegatedSelfTypeRef: FirTypeRef?,
             delegatedEnumSuperTypeRef: FirTypeRef?,
             classKind: ClassKind,
+            containerTypeParameters: List<FirTypeParameterRef>
         ): FirTypeRef {
             var superTypeCallEntry: KtSuperTypeCallEntry? = null
             var delegatedSuperTypeRef: FirTypeRef? = null
@@ -475,6 +482,7 @@ class RawFirBuilder(
                 delegatedSuperTypeRef,
                 delegatedSelfTypeRef ?: delegatedSuperTypeRef,
                 owner = this,
+                containerTypeParameters
             )
             container.declarations += firPrimaryConstructor
             return delegatedSuperTypeRef
@@ -485,6 +493,7 @@ class RawFirBuilder(
             delegatedSuperTypeRef: FirTypeRef,
             delegatedSelfTypeRef: FirTypeRef,
             owner: KtClassOrObject,
+            ownerTypeParameters: List<FirTypeParameterRef>
         ): FirConstructor {
             val constructorCallee = superTypeCallEntry?.calleeExpression?.toFirSourceElement()
             val constructorSource = (this ?: owner).toFirSourceElement()
@@ -519,7 +528,7 @@ class RawFirBuilder(
                 this.status = status
                 symbol = FirConstructorSymbol(callableIdForClassConstructor())
                 delegatedConstructor = firDelegatedCall
-                typeParameters += typeParametersFromSelfType(delegatedSelfTypeRef)
+                typeParameters += constructorTypeParametersFromConstructedClass(ownerTypeParameters)
                 this@toFirConstructor?.extractAnnotationsTo(this)
                 this@toFirConstructor?.extractValueParametersTo(this)
             }
@@ -592,6 +601,7 @@ class RawFirBuilder(
                             correctedEnumSelfTypeRef,
                             delegatedEntrySelfType,
                             owner = ktEnumEntry,
+                            typeParameters
                         )
                         for (declaration in ktEnumEntry.declarations) {
                             declarations += declaration.toFirDeclaration(
@@ -600,6 +610,7 @@ class RawFirBuilder(
                                 ktEnumEntry,
                                 ownerClassBuilder = this,
                                 hasPrimaryConstructor = true,
+                                ownerTypeParameters = emptyList()
                             )
                         }
                     }
@@ -631,6 +642,9 @@ class RawFirBuilder(
                     isData = (classOrObject as? KtClass)?.isData() == true
                     isInline = classOrObject.hasModifier(INLINE_KEYWORD)
                 }
+                withCapturedTypeParameters {
+                    if (!status.isInner) context.capturedTypeParameters = context.capturedTypeParameters.clear()
+
                     val classBuilder = if (status.modality == Modality.SEALED) FirSealedClassBuilder() else FirClassImplBuilder()
                     classBuilder.apply {
                         source = classOrObject.toFirSourceElement()
@@ -643,9 +657,13 @@ class RawFirBuilder(
 
                         classOrObject.extractAnnotationsTo(this)
                         classOrObject.extractTypeParametersTo(this)
+                        typeParameters += context.capturedTypeParameters.map { buildOuterClassTypeParameterRef { symbol = it } }
+
+                        addCapturedTypeParameters(typeParameters.take(classOrObject.typeParameters.size))
 
                         val delegatedSelfType = classOrObject.toDelegatedSelfType(this)
-                    val delegatedSuperType = classOrObject.extractSuperTypeListEntriesTo(this, delegatedSelfType, null, classKind)
+                        val delegatedSuperType =
+                            classOrObject.extractSuperTypeListEntriesTo(this, delegatedSelfType, null, classKind, typeParameters)
 
                         val primaryConstructor = classOrObject.primaryConstructor
                         val firPrimaryConstructor = declarations.firstOrNull() as? FirConstructor
@@ -662,8 +680,11 @@ class RawFirBuilder(
                         for (declaration in classOrObject.declarations) {
                             addDeclaration(
                                 declaration.toFirDeclaration(
-                                delegatedSuperType, delegatedSelfType, classOrObject, classBuilder,
+                                    delegatedSuperType,
-                                hasPrimaryConstructor = primaryConstructor != null,
+                                    delegatedSelfType,
+                                    classOrObject,
+                                    classBuilder, hasPrimaryConstructor = primaryConstructor != null,
+                                    typeParameters
                                 ),
                             )
                         }
@@ -688,6 +709,7 @@ class RawFirBuilder(
                     }.build()
                 }
             }
+        }
 
         override fun visitObjectLiteralExpression(expression: KtObjectLiteralExpression, data: Unit): FirElement {
             val objectDeclaration = expression.objectDeclaration
@@ -698,11 +720,19 @@ class RawFirBuilder(
                     classKind = ClassKind.OBJECT
                     scopeProvider = baseScopeProvider
                     symbol = FirAnonymousObjectSymbol()
+                    typeParameters += context.capturedTypeParameters.map { buildOuterClassTypeParameterRef { symbol = it } }
                     val delegatedSelfType = objectDeclaration.toDelegatedSelfType(this)
                     objectDeclaration.extractAnnotationsTo(this)
-                    val delegatedSuperType = objectDeclaration.extractSuperTypeListEntriesTo(this, delegatedSelfType, null, ClassKind.CLASS)
+                    val delegatedSuperType = objectDeclaration.extractSuperTypeListEntriesTo(
+                        this,
+                        delegatedSelfType,
+                        null,
+                        ClassKind.CLASS,
+                        containerTypeParameters = emptyList()
+                    )
                     typeRef = delegatedSelfType
 
+
                     for (declaration in objectDeclaration.declarations) {
                         declarations += declaration.toFirDeclaration(
                             delegatedSuperType,
@@ -710,6 +740,7 @@ class RawFirBuilder(
                             owner = objectDeclaration,
                             ownerClassBuilder = this,
                             hasPrimaryConstructor = false,
+                            ownerTypeParameters = emptyList()
                         )
                     }
                 }
@@ -789,7 +820,10 @@ class RawFirBuilder(
                 for (valueParameter in function.valueParameters) {
                     valueParameters += valueParameter.convert<FirValueParameter>()
                 }
+                withCapturedTypeParameters {
+                    if (this is FirSimpleFunctionBuilder) addCapturedTypeParameters(this.typeParameters)
                     body = function.buildFirBody()
+                }
                 context.firFunctionTargets.removeLast()
             }.build().also {
                 target.bind(it)
@@ -884,6 +918,7 @@ class RawFirBuilder(
             delegatedSelfTypeRef: FirTypeRef,
             owner: KtClassOrObject,
             hasPrimaryConstructor: Boolean,
+            ownerTypeParameters: List<FirTypeParameterRef>
         ): FirConstructor {
             val target = FirFunctionTarget(labelName = null, isLambda = false)
             return buildConstructor {
@@ -906,7 +941,7 @@ class RawFirBuilder(
                 )
                 this@RawFirBuilder.context.firFunctionTargets += target
                 extractAnnotationsTo(this)
-                typeParameters += typeParametersFromSelfType(delegatedSelfTypeRef)
+                typeParameters += constructorTypeParametersFromConstructedClass(ownerTypeParameters)
                 extractValueParametersTo(this)
                 body = buildFirBody()
                 this@RawFirBuilder.context.firFunctionTargets.removeLast()
@@ -952,6 +987,7 @@ class RawFirBuilder(
                 returnTypeRef = propertyType
                 name = propertyName
                 this.isVar = isVar
+
                 initializer = propertyInitializer
 
                 if (this@toFirProperty.isLocal) {
@@ -974,6 +1010,9 @@ class RawFirBuilder(
                     isLocal = false
                     receiverTypeRef = receiverTypeReference.convertSafe()
                     symbol = FirPropertySymbol(callableIdForName(propertyName))
+                    extractTypeParametersTo(this)
+                    withCapturedTypeParameters {
+                        addCapturedTypeParameters(this.typeParameters)
                         val delegateBuilder = if (hasDelegate()) {
                             FirWrappedDelegateExpressionBuilder().apply {
                                 source = if (stubMode) null else delegateExpression?.toFirSourceElement()
@@ -987,7 +1026,6 @@ class RawFirBuilder(
                             isConst = hasModifier(CONST_KEYWORD)
                             isLateInit = hasModifier(LATEINIT_KEYWORD)
                         }
-                    extractTypeParametersTo(this)
 
                         getter = this@toFirProperty.getter.toFirPropertyAccessor(this@toFirProperty, propertyType, isGetter = true)
                         setter = if (isVar) {
@@ -1004,6 +1042,7 @@ class RawFirBuilder(
                             receiver
                         )
                     }
+                }
                 extractAnnotationsTo(this)
             }
         }