FIR: reuse scopeSession in call resolver & implicit typed ref resolution

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

@@ -60,7 +60,7 @@ fun ConeClassLikeType.wrapSubstitutionScopeIfNeed(
             typeParameter.symbol to (typeArgument as? ConeTypedProjection)?.type
         }.filter { (_, type) -> type != null }.toMap() as Map<ConeTypeParameterSymbol, ConeKotlinType>
 
-        FirClassSubstitutionScope(session, useSiteScope, substitution)
+        FirClassSubstitutionScope(session, useSiteScope, builder, substitution)
     }
 }
 

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

@@ -95,7 +95,7 @@ private fun ConeClassLikeType.wrapSubstitutionScopeIfNeed(
             typeParameter.symbol to (typeArgument as? ConeTypedProjection)?.type
         }.filter { (_, type) -> type != null }.toMap() as Map<ConeTypeParameterSymbol, ConeKotlinType>
 
-        FirClassSubstitutionScope(session, useSiteScope, substitution)
+        FirClassSubstitutionScope(session, useSiteScope, builder, substitution)
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CallResolver.kt

@@ -30,12 +30,10 @@ import org.jetbrains.kotlin.fir.symbols.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
-import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.types.AbstractTypeChecker
@@ -191,7 +189,8 @@ abstract class SessionBasedTowerLevel(val session: FirSession) : TowerScopeLevel
 class MemberScopeTowerLevel(
     session: FirSession,
     val dispatchReceiver: ReceiverValue,
-    val implicitExtensionReceiver: ImplicitReceiverValue? = null
+    val implicitExtensionReceiver: ImplicitReceiverValue? = null,
+    val scopeSession: ScopeSession
 ) : SessionBasedTowerLevel(session) {
 
     private fun <T : ConeSymbol> processMembers(
@@ -201,7 +200,7 @@ class MemberScopeTowerLevel(
     ): ProcessorAction {
         if (implicitExtensionReceiver != null && explicitExtensionReceiver != null) return ProcessorAction.NEXT
         val extensionReceiver = implicitExtensionReceiver ?: explicitExtensionReceiver
-        val scope = dispatchReceiver.type.scope(session, ScopeSession()) ?: return ProcessorAction.NEXT
+        val scope = dispatchReceiver.type.scope(session, scopeSession) ?: return ProcessorAction.NEXT
         if (scope.processScopeMembers { candidate ->
                 if (candidate is ConeCallableSymbol && candidate.hasConsistentExtensionReceiver(extensionReceiver)) {
                     // NB: we do not check dispatchReceiverValue != null here,
@@ -214,7 +213,7 @@ class MemberScopeTowerLevel(
                 }
             }.stop()
         ) return ProcessorAction.STOP
-        val withSynthetic = FirSyntheticPropertiesScope(session, scope, ReturnTypeCalculatorWithJump(session))
+        val withSynthetic = FirSyntheticPropertiesScope(session, scope, ReturnTypeCalculatorWithJump(session, scopeSession))
         return withSynthetic.processScopeMembers { symbol ->
             output.consumeCandidate(symbol, symbol.dispatchReceiverValue(), implicitExtensionReceiver)
         }
@@ -517,29 +516,30 @@ class ExplicitReceiverTowerDataConsumer<T : ConeSymbol>(
         if (checkSkip(group, resultCollector)) return ProcessorAction.NEXT
         return when (kind) {
             TowerDataKind.EMPTY ->
-                MemberScopeTowerLevel(session, explicitReceiver).processElementsByName(
+                MemberScopeTowerLevel(session, explicitReceiver, scopeSession = candidateFactory.inferenceComponents.scopeSession)
-                    token,
+                    .processElementsByName(
-                    name,
+                        token,
-                    explicitReceiver = null,
+                        name,
-                    processor = object : TowerScopeLevel.TowerScopeLevelProcessor<T> {
+                        explicitReceiver = null,
-                        override fun consumeCandidate(
+                        processor = object : TowerScopeLevel.TowerScopeLevelProcessor<T> {
-                            symbol: T,
+                            override fun consumeCandidate(
-                            dispatchReceiverValue: ClassDispatchReceiverValue?,
+                                symbol: T,
-                            implicitExtensionReceiverValue: ImplicitReceiverValue?
+                                dispatchReceiverValue: ClassDispatchReceiverValue?,
-                        ): ProcessorAction {
+                                implicitExtensionReceiverValue: ImplicitReceiverValue?
-                            resultCollector.consumeCandidate(
+                            ): ProcessorAction {
-                                group,
+                                resultCollector.consumeCandidate(
-                                candidateFactory.createCandidate(
+                                    group,
-                                    symbol,
+                                    candidateFactory.createCandidate(
-                                    dispatchReceiverValue,
+                                        symbol,
-                                    implicitExtensionReceiverValue,
+                                        dispatchReceiverValue,
-                                    ExplicitReceiverKind.DISPATCH_RECEIVER
+                                        implicitExtensionReceiverValue,
+                                        ExplicitReceiverKind.DISPATCH_RECEIVER
+                                    )
                                 )
-                            )
+                                return ProcessorAction.NEXT
-                            return ProcessorAction.NEXT
+                            }
                         }
-                    }
+                    )
-                )
             TowerDataKind.TOWER_LEVEL -> {
                 if (token == TowerScopeLevel.Token.Objects) return ProcessorAction.NEXT
                 towerScopeLevel.processElementsByName(
@@ -659,14 +659,14 @@ class CallResolver(val typeCalculator: ReturnTypeCalculator, val components: Inf
         var group = oldGroup
         towerDataConsumer.consume(
             TowerDataKind.TOWER_LEVEL,
-            MemberScopeTowerLevel(session, implicitReceiverValue),
+            MemberScopeTowerLevel(session, implicitReceiverValue, scopeSession = components.scopeSession),
             collector, group++
         )
 
         // This is an equivalent to the old "BothTowerLevelAndImplicitReceiver"
         towerDataConsumer.consume(
             TowerDataKind.TOWER_LEVEL,
-            MemberScopeTowerLevel(session, implicitReceiverValue, implicitReceiverValue),
+            MemberScopeTowerLevel(session, implicitReceiverValue, implicitReceiverValue, components.scopeSession),
             collector, group++
         )
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/InferenceUtil.kt

@@ -46,7 +46,12 @@ open class ConeTypeVariable(name: String) : TypeVariableMarker {
     val defaultType = ConeTypeVariableType(ConeNullability.NOT_NULL, typeConstructor)
 }
 
-class InferenceComponents(val ctx: TypeSystemInferenceExtensionContextDelegate, val session: FirSession, val returnTypeCalculator: ReturnTypeCalculator) {
+class InferenceComponents(
+    val ctx: TypeSystemInferenceExtensionContextDelegate,
+    val session: FirSession,
+    val returnTypeCalculator: ReturnTypeCalculator,
+    val scopeSession: ScopeSession
+) {
     private val approximator = object : AbstractTypeApproximator(ctx) {
         override fun createErrorType(message: String): SimpleTypeMarker {
             return ConeClassErrorType(message)

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

@@ -44,7 +44,10 @@ import org.jetbrains.kotlin.resolve.calls.results.TypeSpecificityComparator
 import org.jetbrains.kotlin.types.model.*
 import org.jetbrains.kotlin.utils.addIfNotNull
 
-open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOnly: Boolean) : FirTransformer<Any?>() {
+open class FirBodyResolveTransformer(
+    val session: FirSession, val implicitTypeOnly: Boolean,
+    val scopeSession: ScopeSession = ScopeSession()
+) : FirTransformer<Any?>() {
 
     val symbolProvider = session.service<FirSymbolProvider>()
 
@@ -183,8 +186,6 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         return result
     }
 
-    val scopeSession = ScopeSession()
-
     val scopes = mutableListOf<FirScope>()
     private val localScopes = mutableListOf<FirLocalScope>()
 
@@ -192,7 +193,7 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
 
     private val implicitReceiverStack = mutableListOf<ImplicitReceiverValue>()
 
-    private val jump = ReturnTypeCalculatorWithJump(session)
+    private val jump = ReturnTypeCalculatorWithJump(session, scopeSession)
 
     private fun <T> storeTypeFromCallee(access: T) where T : FirQualifiedAccess, T : FirExpression {
         access.resultType = typeFromCallee(access)
@@ -243,7 +244,6 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
     }
 
 
-
     private var qualifierStack = mutableListOf<Name>()
     private var qualifierPartsToDrop = 0
 
@@ -520,7 +520,7 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
     }
 
     private val noExpectedType = FirImplicitTypeRefImpl(session, null)
-    private val inferenceComponents = inferenceComponents(session, jump)
+    private val inferenceComponents = inferenceComponents(session, jump, scopeSession)
 
     private fun resolveCallAndSelectCandidate(functionCall: FirFunctionCall, expectedTypeRef: FirTypeRef?): FirFunctionCall {
 
@@ -1095,7 +1095,7 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
     }
 }
 
-private fun inferenceComponents(session: FirSession, jump: ReturnTypeCalculatorWithJump) =
+private fun inferenceComponents(session: FirSession, jump: ReturnTypeCalculatorWithJump, scopeSession: ScopeSession) =
     InferenceComponents(object : ConeInferenceContext, TypeSystemInferenceExtensionContextDelegate {
         override fun findCommonIntegerLiteralTypesSuperType(explicitSupertypes: List<SimpleTypeMarker>): SimpleTypeMarker? {
             //TODO wtf
@@ -1112,10 +1112,10 @@ private fun inferenceComponents(session: FirSession, jump: ReturnTypeCalculatorW
         override fun KotlinTypeMarker.removeExactAnnotation(): KotlinTypeMarker {
             return this
         }
-    }, session, jump)
+    }, session, jump, scopeSession)
 
 
-class ReturnTypeCalculatorWithJump(val session: FirSession) : ReturnTypeCalculator {
+class ReturnTypeCalculatorWithJump(val session: FirSession, val scopeSession: ScopeSession) : ReturnTypeCalculator {
 
 
     val storeType = object : FirTransformer<FirTypeRef>() {
@@ -1172,7 +1172,8 @@ class ReturnTypeCalculatorWithJump(val session: FirSession) : ReturnTypeCalculat
 
         val transformer = FirDesignatedBodyResolveTransformer(
             (listOf(file) + outerClasses.filterNotNull().asReversed() + listOf(declaration)).iterator(),
-            file.session
+            file.session,
+            scopeSession
         )
 
         file.transform<FirElement, Any?>(transformer, null)
@@ -1186,8 +1187,8 @@ class ReturnTypeCalculatorWithJump(val session: FirSession) : ReturnTypeCalculat
 }
 
 
-class FirDesignatedBodyResolveTransformer(val designation: Iterator<FirElement>, session: FirSession) :
+class FirDesignatedBodyResolveTransformer(val designation: Iterator<FirElement>, session: FirSession, scopeSession: ScopeSession) :
-    FirBodyResolveTransformer(session, implicitTypeOnly = true) {
+    FirBodyResolveTransformer(session, implicitTypeOnly = true, scopeSession = scopeSession) {
 
     override fun <E : FirElement> transformElement(element: E, data: Any?): CompositeTransformResult<E> {
         if (designation.hasNext()) {

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.FirNamedFunction
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirMemberFunctionImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorWithJump
 import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
@@ -27,6 +28,7 @@ import org.jetbrains.kotlin.name.Name
 class FirClassSubstitutionScope(
     private val session: FirSession,
     private val useSiteScope: FirScope,
+    scopeSession: ScopeSession,
     substitution: Map<ConeTypeParameterSymbol, ConeKotlinType>
 ) : FirScope() {
 
@@ -49,7 +51,7 @@ class FirClassSubstitutionScope(
         return useSiteScope.processPropertiesByName(name, processor)
     }
 
-    private val typeCalculator by lazy { ReturnTypeCalculatorWithJump(session) }
+    private val typeCalculator by lazy { ReturnTypeCalculatorWithJump(session, scopeSession) }
 
     private fun ConeKotlinType.substitute(): ConeKotlinType? {
         return substitutor.substituteOrNull(this)