FIR: Add FirTypeScope

^KT-35495 Relates

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.inferenceContext
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassSubstitutionScope
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
@@ -43,7 +44,7 @@ inline fun <reified ID : Any, reified FS : FirScope> scopeSessionKey(): ScopeSes
     return object : ScopeSessionKey<ID, FS>() {}
 }
 
-val USE_SITE = scopeSessionKey<FirClassSymbol<*>, FirScope>()
+val USE_SITE = scopeSessionKey<FirClassSymbol<*>, FirTypeScope>()
 
 data class SubstitutionScopeKey(val type: ConeClassLikeType) : ScopeSessionKey<FirClassLikeSymbol<*>, FirClassSubstitutionScope>()
 
@@ -79,11 +80,11 @@ fun createSubstitution(
 
 fun ConeClassLikeType.wrapSubstitutionScopeIfNeed(
     session: FirSession,
-    useSiteMemberScope: FirScope,
+    useSiteMemberScope: FirTypeScope,
     declaration: FirClassLikeDeclaration<*>,
     builder: ScopeSession,
     derivedClassId: ClassId? = null
-): FirScope {
+): FirTypeScope {
     if (this.typeArguments.isEmpty()) return useSiteMemberScope
     return builder.getOrBuild(declaration.symbol, SubstitutionScopeKey(this)) {
         val typeParameters = (declaration as? FirTypeParameterRefsOwner)?.typeParameters.orEmpty()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/KotlinScopeProvider.kt

@@ -38,7 +38,7 @@ class KotlinScopeProvider(
         klass: FirClass<*>,
         useSiteSession: FirSession,
         scopeSession: ScopeSession
-    ): FirScope {
+    ): FirTypeScope {
         return scopeSession.getOrBuild(klass.symbol, USE_SITE) {
             val declaredScope = declaredMemberScope(klass)
             val decoratedDeclaredMemberScope =
@@ -52,14 +52,18 @@ class KotlinScopeProvider(
                         symbol.fir.scope(
                             substitutor(symbol, useSiteSuperType, useSiteSession),
                             useSiteSession, scopeSession, skipPrivateMembers = true, klass.classId
-                        )
+                        ).let {
+                            it as? FirTypeScope ?: error("$it is expected to be FirOverrideAwareScope")
+                        }
                     } else {
                         null
                     }
                 }
             FirClassUseSiteMemberScope(
                 useSiteSession,
-                FirSuperTypeScope.prepareSupertypeScope(useSiteSession, FirStandardOverrideChecker(useSiteSession), scopes),
+                FirSuperTypeScope.prepareOverrideAwareSupertypeScope(
+                    useSiteSession, FirStandardOverrideChecker(useSiteSession), scopes
+                ),
                 decoratedDeclaredMemberScope
             )
         }

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

@@ -8,12 +8,15 @@ package org.jetbrains.kotlin.fir.scopes.impl
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirOverrideChecker
-import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 
-abstract class AbstractFirOverrideScope(val session: FirSession, protected val overrideChecker: FirOverrideChecker) : FirScope() {
+abstract class AbstractFirOverrideScope(
+    val session: FirSession,
+    protected val overrideChecker: FirOverrideChecker
+) : FirTypeScope() {
     //base symbol as key, overridden as value
     val overrideByBase = mutableMapOf<FirCallableSymbol<*>, FirCallableSymbol<*>?>()
 

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

@@ -14,15 +14,17 @@ import org.jetbrains.kotlin.fir.declarations.builder.FirSimpleFunctionBuilder
 import org.jetbrains.kotlin.fir.declarations.builder.FirValueParameterBuilder
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirOverrideChecker
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.Name
 
 abstract class AbstractFirUseSiteMemberScope(
     session: FirSession,
     overrideChecker: FirOverrideChecker,
-    protected val superTypesScope: FirScope,
+    protected val superTypesScope: FirTypeScope,
     protected val declaredMemberScope: FirScope
 ) : AbstractFirOverrideScope(session, overrideChecker) {
 

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

@@ -20,7 +20,8 @@ import org.jetbrains.kotlin.fir.resolve.substitution.chain
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
-import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
+import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -34,12 +35,12 @@ import org.jetbrains.kotlin.name.Name
 
 class FirClassSubstitutionScope(
     private val session: FirSession,
-    private val useSiteMemberScope: FirScope,
+    private val useSiteMemberScope: FirTypeScope,
     scopeSession: ScopeSession,
     private val substitutor: ConeSubstitutor,
     private val skipPrivateMembers: Boolean,
     private val derivedClassId: ClassId? = null
-) : FirScope() {
+) : FirTypeScope() {
 
     private val fakeOverrideFunctions = mutableMapOf<FirFunctionSymbol<*>, FirFunctionSymbol<*>>()
     private val fakeOverrideConstructors = mutableMapOf<FirConstructorSymbol, FirConstructorSymbol>()
@@ -48,7 +49,7 @@ class FirClassSubstitutionScope(
     private val fakeOverrideAccessors = mutableMapOf<FirAccessorSymbol, FirAccessorSymbol>()
 
     constructor(
-        session: FirSession, useSiteMemberScope: FirScope, scopeSession: ScopeSession,
+        session: FirSession, useSiteMemberScope: FirTypeScope, scopeSession: ScopeSession,
         substitution: Map<FirTypeParameterSymbol, ConeKotlinType>,
         skipPrivateMembers: Boolean, derivedClassId: ClassId? = null
     ) : this(session, useSiteMemberScope, scopeSession, substitutorByMap(substitution), skipPrivateMembers, derivedClassId)

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

@@ -6,13 +6,14 @@
 package org.jetbrains.kotlin.fir.scopes.impl
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.name.Name
 
 class FirClassUseSiteMemberScope(
     session: FirSession,
-    superTypesScope: FirScope,
+    superTypesScope: FirTypeScope,
     declaredMemberScope: FirScope
 ) : AbstractFirUseSiteMemberScope(session, FirStandardOverrideChecker(session), superTypesScope, declaredMemberScope) {
 

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

@@ -11,8 +11,10 @@ import org.jetbrains.kotlin.fir.declarations.FirCallableMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.FirOverrideChecker
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeFlexibleType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -28,7 +30,7 @@ import kotlin.collections.HashSet
 class FirSuperTypeScope private constructor(
     session: FirSession,
     overrideChecker: FirOverrideChecker,
-    val scopes: List<FirScope>
+    private val scopes: List<FirScope>,
 ) : AbstractFirOverrideScope(session, overrideChecker) {
 
     private val absentFunctions = mutableSetOf<Name>()
@@ -260,5 +262,13 @@ class FirSuperTypeScope private constructor(
 
             return FirSuperTypeScope(session, overrideChecker, scopes)
         }
+
+        // This methods is needed just to preserve the possibility to move
+        // org.jetbrains.kotlin.fir.scopes.FirScope.processOverriddenFunctions to FirOverrideAwareScope
+        fun prepareOverrideAwareSupertypeScope(
+            session: FirSession,
+            overrideChecker: FirOverrideChecker,
+            scopes: List<FirTypeScope>
+        ): FirTypeScope = prepareSupertypeScope(session, overrideChecker, scopes) as FirTypeScope
     }
 }