[FIR] Substitute supertypes in nested classifier scope

KT-38992

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SupertypeUtils.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.transformers.createSubstitutionForSupertype
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
@@ -33,10 +34,11 @@ fun lookupSuperTypes(
     lookupInterfaces: Boolean,
     deep: Boolean,
     useSiteSession: FirSession,
-    supertypeSupplier: SupertypeSupplier = SupertypeSupplier.Default
+    supertypeSupplier: SupertypeSupplier = SupertypeSupplier.Default,
+    substituteTypes: Boolean = false
 ): List<ConeClassLikeType> {
     return mutableListOf<ConeClassLikeType>().also {
-        klass.symbol.collectSuperTypes(it, mutableSetOf(), deep, lookupInterfaces, useSiteSession, supertypeSupplier)
+        klass.symbol.collectSuperTypes(it, mutableSetOf(), deep, lookupInterfaces, substituteTypes, useSiteSession, supertypeSupplier)
     }
 }
 
@@ -116,6 +118,7 @@ private fun FirClassifierSymbol<*>.collectSuperTypes(
     visitedSymbols: MutableSet<FirClassifierSymbol<*>>,
     deep: Boolean,
     lookupInterfaces: Boolean,
+    substituteSuperTypes: Boolean,
     useSiteSession: FirSession,
     supertypeSupplier: SupertypeSupplier
 ) {
@@ -131,14 +134,30 @@ private fun FirClassifierSymbol<*>.collectSuperTypes(
             if (deep)
                 superClassTypes.forEach {
                     if (it !is ConeClassErrorType) {
-                        it.lookupTag.toSymbol(useSiteSession)?.collectSuperTypes(
-                            list,
-                            visitedSymbols,
-                            deep,
-                            lookupInterfaces,
-                            useSiteSession,
-                            supertypeSupplier
-                        )
+                        if (substituteSuperTypes) {
+                            val substitutedTypes = mutableListOf<ConeClassLikeType>()
+                            it.lookupTag.toSymbol(useSiteSession)?.collectSuperTypes(
+                                substitutedTypes,
+                                visitedSymbols,
+                                deep,
+                                lookupInterfaces,
+                                substituteSuperTypes,
+                                useSiteSession,
+                                supertypeSupplier
+                            )
+                            val substitutor = createSubstitutionForSupertype(it, useSiteSession)
+                            substitutedTypes.mapTo(list) { superType -> substitutor.substituteOrSelf(superType) as ConeClassLikeType }
+                        } else {
+                            it.lookupTag.toSymbol(useSiteSession)?.collectSuperTypes(
+                                list,
+                                visitedSymbols,
+                                deep,
+                                lookupInterfaces,
+                                substituteSuperTypes,
+                                useSiteSession,
+                                supertypeSupplier
+                            )
+                        }
                     }
                 }
         }
@@ -146,7 +165,7 @@ private fun FirClassifierSymbol<*>.collectSuperTypes(
             val expansion =
                 supertypeSupplier.expansionForTypeAlias(fir)?.computePartialExpansion(useSiteSession, supertypeSupplier) ?: return
             expansion.lookupTag.toSymbol(useSiteSession)
-                ?.collectSuperTypes(list, visitedSymbols, deep, lookupInterfaces, useSiteSession, supertypeSupplier)
+                ?.collectSuperTypes(list, visitedSymbols, deep, lookupInterfaces, substituteSuperTypes, useSiteSession, supertypeSupplier)
         }
         else -> error("?!id:1")
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAbstractTreeTransformerWithSuperTypes.kt

@@ -5,12 +5,9 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers
 
-import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.lookupSuperTypes
 import org.jetbrains.kotlin.fir.resolve.providers.getNestedClassifierScope
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -21,8 +18,6 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirMemberTypeParameterScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirNestedClassifierScope
 import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
-import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
-import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLookupTagWithFixedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
@@ -52,13 +47,17 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
     ): CompositeTransformResult<FirStatement> {
         return withScopeCleanup {
             // ? Is it Ok to use original file session here ?
-            val superTypes = lookupSuperTypes(firClass, lookupInterfaces = false, deep = true, useSiteSession = session).asReversed()
+            val superTypes = lookupSuperTypes(
+                firClass,
+                lookupInterfaces = false,
+                deep = true,
+                substituteTypes = true,
+                useSiteSession = session
+            ).asReversed()
             for (superType in superTypes) {
                 session.getNestedClassifierScope(superType.lookupTag)?.let { nestedClassifierScope ->
-                    val klass = superType.lookupTag.toSymbol(session)?.fir as? FirRegularClass
-                    val mapping = klass?.typeParameters?.map { it.symbol }?.zip(superType.typeArguments.map { it as ConeKotlinType })?.toMap()
-                    val substitutor = mapping?.let { ConeSubstitutorByMap(it) } ?: ConeSubstitutor.Empty
                     val scope = if (nestedClassifierScope is FirNestedClassifierScope) {
+                        val substitutor = createSubstitutionForSupertype(superType, session)
                         FirNestedClassifierScopeWithSubstitution(nestedClassifierScope, substitutor)
                     } else {
                         nestedClassifierScope
@@ -93,6 +92,13 @@ private class FirNestedClassifierScopeWithSubstitution(
 ) : FirScope() {
     override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
         val matchedClass = scope.getClassifierByName(name) ?: return
+        val substitutor = substitutor.takeIf { matchedClass.fir.isInner } ?: ConeSubstitutor.Empty
         processor(matchedClass, substitutor)
     }
+}
+
+fun createSubstitutionForSupertype(superType: ConeLookupTagBasedType, session: FirSession): ConeSubstitutor {
+    val klass = superType.lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: return ConeSubstitutor.Empty
+    val mapping = klass.typeParameters.map { it.symbol }.zip(superType.typeArguments.map { it as ConeKotlinType }).toMap()
+    return ConeSubstitutorByMap(mapping)
 }