FIR: Refactor nested scope, fix super-type resolving jump context

#KT-24017 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/descriptors/FirTypeResolveTransformer.kt

@@ -12,26 +12,27 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
-import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirCompositeScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirExplicitImportingScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirNestedClassifierScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirSelfImportingScope
-import org.jetbrains.kotlin.fir.transformInplace
+import org.jetbrains.kotlin.fir.transformSingle
-import org.jetbrains.kotlin.fir.types.ConeClassType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.FirResolvedType
-import org.jetbrains.kotlin.fir.types.FirType
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeImpl
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.compose
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 
-class FirTypeResolveTransformer : FirTransformer<Nothing?>() {
+class FirTypeResolveTransformer(val superTypesOnly: Boolean = false) : FirTransformer<Nothing?>() {
     override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
-        return (element.transformChildren(this, data) as E).compose()
+        return if (superTypesOnly) {
+            element.compose()
+        } else {
+            (element.transformChildren(this, data) as E).compose()
+        }
     }
 
     lateinit var scope: FirCompositeScope
@@ -46,13 +47,21 @@ class FirTypeResolveTransformer : FirTransformer<Nothing?>() {
             )
         )
         packageFqName = file.packageFqName
-        return file.also { it.transformChildren(this, null) }.compose()
+        return super.transformFile(file, data)
+    }
+
+    private fun lookupSuperTypes(klass: FirClass): List<ClassId> {
+        val superTypesBuilder = SuperClassHierarchyBuilder()
+        klass.superTypes.any { it.accept(superTypesBuilder, null) }
+        return superTypesBuilder.classes
     }
 
     override fun transformClass(klass: FirClass, data: Nothing?): CompositeTransformResult<FirDeclaration> {
         klass.transformChildren(SuperTypeResolver(), null)
         className = className.child(klass.name)
+
         scope.scopes += FirNestedClassifierScope(ClassId(packageFqName, className, false), klass.session)
+        scope.scopes.addAll(lookupSuperTypes(klass).map { FirNestedClassifierScope(it, klass.session) })
         val result = super.transformClass(klass, data)
         scope.scopes.also { it.removeAt(it.lastIndex) }
         className = className.parent()
@@ -91,10 +100,53 @@ class FirTypeResolveTransformer : FirTransformer<Nothing?>() {
 
             val classId = (transformedType.type as? ConeClassType)?.fqName ?: return transformedType.compose()
             val firProvider = FirProvider.getInstance(transformedType.session)
-            // ???
+
-            firProvider.getFirClassifierByFqName(classId)!!.transformChildren(this, data)
+            val classes = generateSequence(classId) { it.outerClassId }.toList()
+
+            val transformer = FirTypeResolveTransformer(superTypesOnly = true)
+
+            val file = firProvider.getFirClassifierContainerFile(classes.last())
+
+            file.transformSingle(transformer, null)
+
+            classes.forEach {
+                firProvider.getFirClassifierByFqName(it)!!.transformSingle(transformer, data)
+            }
 
             return transformedType.compose()
         }
     }
+
+    private inner class SuperClassHierarchyBuilder : FirVisitor<Boolean, Nothing?>() {
+
+        override fun visitElement(element: FirElement, data: Nothing?): Boolean {
+            return false
+        }
+
+        private tailrec fun ConeClassType.computePartialExpansion(): ClassId {
+            return when (this) {
+                !is ConeAbbreviatedType -> this.fqName
+                else -> (this.directExpansion as ConeClassType).computePartialExpansion()
+            }
+        }
+
+        val classes = mutableListOf<ClassId>()
+
+        override fun visitResolvedType(resolvedType: FirResolvedType, data: Nothing?): Boolean {
+            val provider = FirProvider.getInstance(resolvedType.session)
+            val targetClassId = resolvedType.coneTypeSafe<ConeClassType>()?.computePartialExpansion() ?: return false
+            val classifier = provider.getFirClassifierByFqName(targetClassId)!!
+            when (classifier) {
+                is FirClass -> {
+                    if (classifier.classKind == ClassKind.CLASS) {
+                        classes += targetClassId
+                        classifier.superTypes.any { it.accept(this, data) }
+                    }
+                }
+                is FirTypeAlias -> classifier.abbreviatedType.accept(this, data)
+            }
+            return true
+        }
+
+    }
 }

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
@@ -19,6 +20,7 @@ interface FirProvider {
 
     fun getFirTypeParameterByFqName(fqName: ClassId, parameterName: Name): FirTypeParameter?
 
+    fun getFirClassifierContainerFile(fqName: ClassId): FirFile
 
     companion object {
         fun getInstance(session: FirSession): FirProvider = session.service()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/impl/FirProviderImpl.kt

@@ -15,6 +15,9 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
 class FirProviderImpl(val session: FirSession) : FirProvider {
+    override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
+        return classifierContainerFileMap[fqName] ?: error("Couldn't find container for $fqName")
+    }
 
     fun recordFile(file: FirFile) {
         val packageName = file.packageFqName
@@ -27,7 +30,9 @@ class FirProviderImpl(val session: FirSession) : FirProvider {
 
             override fun visitClass(klass: FirClass) {
                 val fqName = containerFqName.child(klass.name)
-                classifierMap[ClassId(packageName, fqName, false)] = klass
+                val classId = ClassId(packageName, fqName, false)
+                classifierMap[classId] = klass
+                classifierContainerFileMap[classId] = file
 
                 containerFqName = fqName
                 klass.acceptChildren(this)
@@ -36,13 +41,16 @@ class FirProviderImpl(val session: FirSession) : FirProvider {
 
             override fun visitTypeAlias(typeAlias: FirTypeAlias) {
                 val fqName = containerFqName.child(typeAlias.name)
-                classifierMap[ClassId(packageName, fqName, false)] = typeAlias
+                val classId = ClassId(packageName, fqName, false)
+                classifierMap[classId] = typeAlias
+                classifierContainerFileMap[classId] = file
             }
         })
     }
 
     private val fileMap = mutableMapOf<FqName, List<FirFile>>()
     private val classifierMap = mutableMapOf<ClassId, FirMemberDeclaration>()
+    private val classifierContainerFileMap = mutableMapOf<ClassId, FirFile>()
 
     override fun getFirFilesByPackage(fqName: FqName): List<FirFile> {
         return fileMap[fqName].orEmpty()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirNestedClassifierScope.kt

@@ -6,16 +6,9 @@
 package org.jetbrains.kotlin.fir.scopes.impl
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.types.ConeAbbreviatedType
-import org.jetbrains.kotlin.fir.types.ConeClassType
-import org.jetbrains.kotlin.fir.types.FirResolvedType
-import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
@@ -23,40 +16,12 @@ class FirNestedClassifierScope(val classId: ClassId, val session: FirSession) :
 
     private val firProvider = FirProvider.getInstance(session)
 
-    fun ClassId.getFir(): FirMemberDeclaration? {
+    private fun ClassId.getFir(): FirMemberDeclaration? {
         return firProvider.getFirClassifierByFqName(this)
     }
 
-    private tailrec fun ConeClassType.computePartialExpansion(): ClassId {
-        return when (this) {
-            !is ConeAbbreviatedType -> this.fqName
-            else -> (this.directExpansion as ConeClassType).computePartialExpansion()
-        }
-    }
-
-    private val superScopes by lazy {
-        val self = classId.getFir()
-        when (self) {
-            is FirClass -> {
-                val superTypes = self.superTypes as List<FirResolvedType>
-                FirCompositeScope(superTypes.mapTo(ArrayList(superTypes.size)) {
-                    FirNestedClassifierScope(it.coneTypeUnsafe<ConeClassType>().computePartialExpansion(), session)
-                })
-            }
-            is FirTypeAlias -> {
-                val expansionTarget = self.abbreviatedType.coneTypeUnsafe<ConeClassType>().computePartialExpansion()
-                FirNestedClassifierScope(expansionTarget, session)
-            }
-            else -> error("!")
-        }
-    }
-
     override fun processClassifiersByName(name: Name, processor: (ClassId) -> Boolean): Boolean {
         val child = ClassId(classId.packageFqName, classId.relativeClassName.child(name), false)
-        if (child.getFir() != null && !processor(child)) return false
+        return child.getFir() == null || processor(child)
-
-        return superScopes.processClassifiersByName(name, processor)
     }
-
-
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirResolvedType.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
-import org.jetbrains.kotlin.types.KotlinType
 
 interface FirResolvedType : FirTypeWithNullability {
     val type: ConeKotlinType
@@ -15,4 +14,5 @@ interface FirResolvedType : FirTypeWithNullability {
         visitor.visitResolvedType(this, data)
 }
 
-inline fun <reified T> FirType.coneTypeUnsafe() = (this as FirResolvedType).type as T
+inline fun <reified T : ConeKotlinType> FirType.coneTypeUnsafe() = (this as FirResolvedType).type as T
+inline fun <reified T : ConeKotlinType> FirType.coneTypeSafe() = (this as FirResolvedType).type as? T

compiler/testData/fir/resolve/multifile/NestedSuperType.1.kt

@@ -0,0 +1,7 @@
+package b
+
+import c.C
+
+open class B : C() {
+    open class NestedInB : NestedInC()
+}

compiler/testData/fir/resolve/multifile/NestedSuperType.2.kt

@@ -0,0 +1,5 @@
+package c
+
+open class C {
+    open class NestedInC
+}

compiler/testData/fir/resolve/multifile/NestedSuperType.kt

@@ -0,0 +1,8 @@
+package a
+
+import b.B
+
+class A : B() {
+    class NestedInA1 : NestedInB()
+    class NestedInA2 : NestedInC()
+}

compiler/testData/fir/resolve/multifile/NestedSuperType.txt

@@ -0,0 +1,9 @@
+FILE: NestedSuperType.kt
+    unknown final class A() : R/b.B/ {
+        unknown final class NestedInA1() : R/b.B.NestedInB/ {
+        }
+
+        unknown final class NestedInA2() : R/c.C.NestedInC/ {
+        }
+
+    }

compiler/tests/org/jetbrains/kotlin/fir/FirResolveTestCaseGenerated.java

@@ -69,6 +69,12 @@ public class FirResolveTestCaseGenerated extends AbstractFirResolveTestCase {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/fir/resolve/multifile"), Pattern.compile("^([^.]+)\\.kt$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("NestedSuperType.kt")
+        public void testNestedSuperType() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/fir/resolve/multifile/NestedSuperType.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("simpleAliasedImport.kt")
         public void testSimpleAliasedImport() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/fir/resolve/multifile/simpleAliasedImport.kt");