Organize FirClassDeclaredMemberScope & FirNestedClassifierScope in the similar way

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

@@ -41,7 +41,7 @@ inline fun <K, V, VA : V> MutableMap<K, V>.getOrPut(key: K, defaultValue: (K) ->
 val FirSession.firSymbolProvider: FirSymbolProvider by componentArrayAccessor()
 val FirSession.firProvider: FirProvider by componentArrayAccessor()
 val FirSession.correspondingSupertypesCache: FirCorrespondingSupertypesCache by componentArrayAccessor()
-val FirSession.declaredMemberScopeProvider: FirClassDeclaredMemberScopeProvider by componentArrayAccessor()
+val FirSession.memberScopeProvider: FirClassDeclaredMemberScopeProvider by componentArrayAccessor()
 
 fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): FirClassLikeSymbol<*>? {
     val firSymbolProvider = useSiteSession.firSymbolProvider

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

@@ -6,9 +6,19 @@
 package org.jetbrains.kotlin.fir.resolve.transformers
 
 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.declarations.isCompanion
+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.scopes.impl.FirCompositeScope
 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.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.name.ClassId
 
 abstract class FirAbstractTreeTransformerWithSuperTypes(
     phase: FirResolvePhase,
@@ -27,6 +37,32 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
         return result
     }
 
+    private fun nestedClassifierScope(classId: ClassId): FirNestedClassifierScope? {
+        val classSymbol = session.firSymbolProvider.getClassLikeSymbolByFqName(classId) as? FirClassSymbol ?: return null
+        return nestedClassifierScope(classSymbol.fir)
+    }
+
+    protected fun resolveNestedClassesSupertypes(
+        regularClass: FirRegularClass,
+        data: Nothing?
+    ): CompositeTransformResult<FirStatement> {
+        return withScopeCleanup {
+            // ? Is it Ok to use original file session here ?
+            lookupSuperTypes(regularClass, lookupInterfaces = false, deep = true, useSiteSession = session)
+                .asReversed().mapNotNullTo(towerScope.scopes) {
+                    nestedClassifierScope(it.lookupTag.classId)
+                }
+            val companionObjects = regularClass.declarations.filterIsInstance<FirRegularClass>().filter { it.isCompanion }
+            for (companionObject in companionObjects) {
+                towerScope.scopes += nestedClassifierScope(companionObject)
+            }
+            towerScope.scopes += nestedClassifierScope(regularClass)
+            regularClass.addTypeParametersScope()
+
+            transformElement(regularClass, data)
+        }
+    }
+
     protected fun FirMemberDeclaration.addTypeParametersScope() {
         val scopes = towerScope.scopes
         if (typeParameters.isNotEmpty()) {

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

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.scopes.FirPosition
 import org.jetbrains.kotlin.fir.scopes.addImportingScopes
-import org.jetbrains.kotlin.fir.scopes.impl.FirNestedClassifierScope
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -52,14 +51,6 @@ class FirSupertypeResolverTransformer : FirAbstractTreeTransformer(phase = FirRe
         return result.compose()
     }
 
-    // This and transformProperty functions are required to forbid supertype resolving for local classes
-//    override fun transformDeclarationWithBody(
-//        declarationWithBody: FirDeclarationWithBody,
-//        data: Nothing?
-//    ): CompositeTransformResult<FirDeclaration> {
-//        return declarationWithBody.compose()
-//    }
-
     override fun transformAnonymousInitializer(
         anonymousInitializer: FirAnonymousInitializer,
         data: Nothing?
@@ -215,29 +206,6 @@ class FirSupertypeResolverTransformer : FirAbstractTreeTransformer(phase = FirRe
             val fir = session.firSymbolProvider.getClassLikeSymbolByFqName(this)?.fir ?: return false
             return fir.areSupertypesComputing()
         }
-
-        private fun resolveNestedClassesSupertypes(
-            regularClass: FirRegularClass,
-            data: Nothing?
-        ): CompositeTransformResult<FirStatement> {
-            return withScopeCleanup {
-                // ? Is it Ok to use original file session here ?
-                val firProvider = FirProvider.getInstance(session)
-                val classId = regularClass.symbol.classId
-                lookupSuperTypes(regularClass, lookupInterfaces = false, deep = true, useSiteSession = session)
-                    .asReversed().mapTo(towerScope.scopes) {
-                        FirNestedClassifierScope(it.lookupTag.classId, FirSymbolProvider.getInstance(session))
-                    }
-                val companionObjects = regularClass.declarations.filterIsInstance<FirRegularClass>().filter { it.isCompanion }
-                for (companionObject in companionObjects) {
-                    towerScope.scopes += FirNestedClassifierScope(companionObject.symbol.classId, firProvider)
-                }
-                towerScope.scopes += FirNestedClassifierScope(classId, firProvider)
-                regularClass.addTypeParametersScope()
-
-                transformElement(regularClass, data)
-            }
-        }
     }
 }
 
@@ -248,5 +216,6 @@ private fun ClassId.outerClasses() = generateSequence(this, ClassId::getOuterCla
 
 private fun FirClassLikeDeclaration<*>.areSupertypesComputed() =
     supertypesComputationStatus == SupertypesComputationStatus.COMPUTED
+
 private fun FirClassLikeDeclaration<*>.areSupertypesComputing() =
     supertypesComputationStatus == SupertypesComputationStatus.COMPUTING

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

@@ -8,13 +8,9 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 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.FirProvider
+import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.lookupSuperTypes
 import org.jetbrains.kotlin.fir.scopes.FirPosition
 import org.jetbrains.kotlin.fir.scopes.addImportingScopes
-import org.jetbrains.kotlin.fir.scopes.impl.FirNestedClassifierScope
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
@@ -44,23 +40,7 @@ open class FirTypeResolveTransformer : FirAbstractTreeTransformerWithSuperTypes(
             }
         }
 
-        return withScopeCleanup {
+        return resolveNestedClassesSupertypes(regularClass, data)
-            val session = session
-            val firProvider = FirProvider.getInstance(session)
-            val classId = regularClass.symbol.classId
-            lookupSuperTypes(regularClass, lookupInterfaces = false, deep = true, useSiteSession = session)
-                .asReversed().mapTo(towerScope.scopes) {
-                    FirNestedClassifierScope(it.lookupTag.classId, FirSymbolProvider.getInstance(session))
-                }
-            val companionObjects = regularClass.declarations.filterIsInstance<FirRegularClass>().filter { it.isCompanion }
-            for (companionObject in companionObjects) {
-                towerScope.scopes += FirNestedClassifierScope(companionObject.symbol.classId, firProvider)
-            }
-            towerScope.scopes += FirNestedClassifierScope(classId, firProvider)
-            regularClass.addTypeParametersScope()
-
-            transformDeclaration(regularClass, data) as CompositeTransformResult<FirStatement>
-        }
     }
 
     override fun transformConstructor(constructor: FirConstructor, data: Nothing?): CompositeTransformResult<FirDeclaration> {

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

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.FirCallableMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.resolve.declaredMemberScopeProvider
+import org.jetbrains.kotlin.fir.resolve.memberScopeProvider
 import org.jetbrains.kotlin.fir.scopes.FirPosition
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
@@ -20,22 +20,39 @@ import org.jetbrains.kotlin.name.Name
 
 class FirClassDeclaredMemberScopeProvider : FirSessionComponent {
 
-    val cache = mutableMapOf<FirRegularClass, FirClassDeclaredMemberScope>()
+    private val declaredMemberCache = mutableMapOf<FirRegularClass, FirClassDeclaredMemberScope>()
+    private val nestedClassifierCache = mutableMapOf<FirRegularClass, FirNestedClassifierScope>()
+
     fun declaredMemberScope(klass: FirRegularClass): FirClassDeclaredMemberScope {
-        return cache.getOrPut(klass) {
+        return declaredMemberCache.getOrPut(klass) {
             FirClassDeclaredMemberScope(klass)
         }
     }
+
+    fun nestedClassifierScope(klass: FirRegularClass): FirNestedClassifierScope {
+        return nestedClassifierCache.getOrPut(klass) {
+            FirNestedClassifierScope(klass)
+        }
+    }
 }
 
 fun declaredMemberScope(klass: FirRegularClass): FirClassDeclaredMemberScope {
     return klass
         .session
-        .declaredMemberScopeProvider
+        .memberScopeProvider
         .declaredMemberScope(klass)
 }
 
+fun nestedClassifierScope(klass: FirRegularClass): FirNestedClassifierScope {
+    return klass
+        .session
+        .memberScopeProvider
+        .nestedClassifierScope(klass)
+}
+
 class FirClassDeclaredMemberScope(klass: FirRegularClass) : FirScope() {
+    private val nestedClassifierScope = nestedClassifierScope(klass)
+
     private val callablesIndex: Map<Name, List<FirCallableSymbol<*>>> = run {
         val result = mutableMapOf<Name, MutableList<FirCallableSymbol<*>>>()
         for (declaration in klass.declarations) {
@@ -55,15 +72,6 @@ class FirClassDeclaredMemberScope(klass: FirRegularClass) : FirScope() {
         }
         result
     }
-    private val classIndex: Map<Name, FirClassSymbol> = run {
-        val result = mutableMapOf<Name, FirClassSymbol>()
-        for (declaration in klass.declarations) {
-            if (declaration is FirRegularClass) {
-                result[declaration.name] = declaration.symbol
-            }
-        }
-        result
-    }
 
     override fun processFunctionsByName(name: Name, processor: (FirFunctionSymbol<*>) -> ProcessorAction): ProcessorAction {
         val symbols = callablesIndex[name] ?: emptyList()
@@ -89,12 +97,5 @@ class FirClassDeclaredMemberScope(klass: FirRegularClass) : FirScope() {
         name: Name,
         position: FirPosition,
         processor: (FirClassifierSymbol<*>) -> ProcessorAction
-    ): ProcessorAction {
+    ): ProcessorAction = nestedClassifierScope.processClassifiersByName(name, position, processor)
-        val matchedClass = classIndex[name]
-        if (matchedClass != null && !processor(matchedClass)) {
-            return STOP
-        }
-
-        return super.processClassifiersByName(name, position, processor)
-    }
 }

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

@@ -5,27 +5,32 @@
 
 package org.jetbrains.kotlin.fir.scopes.impl
 
-import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.scopes.FirPosition
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
-class FirNestedClassifierScope(
+class FirNestedClassifierScope(val klass: FirRegularClass) : FirScope() {
-    val classId: ClassId,
+
-    private val symbolProvider: FirSymbolProvider
+    private val classIndex: Map<Name, FirClassSymbol> = run {
-) : FirScope() {
+        val result = mutableMapOf<Name, FirClassSymbol>()
+        for (declaration in klass.declarations) {
+            if (declaration is FirRegularClass) {
+                result[declaration.name] = declaration.symbol
+            }
+        }
+        result
+    }
 
     override fun processClassifiersByName(
         name: Name,
         position: FirPosition,
         processor: (FirClassifierSymbol<*>) -> ProcessorAction
     ): ProcessorAction {
-        val child = classId.createNestedClassId(name)
+        val matchedClass = classIndex[name] ?: return ProcessorAction.NONE
-        val symbol = symbolProvider.getClassLikeSymbolByFqName(child) ?: return ProcessorAction.NONE
+        return processor(matchedClass)
-
-        return processor(symbol)
     }
 }