FIR: move lambda / reference specific contexts to general context parts

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

@@ -250,10 +250,7 @@ class FirCallResolver(
         // No reset here!
         val localCollector = CandidateCollector(components, components.resolutionStageRunner)
 
-        val towerDataContext =
-            transformer.context.towerDataContextForCallableReferences[callableReferenceAccess] ?: transformer.context.towerDataContext
-
-        val result = transformer.context.withSpecialTowerDataContext(towerDataContext) {
+        val result = transformer.context.withCallableReferenceTowerDataContext(callableReferenceAccess) {
             towerResolver.runResolver(
                 info,
                 transformer.resolutionContext,

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

@@ -10,6 +10,7 @@ import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataMode.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
@@ -43,7 +44,6 @@ abstract class BodyResolveComponents : SessionHolder {
     abstract val towerDataElements: List<FirTowerDataElement>
     abstract val towerDataContext: FirTowerDataContext
     abstract val localScopes: FirLocalScopes
-    abstract val towerDataContextForAnonymousFunctions: TowerDataContextForAnonymousFunctions
     abstract val noExpectedType: FirTypeRef
     abstract val symbolProvider: FirSymbolProvider
     abstract val file: FirFile
@@ -175,9 +175,12 @@ class FirTowerDataContextsForClassParts(
     val forCompanionObject: FirTowerDataContext get() = modeMap.getValue(COMPANION_OBJECT)
     val forConstructorHeaders: FirTowerDataContext get() = modeMap.getValue(CONSTRUCTOR_HEADER)
 
+    val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
+    val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext> = mutableMapOf()
+
     var special: FirTowerDataContext
         get() = modeMap.getValue(SPECIAL)
-        set(value) {
+        private set(value) {
             mode = SPECIAL
             modeMap[SPECIAL] = value
         }
@@ -187,6 +190,17 @@ class FirTowerDataContextsForClassParts(
             modeMap[mode] = value
         }
 
+    fun setAnonymousFunctionContext(symbol: FirAnonymousFunctionSymbol) {
+        special = towerDataContextForAnonymousFunctions.getValue(symbol)
+    }
+
+    fun setCallableReferenceContextIfAny(access: FirCallableReferenceAccess) {
+        val context = towerDataContextForCallableReferences[access]
+        if (context != null) {
+            special = context
+        }
+    }
+
     fun withMemberDeclarationContext(newContextForMemberDeclarations: FirTowerDataContext): FirTowerDataContextsForClassParts =
         FirTowerDataContextsForClassParts(
             newContextForMemberDeclarations, modeMap[NESTED_CLASS], modeMap[COMPANION_OBJECT], modeMap[CONSTRUCTOR_HEADER],

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -222,8 +222,7 @@ class FirCallCompleter(
                 FirBuilderInferenceSession(transformer.resolutionContext, stubsForPostponedVariables as Map<ConeTypeVariable, ConeStubType>)
             }
 
-            val localContext = components.towerDataContextForAnonymousFunctions[lambdaArgument.symbol] ?: error("")
-            transformer.context.withSpecialTowerDataContext(localContext) {
+            transformer.context.withAnonymousFunctionTowerDataContext(lambdaArgument.symbol) {
                 if (builderInferenceSession != null) {
                     transformer.context.withInferenceSession(builderInferenceSession) {
                         lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -35,23 +35,27 @@ class BodyResolveContext(
     lateinit var file: FirFile
         internal set
 
-    val implicitReceiverStack: ImplicitReceiverStack get() = towerDataContext.implicitReceiverStack
-
-    val towerDataContext: FirTowerDataContext
-        get() = towerDataContextsForClassParts.currentContext
-
     @set:PrivateForInline
     var towerDataContextsForClassParts: FirTowerDataContextsForClassParts =
         FirTowerDataContextsForClassParts(forMemberDeclarations = FirTowerDataContext())
 
-    val containerIfAny: FirDeclaration?
-        get() = containers.lastOrNull()
+    val towerDataContext: FirTowerDataContext
+        get() = towerDataContextsForClassParts.currentContext
+
+    val implicitReceiverStack: ImplicitReceiverStack
+        get() = towerDataContext.implicitReceiverStack
+
+    val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext>
+        get() = towerDataContextsForClassParts.towerDataContextForAnonymousFunctions
+
+    val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext>
+        get() = towerDataContextsForClassParts.towerDataContextForCallableReferences
 
     @set:PrivateForInline
     var containers: PersistentList<FirDeclaration> = persistentListOf()
 
-    val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
-    val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext> = mutableMapOf()
+    val containerIfAny: FirDeclaration?
+        get() = containers.lastOrNull()
 
     @set:PrivateForInline
     var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
@@ -115,9 +119,16 @@ class BodyResolveContext(
         }
     }
 
-    inline fun <T> withSpecialTowerDataContext(context: FirTowerDataContext, f: () -> T): T {
+    inline fun <T> withAnonymousFunctionTowerDataContext(symbol: FirAnonymousFunctionSymbol, f: () -> T): T {
         return withTowerModeCleanup {
-            towerDataContextsForClassParts.special = context
+            towerDataContextsForClassParts.setAnonymousFunctionContext(symbol)
+            f()
+        }
+    }
+
+    inline fun <T> withCallableReferenceTowerDataContext(access: FirCallableReferenceAccess, f: () -> T): T {
+        return withTowerModeCleanup {
+            towerDataContextsForClassParts.setCallableReferenceContextIfAny(access)
             f()
         }
     }
@@ -193,11 +204,9 @@ class BodyResolveContext(
         towerDataContextForAnonymousFunctions.remove(anonymousFunction.symbol)
     }
 
-    fun cleanContextForAnonymousFunction() {
+    fun clear() {
         towerDataContextForAnonymousFunctions.clear()
-    }
-
-    fun cleanDataFlowContext() {
+        towerDataContextForCallableReferences.clear()
         dataFlowAnalyzerContext.reset()
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt

@@ -140,7 +140,6 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         override val file: FirFile get() = context.file
         override val implicitReceiverStack: ImplicitReceiverStack get() = context.implicitReceiverStack
         override val containingDeclarations: List<FirDeclaration> get() = context.containers
-        override val towerDataContextForAnonymousFunctions: TowerDataContextForAnonymousFunctions get() = context.towerDataContextForAnonymousFunctions
         override val returnTypeCalculator: ReturnTypeCalculator get() = context.returnTypeCalculator
         override val container: FirDeclaration get() = context.containerIfAny!!
 

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

@@ -35,7 +35,6 @@ open class FirBodyResolveTransformer(
     val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve(),
     outerBodyResolveContext: BodyResolveContext? = null
 ) : FirAbstractBodyResolveTransformer(phase) {
-    private var packageFqName = FqName.ROOT
 
     final override val context: BodyResolveContext =
         outerBodyResolveContext ?: BodyResolveContext(returnTypeCalculator, DataFlowAnalyzerContext.empty(session))
@@ -50,12 +49,9 @@ open class FirBodyResolveTransformer(
 
     override fun transformFile(file: FirFile, data: ResolutionMode): CompositeTransformResult<FirFile> {
         checkSessionConsistency(file)
-        context.cleanContextForAnonymousFunction()
-        context.towerDataContextForCallableReferences.clear()
-        context.cleanDataFlowContext()
+        context.clear()
         @OptIn(PrivateForInline::class)
         context.file = file
-        packageFqName = file.packageFqName
         return withScopeCleanup(context.fileImportsScope) {
             context.withTowerDataCleanup {
                 val importingScopes = createImportingScopes(file, session, components.scopeSession)