Unite different tower contexts in FirTowerDataContextsForClassParts

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

@@ -253,7 +253,7 @@ class FirCallResolver(
         val towerDataContext =
             transformer.context.towerDataContextForCallableReferences[callableReferenceAccess] ?: transformer.context.towerDataContext
 
-        val result = transformer.context.withTowerDataContext(towerDataContext) {
+        val result = transformer.context.withSpecialTowerDataContext(towerDataContext) {
             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.resolve.FirTowerDataMode.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
@@ -27,6 +28,7 @@ import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.*
 
 interface SessionHolder {
     val session: FirSession
@@ -141,13 +143,56 @@ fun FirScope.asTowerDataElement(isLocal: Boolean): FirTowerDataElement =
 
 typealias TowerDataContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol, FirTowerDataContext>
 
+enum class FirTowerDataMode {
+    MEMBER_DECLARATION,
+    NESTED_CLASS,
+    COMPANION_OBJECT,
+    CONSTRUCTOR_HEADER,
+    SPECIAL,
+}
+
 class FirTowerDataContextsForClassParts(
-    val forNestedClasses: FirTowerDataContext,
+    forMemberDeclarations: FirTowerDataContext,
-    val forCompanionObject: FirTowerDataContext,
+    forNestedClasses: FirTowerDataContext? = null,
-    val forConstructorHeaders: FirTowerDataContext,
+    forCompanionObject: FirTowerDataContext? = null,
-    val primaryConstructorPureParametersScope: FirLocalScope?,
+    forConstructorHeaders: FirTowerDataContext? = null,
-    val primaryConstructorAllParametersScope: FirLocalScope?,
+    val primaryConstructorPureParametersScope: FirLocalScope? = null,
-)
+    val primaryConstructorAllParametersScope: FirLocalScope? = null,
+) {
+    private val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
+
+    init {
+        modeMap[MEMBER_DECLARATION] = forMemberDeclarations
+        modeMap[NESTED_CLASS] = forNestedClasses
+        modeMap[COMPANION_OBJECT] = forCompanionObject
+        modeMap[CONSTRUCTOR_HEADER] = forConstructorHeaders
+    }
+
+    var mode: FirTowerDataMode = MEMBER_DECLARATION
+
+    val forMemberDeclaration: FirTowerDataContext get() = modeMap.getValue(MEMBER_DECLARATION)
+    val forNestedClasses: FirTowerDataContext get() = modeMap.getValue(NESTED_CLASS)
+    val forCompanionObject: FirTowerDataContext get() = modeMap.getValue(COMPANION_OBJECT)
+    val forConstructorHeaders: FirTowerDataContext get() = modeMap.getValue(CONSTRUCTOR_HEADER)
+
+    var special: FirTowerDataContext
+        get() = modeMap.getValue(SPECIAL)
+        set(value) {
+            mode = SPECIAL
+            modeMap[SPECIAL] = value
+        }
+    var currentContext: FirTowerDataContext
+        get() = modeMap.getValue(mode)
+        set(value) {
+            modeMap[mode] = value
+        }
+
+    fun withMemberDeclarationContext(newContextForMemberDeclarations: FirTowerDataContext): FirTowerDataContextsForClassParts =
+        FirTowerDataContextsForClassParts(
+            newContextForMemberDeclarations, modeMap[NESTED_CLASS], modeMap[COMPANION_OBJECT], modeMap[CONSTRUCTOR_HEADER],
+            primaryConstructorPureParametersScope, primaryConstructorAllParametersScope
+        )
+}
 
 // --------------------------------------- Utils ---------------------------------------
 

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

@@ -223,7 +223,7 @@ class FirCallCompleter(
             }
 
             val localContext = components.towerDataContextForAnonymousFunctions[lambdaArgument.symbol] ?: error("")
-            transformer.context.withTowerDataContext(localContext) {
+            transformer.context.withSpecialTowerDataContext(localContext) {
                 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

@@ -10,10 +10,7 @@ import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContextsForClassParts
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataElement
-import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
@@ -25,7 +22,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.utils.sure
 
 class BodyResolveContext(
     val returnTypeCalculator: ReturnTypeCalculator,
@@ -41,11 +37,12 @@ class BodyResolveContext(
 
     val implicitReceiverStack: ImplicitReceiverStack get() = towerDataContext.implicitReceiverStack
 
-    @set:PrivateForInline
+    val towerDataContext: FirTowerDataContext
-    var towerDataContext: FirTowerDataContext = FirTowerDataContext()
+        get() = towerDataContextsForClassParts.currentContext
 
     @set:PrivateForInline
-    var towerDataContextsForClassParts: FirTowerDataContextsForClassParts? = null
+    var towerDataContextsForClassParts: FirTowerDataContextsForClassParts =
+        FirTowerDataContextsForClassParts(forMemberDeclarations = FirTowerDataContext())
 
     val containerIfAny: FirDeclaration?
         get() = containers.lastOrNull()
@@ -84,23 +81,11 @@ class BodyResolveContext(
         }
     }
 
-    fun getTowerDataContextForStaticNestedClassesUnsafe(): FirTowerDataContext =
-        firTowerDataContextsForClassParts().forNestedClasses
-
-    fun getTowerDataContextForCompanionUnsafe(): FirTowerDataContext =
-        firTowerDataContextsForClassParts().forCompanionObject
-
-    fun getTowerDataContextForConstructorResolution(): FirTowerDataContext =
-        firTowerDataContextsForClassParts().forConstructorHeaders
-
     fun getPrimaryConstructorPureParametersScope(): FirLocalScope? =
-        towerDataContextsForClassParts?.primaryConstructorPureParametersScope
+        towerDataContextsForClassParts.primaryConstructorPureParametersScope
 
     fun getPrimaryConstructorAllParametersScope(): FirLocalScope? =
-        towerDataContextsForClassParts?.primaryConstructorAllParametersScope
+        towerDataContextsForClassParts.primaryConstructorAllParametersScope
-
-    private fun firTowerDataContextsForClassParts() =
-        towerDataContextsForClassParts.sure { "towerDataContextForStaticNestedClasses should not be null" }
 
     @OptIn(PrivateForInline::class)
     inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
@@ -113,20 +98,36 @@ class BodyResolveContext(
         }
     }
 
-    inline fun <T> withTowerDataContext(context: FirTowerDataContext, f: () -> T): T {
-        return withTowerDataCleanup {
-            replaceTowerDataContext(context)
-            f()
-        }
-    }
-
     @OptIn(PrivateForInline::class)
     inline fun <R> withTowerDataCleanup(l: () -> R): R {
         val initialContext = towerDataContext
         return try {
             l()
         } finally {
-            towerDataContext = initialContext
+            replaceTowerDataContext(initialContext)
+        }
+    }
+
+    inline fun <T> withTowerDataMode(mode: FirTowerDataMode, f: () -> T): T {
+        return withTowerModeCleanup {
+            towerDataMode = mode
+            f()
+        }
+    }
+
+    inline fun <T> withSpecialTowerDataContext(context: FirTowerDataContext, f: () -> T): T {
+        return withTowerModeCleanup {
+            towerDataContextsForClassParts.special = context
+            f()
+        }
+    }
+
+    inline fun <R> withTowerModeCleanup(l: () -> R): R {
+        val initialMode = towerDataMode
+        return try {
+            l()
+        } finally {
+            towerDataMode = initialMode
         }
     }
 
@@ -140,9 +141,15 @@ class BodyResolveContext(
         }
     }
 
+    var towerDataMode: FirTowerDataMode
+        get() = towerDataContextsForClassParts.mode
+        set(value) {
+            towerDataContextsForClassParts.mode = value
+        }
+
     @OptIn(PrivateForInline::class)
     fun replaceTowerDataContext(newContext: FirTowerDataContext) {
-        towerDataContext = newContext
+        towerDataContextsForClassParts.currentContext = newContext
     }
 
     fun addNonLocalTowerDataElement(element: FirTowerDataElement) {
@@ -199,12 +206,13 @@ class BodyResolveContext(
     fun createSnapshotForLocalClasses(
         returnTypeCalculator: ReturnTypeCalculator,
         targetedLocalClasses: Set<FirClass<*>>
-    ): BodyResolveContext = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
+    ): BodyResolveContext =
-        file = this@BodyResolveContext.file
+        BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
-        towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
+            file = this@BodyResolveContext.file
-        towerDataContextForCallableReferences.putAll(this@BodyResolveContext.towerDataContextForCallableReferences)
+            towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
-        containers = this@BodyResolveContext.containers
+            towerDataContextForCallableReferences.putAll(this@BodyResolveContext.towerDataContextForCallableReferences)
-        towerDataContext = this@BodyResolveContext.towerDataContext
+            containers = this@BodyResolveContext.containers
-        anonymousFunctionsAnalyzedInDependentContext.addAll(this@BodyResolveContext.anonymousFunctionsAnalyzedInDependentContext)
+            replaceTowerDataContext(this@BodyResolveContext.towerDataContext)
-    }
+            anonymousFunctionsAnalyzedInDependentContext.addAll(this@BodyResolveContext.anonymousFunctionsAnalyzedInDependentContext)
+        }
 }

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

@@ -51,6 +51,14 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         }
     }
 
+    protected inline fun <T> withLocalScope(localScope: FirLocalScope?, crossinline l: () -> T): T {
+        if (localScope == null) return l()
+        return context.withTowerDataCleanup {
+            addLocalScope(localScope)
+            l()
+        }
+    }
+
     protected fun addNewLocalScope() {
         context.addLocalScope(FirLocalScope())
     }

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

@@ -115,7 +115,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     }
 
     override fun transformEnumEntry(enumEntry: FirEnumEntry, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        context.withTowerDataContext(context.getTowerDataContextForConstructorResolution()) {
+        context.withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
             return (enumEntry.transformChildren(this, data) as FirEnumEntry).compose()
         }
     }
@@ -185,7 +185,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         return withFullBodyResolve {
             withLocalScopeCleanup {
                 val primaryConstructorParametersScope = context.getPrimaryConstructorAllParametersScope()
-                context.withTowerDataContext(context.getTowerDataContextForConstructorResolution()) {
+                context.withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
                     context.withContainer(field) {
                         withLocalScopeCleanup {
                             addLocalScope(primaryConstructorParametersScope)
@@ -400,15 +400,13 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             return regularClass.runAllPhasesForLocalClass(transformer, components, data).compose()
         }
 
-        return context.withTowerDataCleanup {
+        return context.withTowerModeCleanup {
             if (!regularClass.isInner && context.containerIfAny is FirRegularClass) {
-                context.replaceTowerDataContext(
+                if (regularClass.isCompanion) {
-                    if (regularClass.isCompanion) {
+                    context.towerDataMode = FirTowerDataMode.COMPANION_OBJECT
-                        context.getTowerDataContextForCompanionUnsafe()
+                } else {
-                    } else {
+                    context.towerDataMode = FirTowerDataMode.NESTED_CLASS
-                        context.getTowerDataContextForStaticNestedClassesUnsafe()
+                }
-                    }
-                )
             }
 
             doTransformRegularClass(regularClass, data)
@@ -642,7 +640,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             /*
              * Default values of constructor can't access members of constructing class
              */
-            context.withTowerDataContext(context.getTowerDataContextForConstructorResolution()) {
+            context.withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
                 if (owningClass != null && !constructor.isPrimary) {
                     context.addReceiver(
                         null,
@@ -669,8 +667,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
              * Delegated constructor call is called before constructor body, so we need to
              *   analyze it before body, so body can access smartcasts from that call
              */
-            context.withTowerDataCleanup {
+            withLocalScope(scopeWithValueParameters) {
-                addLocalScope(scopeWithValueParameters)
                 constructor.transformDelegatedConstructor(transformer, data)
             }
 
@@ -681,7 +678,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                      *   In it's body we don't have this receiver for building class, so we need to use
                      *   special towerDataContext
                      */
-                    context.withTowerDataContext(context.getTowerDataContextForConstructorResolution()) {
+                    context.withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
                         addLocalScope(scopeWithValueParameters)
                         constructor.transformBody(transformer, data)
                     }
@@ -914,7 +911,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         owner: FirClass<*>,
         type: ConeKotlinType,
         block: () -> T
-    ): T = context.withTowerDataCleanup {
+    ): T {
         val towerElementsForClass = components.collectTowerDataElementsForClass(owner, type)
 
         val base = context.towerDataContext.addNonLocalTowerDataElements(towerElementsForClass.superClassesStaticsAndCompanionReceivers)
@@ -952,9 +949,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 null to null
             }
 
-        components.context.replaceTowerDataContext(forMembersResolution)
-
         val newContexts = FirTowerDataContextsForClassParts(
+            forMembersResolution,
             newTowerDataContextForStaticNestedClasses,
             statics,
             scopeForConstructorHeader,
@@ -962,7 +958,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             primaryConstructorAllParametersScope
         )
 
-        context.withNewTowerDataForClassParts(newContexts) {
+        return context.withNewTowerDataForClassParts(newContexts) {
             block()
         }
     }

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

@@ -827,24 +827,21 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         var result = delegatedConstructorCall
         try {
             val lastDispatchReceiver = implicitReceiverStack.lastDispatchReceiver()
-            context.withTowerDataCleanup {
-                if ((context.containerIfAny as? FirConstructor)?.isPrimary == true) {
-                    context.replaceTowerDataContext(context.getTowerDataContextForConstructorResolution())
-                    context.getPrimaryConstructorAllParametersScope()?.let(context::addLocalScope)
-                }
 
-                // it's just a constructor parameters scope created in
+            // because there's a `context.saveContextForAnonymousFunction(anonymousFunction)`
-                // `FirDeclarationResolveTransformer::doTransformConstructor()`
+            // call inside of the FirDeclarationResolveTransformer and accessing `this`
-                val parametersScope = context.towerDataContext.localScopes.lastOrNull()
+            // inside a lambda which is a value parameter of a constructor delegate
-
+            // is prohibited
-                // because there's a `context.saveContextForAnonymousFunction(anonymousFunction)`
+            context.withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
-                // call inside of the FirDeclarationResolveTransformer and accessing `this`
+                context.withTowerDataCleanup {
-                // inside a lambda which is a value parameter of a constructor delegate
+                    if ((context.containerIfAny as? FirConstructor)?.isPrimary == true) {
-                // is prohibited
+                        context.getPrimaryConstructorAllParametersScope()?.let(context::addLocalScope)
-                context.withTowerDataContext(context.getTowerDataContextForConstructorResolution()) {
-                    parametersScope?.let {
-                        addLocalScope(it)
                     }
+
+                    // it's just a constructor parameters scope created in
+                    // `FirDeclarationResolveTransformer::doTransformConstructor()`
+                    context.towerDataContextsForClassParts.forMemberDeclaration.localScopes.lastOrNull()?.let(context::addLocalScope)
+
                     if (containingClass is FirRegularClass && !containingConstructor.isPrimary) {
                         context.addReceiver(
                             null,