FIR: Refactor scope structure on body resolve stage

Before this change:
- Local scopes, implicit receivers and type parameter scopes were separated.
- Static scopes for super classes were not present as a concept.
Instead of them, all-inherited-static for the current class has been used.
- During call resolution we were processing implicit receivers first and then non-local scopes,
while we should process them in the order of their syntax appearance from the closest
to the most distant

All these facts affect semantics (see test data changed here and the following commits)

The architecture changes are the following:
- FirTowerDataElement introduced as tower level that is used in resolution
  (effectively it's a union type between scope and implicit receiver + isLocal flag)
- FirTowerDataContext introduced for sake of encapsulation of tower data elements' list
  (it also has redundant implicitReceiverStack and localScopes)
- For each regular class we collect relevant tower data elements and add them to the current context
- Also, we preserve a special tower data context for static entities of the class
  (it doesn't have class' dispatch receiver and generic parameters)

compiler/fir/analysis-tests/testData/resolve/callResolution/companionVsSuperStatic.kt

+// FILE: JavaClass.java
+
+public class JavaClass {
+    public static String FIELD = "";
+}
+
+// FILE: main.kt
+
+class A : JavaClass() {
+    companion object {
+        val FIELD = 1
+    }
+
+    fun foo() {
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>FIELD<!>
+    }
+}

compiler/fir/analysis-tests/testData/resolve/callResolution/companionVsSuperStatic.txt

+FILE: main.kt
+    public final class A : R|JavaClass| {
+        public constructor(): R|A| {
+            super<R|JavaClass|>()
+        }
+
+        public final companion object Companion : R|kotlin/Any| {
+            private constructor(): R|A.Companion| {
+                super<R|kotlin/Any|>()
+            }
+
+            public final val FIELD: R|kotlin/Int| = Int(1)
+                public get(): R|kotlin/Int|
+
+        }
+
+        public final fun foo(): R|kotlin/Unit| {
+            this@R|/A.Companion|.R|/A.Companion.FIELD|
+        }
+
+    }

compiler/fir/analysis-tests/testData/resolve/objectInnerClass.txt

@@ -169,14 +169,14 @@ FILE: objectInnerClass.kt
             }
 
             public[local] final fun caseForBase(): R|kotlin/Unit| {
-                lval base: R|Case3.<anonymous>.Base| = R|/Case3.<anonymous>.Base.Base|(R|/B.B|())
+                lval base: R|Case3.<anonymous>.Base| = this@R|/anonymous|.R|/Case3.<anonymous>.Base.Base|(R|/B.B|())
                 R|<local>/base|.R|/Case3.<anonymous>.Base.baseFun|()
                 R|<local>/base|.R|/Case3.<anonymous>.Base.property|
                 (this@R|/anonymous|, R|<local>/base|).R|/anonymous.hoo|()
             }
 
             public[local] final fun caseForChild(): R|kotlin/Unit| {
-                lval child: R|Case3.<anonymous>.Child| = R|/Case3.<anonymous>.Child.Child|(R|/B.B|())
+                lval child: R|Case3.<anonymous>.Child| = this@R|/anonymous|.R|/Case3.<anonymous>.Child.Child|(R|/B.B|())
                 R|<local>/child|.R|/Case3.<anonymous>.Base.baseFun|()
                 R|<local>/child|.R|/Case3.<anonymous>.Base.property|
                 R|<local>/child|.R|/Case3.<anonymous>.Child.foo|()

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -589,6 +589,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/companionInvoke.kt");
         }
 
+        @TestMetadata("companionVsSuperStatic.kt")
+        public void testCompanionVsSuperStatic() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/companionVsSuperStatic.kt");
+        }
+
         @TestMetadata("debugExpressionType.kt")
         public void testDebugExpressionType() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/debugExpressionType.kt");

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -589,6 +589,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/companionInvoke.kt");
         }
 
+        @TestMetadata("companionVsSuperStatic.kt")
+        public void testCompanionVsSuperStatic() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/companionVsSuperStatic.kt");
+        }
+
         @TestMetadata("debugExpressionType.kt")
         public void testDebugExpressionType() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/debugExpressionType.kt");

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSymbolOwner
@@ -19,9 +20,10 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 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
 
 interface SessionHolder {
     val session: FirSession
@@ -33,9 +35,9 @@ interface BodyResolveComponents : SessionHolder {
     val implicitReceiverStack: ImplicitReceiverStack
     val containingDeclarations: List<FirDeclaration>
     val fileImportsScope: List<FirScope>
-    val typeParametersScopes: List<FirScope>
+    val towerDataElements: List<FirTowerDataElement>
     val localScopes: FirLocalScopes
-    val localContextForAnonymousFunctions: LocalContextForAnonymousFunctions
+    val towerDataContextForAnonymousFunctions: TowerDataContextForAnonymousFunctions
     val noExpectedType: FirTypeRef
     val symbolProvider: FirSymbolProvider
     val file: FirFile
@@ -56,14 +58,86 @@ interface BodyResolveComponents : SessionHolder {
 }
 
 typealias FirLocalScopes = PersistentList<FirLocalScope>
-typealias FirTypeParametersScopes = PersistentList<FirScope>
 
-class FirLocalContext(
-    val localScopes: FirLocalScopes,
-    val implicitReceiverStack: MutableImplicitReceiverStack
-)
+class FirTowerDataContext private constructor(
+    val towerDataElements: PersistentList<FirTowerDataElement>,
+    // These properties are effectively redundant, their content should be consistent with `towerDataElements`,
+    // i.e. implicitReceiverStack == towerDataElements.mapNotNull { it.receiver }
+    // i.e. localScopes == towerDataElements.mapNotNull { it.scope?.takeIf { it.isLocal } }
+    val implicitReceiverStack: PersistentImplicitReceiverStack,
+    val localScopes: FirLocalScopes
+) {
+
+    constructor() : this(
+        persistentListOf(),
+        PersistentImplicitReceiverStack(),
+        persistentListOf()
+    )
+
+    fun setLastLocalScope(newLastScope: FirLocalScope): FirTowerDataContext {
+        val oldLastScope = localScopes.last()
+        val indexOfLastLocalScope = towerDataElements.indexOfLast { it.scope === oldLastScope }
+
+        return FirTowerDataContext(
+            towerDataElements.set(indexOfLastLocalScope, newLastScope.asTowerDataElement(isLocal = true)),
+            implicitReceiverStack,
+            localScopes.set(localScopes.lastIndex, newLastScope)
+        )
+    }
+
+    fun addTowerDataElements(newElements: List<FirTowerDataElement>): FirTowerDataContext {
+        return FirTowerDataContext(
+            towerDataElements.addAll(newElements),
+            implicitReceiverStack.addAll(newElements.mapNotNull { it.implicitReceiver }),
+            localScopes
+        )
+    }
+
+    fun addLocalScope(localScope: FirLocalScope): FirTowerDataContext {
+        return FirTowerDataContext(
+            towerDataElements.add(localScope.asTowerDataElement(isLocal = true)),
+            implicitReceiverStack,
+            localScopes.add(localScope)
+        )
+    }
+
+    fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>): FirTowerDataContext {
+        return FirTowerDataContext(
+            towerDataElements.add(implicitReceiverValue.asTowerDataElement()),
+            implicitReceiverStack.add(name, implicitReceiverValue),
+            localScopes
+        )
+    }
+
+    fun addNonLocalScopeIfNotNull(scope: FirScope?): FirTowerDataContext {
+        if (scope == null) return this
+        return addNonLocalScope(scope)
+    }
+
+    private fun addNonLocalScope(scope: FirScope): FirTowerDataContext {
+        return FirTowerDataContext(
+            towerDataElements.add(scope.asTowerDataElement(isLocal = false)),
+            implicitReceiverStack,
+            localScopes
+        )
+    }
+}
+
+class FirTowerDataElement(val scope: FirScope?, val implicitReceiver: ImplicitReceiverValue<*>?, val isLocal: Boolean)
+
+fun ImplicitReceiverValue<*>.asTowerDataElement(): FirTowerDataElement =
+    FirTowerDataElement(scope = null, this, isLocal = false)
+
+fun FirScope.asTowerDataElement(isLocal: Boolean): FirTowerDataElement =
+    FirTowerDataElement(this, implicitReceiver = null, isLocal)
 
-typealias LocalContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol, FirLocalContext>
+typealias TowerDataContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol, FirTowerDataContext>
+
+class FirTowerDataContextsForClassParts(
+    val forNestedClasses: FirTowerDataContext,
+    val forConstructorHeaders: FirTowerDataContext,
+    val primaryConstructorParametersScope: FirLocalScope?,
+)
 
 // --------------------------------------- Utils ---------------------------------------
 
@@ -81,22 +155,10 @@ fun SessionHolder.collectImplicitReceivers(
     val implicitCompanionValues = mutableListOf<ImplicitReceiverValue<*>>()
     val implicitReceiverValue = when (owner) {
         is FirClass<*> -> {
-            // Questionable: performance
-            (owner as? FirRegularClass)?.companionObject?.let { companion ->
-                implicitCompanionValues += ImplicitDispatchReceiverValue(
-                    companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION
-                )
-            }
-            lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session).mapNotNull {
-                val superClass = (it as? ConeClassLikeType)?.lookupTag?.toSymbol(session)?.fir as? FirRegularClass
-                superClass?.companionObject?.let { companion ->
-                    implicitCompanionValues += ImplicitDispatchReceiverValue(
-                        companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION_FROM_SUPERTYPE
-                    )
-                }
-            }
-            // ---
-            ImplicitDispatchReceiverValue(owner.symbol, type, session, scopeSession)
+            val towerElementsForClass = collectTowerDataElementsForClass(owner, type)
+            implicitCompanionValues.addAll(towerElementsForClass.implicitCompanionValues)
+
+            towerElementsForClass.thisReceiver
         }
         is FirFunction<*> -> {
             ImplicitExtensionReceiverValue(owner.symbol, type, session, scopeSession)
@@ -110,3 +172,54 @@ fun SessionHolder.collectImplicitReceivers(
     }
     return ImplicitReceivers(implicitReceiverValue, implicitCompanionValues.asReversed())
 }
+
+fun SessionHolder.collectTowerDataElementsForClass(owner: FirClass<*>, defaultType: ConeKotlinType): TowerElementsForClass {
+    val allImplicitCompanionValues = mutableListOf<ImplicitReceiverValue<*>>()
+
+    val companionObject = (owner as? FirRegularClass)?.companionObject
+    val companionReceiver = companionObject?.let { companion ->
+        ImplicitDispatchReceiverValue(
+            companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION
+        )
+    }
+    allImplicitCompanionValues.addIfNotNull(companionReceiver)
+
+    val superClassesStaticsAndCompanionReceivers = mutableListOf<FirTowerDataElement>()
+    for (superType in lookupSuperTypes(owner, lookupInterfaces = false, deep = true, useSiteSession = session)) {
+        val superClass = superType.fullyExpandedType(session).lookupTag.toSymbol(session)?.fir as? FirRegularClass ?: continue
+
+        superClass.staticScope(this)?.asTowerDataElement(isLocal = false)?.let(superClassesStaticsAndCompanionReceivers::add)
+
+        (superClass as? FirRegularClass)?.companionObject?.let { companion ->
+            val superCompanionReceiver = ImplicitDispatchReceiverValue(
+                companion.symbol, session, scopeSession, kind = ImplicitDispatchReceiverKind.COMPANION_FROM_SUPERTYPE
+            )
+
+            superClassesStaticsAndCompanionReceivers += superCompanionReceiver.asTowerDataElement()
+            allImplicitCompanionValues += superCompanionReceiver
+        }
+    }
+
+    val thisReceiver = ImplicitDispatchReceiverValue(owner.symbol, defaultType, session, scopeSession)
+
+    return TowerElementsForClass(
+        thisReceiver,
+        owner.staticScope(this),
+        companionReceiver,
+        companionObject?.staticScope(this),
+        superClassesStaticsAndCompanionReceivers,
+        allImplicitCompanionValues
+    )
+}
+
+private fun FirClass<*>.staticScope(sessionHolder: SessionHolder) =
+    scopeProvider.getStaticScope(this, sessionHolder.session, sessionHolder.scopeSession)
+
+class TowerElementsForClass(
+    val thisReceiver: ImplicitReceiverValue<*>,
+    val staticScope: FirScope?,
+    val companionReceiver: ImplicitReceiverValue<*>?,
+    val companionStaticScope: FirScope?,
+    val superClassesStaticsAndCompanionReceivers: List<FirTowerDataElement>,
+    val implicitCompanionValues: List<ImplicitReceiverValue<*>>
+)

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

@@ -31,6 +31,10 @@ class PersistentImplicitReceiverStack private constructor(
         persistentListOf(),
     )
 
+    fun addAll(receivers: List<ImplicitReceiverValue<*>>): PersistentImplicitReceiverStack {
+        return receivers.fold(this) { acc, value -> acc.add(name = null, value) }
+    }
+
     fun add(name: Name?, value: ImplicitReceiverValue<*>): PersistentImplicitReceiverStack {
         val stack = stack.add(value)
         val originalTypes = originalTypes.add(value.type)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/tower/TowerGroup.kt

@@ -22,12 +22,10 @@ sealed class TowerGroupKind(private val index: Int) : Comparable<TowerGroupKind>
 
     class Local(depth: Int) : WithDepth(40, depth)
 
-    class Implicit(depth: Int) : WithDepth(50, depth)
+    class ImplicitOrNonLocal(depth: Int, val kindForDebugSake: String) : WithDepth(50, depth)
 
     object InvokeExtension : TowerGroupKind(60)
 
-    class Top(depth: Int) : WithDepth(70, depth)
-
     object Last : TowerGroupKind(Integer.MAX_VALUE)
 
     override fun compareTo(other: TowerGroupKind): Int {
@@ -38,6 +36,12 @@ sealed class TowerGroupKind(private val index: Int) : Comparable<TowerGroupKind>
         }
         return 0
     }
+
+    companion object {
+        // These two groups intentionally have the same priority
+        fun Implicit(depth: Int): TowerGroupKind = ImplicitOrNonLocal(depth, "Implicit")
+        fun NonLocal(depth: Int): TowerGroupKind = ImplicitOrNonLocal(depth, "NonLocal")
+    }
 }
 
 @Suppress("FunctionName", "unused", "PropertyName")
@@ -64,8 +68,7 @@ private constructor(
         fun Local(depth: Int) = kindOf(TowerGroupKind.Local(depth))
 
         fun Implicit(depth: Int) = kindOf(TowerGroupKind.Implicit(depth))
-
-        fun Top(depth: Int) = kindOf(TowerGroupKind.Top(depth))
+        fun NonLocal(depth: Int) = kindOf(TowerGroupKind.NonLocal(depth))
 
         fun TopPrioritized(depth: Int) = kindOf(TowerGroupKind.TopPrioritized(depth))
 
@@ -79,11 +82,10 @@ private constructor(
     fun Local(depth: Int) = kindOf(TowerGroupKind.Local(depth))
 
     fun Implicit(depth: Int) = kindOf(TowerGroupKind.Implicit(depth))
+    fun NonLocal(depth: Int) = kindOf(TowerGroupKind.NonLocal(depth))
 
     val InvokeExtension get() = kindOf(TowerGroupKind.InvokeExtension)
 
-    fun Top(depth: Int) = kindOf(TowerGroupKind.Top(depth))
-
     // Treating `a.foo()` common calls as more prioritized than `a.foo.invoke()`
     // It's not the same as TowerGroupKind because it's not about tower levels, but rather a different dimension semantically.
     // It could be implemented via another TowerGroupKind, but it's not clear what priority should be assigned to the new TowerGroupKind

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

@@ -190,8 +190,8 @@ class FirCallCompleter(
             lambdaArgument.replaceValueParameters(lambdaArgument.valueParameters + listOfNotNull(itParam))
             lambdaArgument.replaceReturnTypeRef(expectedReturnTypeRef ?: noExpectedType)
 
-            val localContext = localContextForAnonymousFunctions.getValue(lambdaArgument.symbol)
-            transformer.context.withLocalContext(localContext) {
+            val localContext = towerDataContextForAnonymousFunctions.getValue(lambdaArgument.symbol)
+            transformer.context.withTowerDataContext(localContext) {
                 lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
             }
             transformer.context.dropContextForAnonymousFunction(lambdaArgument)

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

@@ -10,6 +10,7 @@ import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionStageRunner
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
@@ -24,6 +25,8 @@ import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.sure
 
 abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAbstractPhaseTransformer<ResolutionMode>(phase) {
     abstract val context: BodyResolveContext
@@ -42,7 +45,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         }
 
     protected inline fun <T> withLocalScopeCleanup(crossinline l: () -> T): T {
-        return context.withLocalScopesCleanup(l)
+        return context.withTowerDataCleanup(l)
     }
 
     protected fun addLocalScope(localScope: FirLocalScope?) {
@@ -97,17 +100,16 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         val fileImportsScope: MutableList<FirScope> = mutableListOf()
 
         @set:PrivateForInline
-        var typeParametersScopes: FirTypeParametersScopes = persistentListOf()
+        lateinit var file: FirFile
+            internal set
 
-        @set:PrivateForInline
-        var localScopes: FirLocalScopes = persistentListOf()
+        val implicitReceiverStack: ImplicitReceiverStack get() = towerDataContext.implicitReceiverStack
 
         @set:PrivateForInline
-        lateinit var file: FirFile
-            internal set
+        var towerDataContext: FirTowerDataContext = FirTowerDataContext()
 
         @set:PrivateForInline
-        var implicitReceiverStack: MutableImplicitReceiverStack = ImplicitReceiverStackImpl()
+        var towerDataContextsForClassParts: FirTowerDataContextsForClassParts? = null
 
         val containerIfAny: FirDeclaration?
             get() = containers.lastOrNull()
@@ -115,48 +117,81 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         @set:PrivateForInline
         var containers: PersistentList<FirDeclaration> = persistentListOf()
 
-        val localContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirLocalContext> = mutableMapOf()
+        val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
 
         @OptIn(PrivateForInline::class)
-        inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
-            val oldContainers = containers
-            containers = containers.add(declaration)
+        inline fun <T> withNewTowerDataForClassParts(newContexts: FirTowerDataContextsForClassParts, f: () -> T): T {
+            val old = towerDataContextsForClassParts
+
+            towerDataContextsForClassParts = newContexts
+
             return try {
                 f()
             } finally {
-                containers = oldContainers
+
+                towerDataContextsForClassParts = old
             }
         }
 
-        @OptIn(PrivateForInline::class)
-        inline fun <T> withLocalContext(localContext: FirLocalContext, f: () -> T): T {
-            val existedStack = this.implicitReceiverStack
-            val existedLocalScopes = this.localScopes
+        fun getTowerDataContextForStaticNestedClassesUnsafe(): FirTowerDataContext =
+            firTowerDataContextsForClassParts().forNestedClasses
 
-            implicitReceiverStack = localContext.implicitReceiverStack
-            localScopes = localContext.localScopes
+        fun getTowerDataContextForConstructorResolution(): FirTowerDataContext =
+            firTowerDataContextsForClassParts().forConstructorHeaders
 
+        fun getPrimaryConstructorParametersScope(): FirLocalScope? =
+            towerDataContextsForClassParts?.primaryConstructorParametersScope
+
+        private fun firTowerDataContextsForClassParts() =
+            towerDataContextsForClassParts.sure { "towerDataContextForStaticNestedClasses should not be null" }
+
+        @OptIn(PrivateForInline::class)
+        inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
+            val oldContainers = containers
+            containers = containers.add(declaration)
             return try {
                 f()
             } finally {
-                implicitReceiverStack = existedStack
-                localScopes = existedLocalScopes
+                containers = oldContainers
+            }
+        }
+
+        inline fun <T> withTowerDataContext(context: FirTowerDataContext, f: () -> T): T {
+            return withTowerDataCleanup {
+                replaceTowerDataContext(context)
+                f()
             }
         }
 
         @OptIn(PrivateForInline::class)
-        inline fun <R> withLocalScopesCleanup(l: () -> R): R {
-            val initialLocalScopes = localScopes
+        inline fun <R> withTowerDataCleanup(l: () -> R): R {
+            val initialContext = towerDataContext
             return try {
                 l()
             } finally {
-                localScopes = initialLocalScopes
+                towerDataContext = initialContext
             }
         }
 
         @OptIn(PrivateForInline::class)
+        fun replaceTowerDataContext(newContext: FirTowerDataContext) {
+            towerDataContext = newContext
+        }
+
+        fun addTowerDataElement(element: FirTowerDataElement) {
+            replaceTowerDataContext(towerDataContext.addTowerDataElements(listOf(element)))
+        }
+
+        fun addTowerDataElements(newElements: List<FirTowerDataElement>) {
+            replaceTowerDataContext(towerDataContext.addTowerDataElements(newElements))
+        }
+
         fun addLocalScope(localScope: FirLocalScope) {
-            localScopes = localScopes.add(localScope)
+            replaceTowerDataContext(towerDataContext.addLocalScope(localScope))
+        }
+
+        fun addReceiver(name: Name?, implicitReceiverValue: ImplicitReceiverValue<*>) {
+            replaceTowerDataContext(towerDataContext.addReceiver(name, implicitReceiverValue))
         }
 
         fun storeClassIfNotNested(klass: FirRegularClass) {
@@ -173,21 +208,20 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         }
 
         fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-            localContextForAnonymousFunctions[anonymousFunction.symbol] = FirLocalContext(localScopes, implicitReceiverStack.snapshot())
+            towerDataContextForAnonymousFunctions[anonymousFunction.symbol] = towerDataContext
         }
 
         fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-            localContextForAnonymousFunctions.remove(anonymousFunction.symbol)
+            towerDataContextForAnonymousFunctions.remove(anonymousFunction.symbol)
         }
 
         fun cleanContextForAnonymousFunction() {
-            localContextForAnonymousFunctions.clear()
+            towerDataContextForAnonymousFunctions.clear()
         }
 
-        @OptIn(PrivateForInline::class)
         private inline fun updateLastScope(transform: FirLocalScope.() -> FirLocalScope) {
-            val lastScope = localScopes.lastOrNull() ?: return
-            localScopes = localScopes.set(localScopes.size - 1, lastScope.transform())
+            val lastScope = towerDataContext.localScopes.lastOrNull() ?: return
+            replaceTowerDataContext(towerDataContext.setLastLocalScope(lastScope.transform()))
         }
 
         @OptIn(PrivateForInline::class)
@@ -195,13 +229,10 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             returnTypeCalculator: ReturnTypeCalculator,
             targetedLocalClasses: Set<FirClass<*>>
         ) = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses).apply {
-            fileImportsScope.addAll(this@BodyResolveContext.fileImportsScope)
-            typeParametersScopes = this@BodyResolveContext.typeParametersScopes
-            localScopes = this@BodyResolveContext.localScopes
             file = this@BodyResolveContext.file
-            implicitReceiverStack = this@BodyResolveContext.implicitReceiverStack
-            localContextForAnonymousFunctions.putAll(this@BodyResolveContext.localContextForAnonymousFunctions)
+            towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
             containers = this@BodyResolveContext.containers
+            towerDataContext = this@BodyResolveContext.towerDataContext
         }
     }
 
@@ -212,12 +243,12 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         val context: BodyResolveContext
     ) : BodyResolveComponents {
         override val fileImportsScope: List<FirScope> get() = context.fileImportsScope
-        override val typeParametersScopes: List<FirScope> get() = context.typeParametersScopes
-        override val localScopes: FirLocalScopes get() = context.localScopes
+        override val towerDataElements: List<FirTowerDataElement> get() = context.towerDataContext.towerDataElements
+        override val localScopes: FirLocalScopes get() = context.towerDataContext.localScopes
         override val file: FirFile get() = context.file
         override val implicitReceiverStack: ImplicitReceiverStack get() = context.implicitReceiverStack
         override val containingDeclarations: List<FirDeclaration> get() = context.containers
-        override val localContextForAnonymousFunctions: LocalContextForAnonymousFunctions get() = context.localContextForAnonymousFunctions
+        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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.asTowerDataElement
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
@@ -50,6 +51,8 @@ open class FirBodyResolveTransformer(
         return withScopeCleanup(context.fileImportsScope) {
             val importingScopes = createImportingScopes(file, session, components.scopeSession)
             context.fileImportsScope += importingScopes
+            context.addTowerDataElements(importingScopes.map { it.asTowerDataElement(isLocal = false) })
+
             file.replaceResolvePhase(transformerPhase)
             @Suppress("UNCHECKED_CAST")
             transformDeclarationContent(file, data) as CompositeTransformResult<FirFile>

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

@@ -5,13 +5,15 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirFunctionTarget
+import org.jetbrains.kotlin.fir.copy
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
+import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildUnitExpression
@@ -19,10 +21,12 @@ import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.impl.FirEmptyControlFlowGraphReference
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
 import org.jetbrains.kotlin.fir.resolve.inference.FirDelegatedPropertyInferenceSession
 import org.jetbrains.kotlin.fir.resolve.inference.extractLambdaInfoFromFunctionalType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.*
+import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirMemberTypeParameterScope
 import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType
@@ -35,8 +39,6 @@ import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
 
 open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
-    private var primaryConstructorParametersScope: FirLocalScope? = null
-
     private var containingClass: FirRegularClass? = null
 
     private fun transformDeclarationContent(
@@ -66,22 +68,23 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         }
     }
 
-    protected inline fun <T> withTypeParametersOf(declaration: FirMemberDeclaration, crossinline l: () -> T): T {
-        if (declaration.typeParameters.isEmpty()) return l()
+    protected fun createTypeParameterScope(declaration: FirMemberDeclaration): FirMemberTypeParameterScope? {
+        if (declaration.typeParameters.isEmpty()) return null
 
         for (typeParameter in declaration.typeParameters) {
             (typeParameter as? FirTypeParameter)?.replaceResolvePhase(FirResolvePhase.STATUS)
             typeParameter.transformChildren(transformer, ResolutionMode.ContextIndependent)
         }
 
-        val before = context.typeParametersScopes
-        @OptIn(PrivateForInline::class)
-        context.typeParametersScopes = context.typeParametersScopes.add(FirMemberTypeParameterScope(declaration))
-        return try {
+        return FirMemberTypeParameterScope(declaration)
+    }
+
+    protected inline fun <T> withTypeParametersOf(declaration: FirMemberDeclaration, crossinline l: () -> T): T {
+        val scope = createTypeParameterScope(declaration) ?: return l()
+
+        return context.withTowerDataCleanup {
+            context.addTowerDataElement(scope.asTowerDataElement(isLocal = false))
             l()
-        } finally {
-            @OptIn(PrivateForInline::class)
-            context.typeParametersScopes = before
         }
     }
 
@@ -111,11 +114,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 withLocalScopeCleanup {
                     context.withContainer(property) {
                         if (property.delegate != null) {
-                            addLocalScope(primaryConstructorParametersScope)
+                            addLocalScope(context.getPrimaryConstructorParametersScope())
                             transformPropertyWithDelegate(property)
                         } else {
                             withLocalScopeCleanup {
-                                addLocalScope(primaryConstructorParametersScope)
+                                addLocalScope(context.getPrimaryConstructorParametersScope())
                                 property.transformChildrenWithoutAccessors(returnTypeRef)
                                 property.transformInitializer(integerLiteralTypeApproximator, null)
                             }
@@ -323,8 +326,13 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             return regularClass.runAllPhasesForLocalClass(transformer, components, data).compose()
         }
 
+        return context.withTowerDataCleanup {
+            if (!regularClass.isInner && context.containerIfAny is FirRegularClass) {
+                context.replaceTowerDataContext(
+                    context.getTowerDataContextForStaticNestedClassesUnsafe()
+                )
+            }
 
-        return withTypeParametersOf(regularClass) {
             doTransformRegularClass(regularClass, data)
         }
     }
@@ -339,20 +347,10 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             dataFlowAnalyzer.enterClass()
         }
 
-        val oldConstructorScope = primaryConstructorParametersScope
         val oldContainingClass = containingClass
-        primaryConstructorParametersScope = null
         containingClass = regularClass
         val type = regularClass.defaultType()
-        val result = withLabelAndReceiverType(regularClass.name, regularClass, type) {
-            val constructor = regularClass.declarations.firstOrNull() as? FirConstructor
-            if (constructor?.isPrimary == true) {
-                primaryConstructorParametersScope = FirLocalScope().let {
-                    var scope = it
-                    constructor.valueParameters.forEach { scope = scope.storeVariable(it) }
-                    scope
-                }
-            }
+        val result = withScopesForClass(regularClass.name, regularClass, type) {
             transformDeclarationContent(regularClass, data).single as FirRegularClass
         }
 
@@ -366,7 +364,6 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         }
 
         containingClass = oldContainingClass
-        primaryConstructorParametersScope = oldConstructorScope
         return (@Suppress("UNCHECKED_CAST")
         result.compose())
     }
@@ -386,7 +383,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 this.type = type
             }
         }
-        var result = withLabelAndReceiverType(null, anonymousObject, type) {
+        var result = withScopesForClass(null, anonymousObject, type) {
             transformDeclarationContent(anonymousObject, data).single as FirAnonymousObject
         }
         if (!implicitTypeOnly && result.controlFlowGraphReference == FirEmptyControlFlowGraphReference) {
@@ -521,9 +518,10 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         if (implicitTypeOnly) return anonymousInitializer.compose()
         return withLocalScopeCleanup {
             dataFlowAnalyzer.enterInitBlock(anonymousInitializer)
-            addLocalScope(primaryConstructorParametersScope)
+            addLocalScope(context.getPrimaryConstructorParametersScope())
             addLocalScope(FirLocalScope())
-            val result = transformDeclarationContent(anonymousInitializer, ResolutionMode.ContextIndependent).single as FirAnonymousInitializer
+            val result =
+                transformDeclarationContent(anonymousInitializer, ResolutionMode.ContextIndependent).single as FirAnonymousInitializer
             val graph = dataFlowAnalyzer.exitInitBlock(result)
             result.transformControlFlowGraphReference(ControlFlowGraphReferenceTransformer, graph).compose()
         }
@@ -683,27 +681,79 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         return this
     }
 
+    private inline fun <T> withScopesForClass(
+        labelName: Name?,
+        owner: FirClass<*>,
+        type: ConeKotlinType,
+        block: () -> T
+    ): T = context.withTowerDataCleanup {
+        val towerElementsForClass = components.collectTowerDataElementsForClass(owner, type)
+
+        val staticsAndCompanion =
+            context.towerDataContext
+                .addTowerDataElements(towerElementsForClass.superClassesStaticsAndCompanionReceivers)
+                .run {
+                    if (towerElementsForClass.companionReceiver != null)
+                        addReceiver(null, towerElementsForClass.companionReceiver)
+                    else
+                        this
+                }
+                .addNonLocalScopeIfNotNull(towerElementsForClass.companionStaticScope)
+                .addNonLocalScopeIfNotNull(towerElementsForClass.staticScope)
+
+        val typeParameterScope = (owner as? FirRegularClass)?.let(this::createTypeParameterScope)
+
+        val forMembersResolution =
+            staticsAndCompanion
+                .addReceiver(labelName, towerElementsForClass.thisReceiver)
+                .addNonLocalScopeIfNotNull(typeParameterScope)
+
+        val scopeForConstructorHeader =
+            staticsAndCompanion.addNonLocalScopeIfNotNull(typeParameterScope)
+
+        val newTowerDataContextForStaticNestedClasses =
+            if ((owner as? FirRegularClass)?.classKind?.isSingleton == true)
+                forMembersResolution
+            else
+                staticsAndCompanion
+
+        val constructor = (owner as? FirRegularClass)?.declarations?.firstOrNull() as? FirConstructor
+        val primaryConstructorParametersScope =
+            if (constructor?.isPrimary == true) {
+                constructor.valueParameters.fold(FirLocalScope()) { acc, param -> acc.storeVariable(param) }
+            } else null
+
+        components.context.replaceTowerDataContext(forMembersResolution)
+
+        val newContexts =
+            FirTowerDataContextsForClassParts(
+                newTowerDataContextForStaticNestedClasses,
+                scopeForConstructorHeader,
+                primaryConstructorParametersScope
+            )
+
+        context.withNewTowerDataForClassParts(newContexts) {
+            block()
+        }
+    }
+
     protected inline fun <T> withLabelAndReceiverType(
         labelName: Name?,
-        owner: FirDeclaration,
+        owner: FirCallableDeclaration<*>,
         type: ConeKotlinType?,
         block: () -> T
-    ): T {
-        val (implicitReceiverValue, implicitCompanionValues) = components.collectImplicitReceivers(type, owner)
-        implicitCompanionValues.forEach { value ->
-            context.implicitReceiverStack.add(null, value)
+    ): T = context.withTowerDataCleanup {
+        if (type != null) {
+            val receiver = ImplicitExtensionReceiverValue(
+                owner.symbol,
+                type,
+                components.session,
+                components.scopeSession
+            )
+            context.addReceiver(labelName, receiver)
         }
-        implicitReceiverValue?.let { context.implicitReceiverStack.add(labelName, it) }
 
-        try {
-            return block()
-        } finally {
-            if (type != null) {
-                context.implicitReceiverStack.pop(labelName)
-                for (i in implicitCompanionValues.indices)
-                    context.implicitReceiverStack.pop(null)
-            }
-        }
+        block()
     }
 
     private fun storeVariableReturnType(variable: FirVariable<*>) {

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

@@ -9,9 +9,9 @@ import com.intellij.openapi.progress.ProcessCanceledException
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
+import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildErrorExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
@@ -309,7 +309,8 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
                 }
                 else -> buildErrorExpression {
                     source = operatorCall.source
-                    diagnostic = ConeOperatorAmbiguityError(listOf(operatorCallReference.resolvedSymbol, assignCallReference.resolvedSymbol))
+                    diagnostic =
+                        ConeOperatorAmbiguityError(listOf(operatorCallReference.resolvedSymbol, assignCallReference.resolvedSymbol))
                 }.compose()
             }
         }
@@ -541,7 +542,8 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
             FirConstKind.Float -> constructLiteralType(StandardClassIds.Float)
             FirConstKind.Double -> constructLiteralType(StandardClassIds.Double)
             FirConstKind.IntegerLiteral, FirConstKind.UnsignedIntegerLiteral -> {
-                val integerLiteralType = ConeIntegerLiteralTypeImpl(constExpression.value as Long, isUnsigned = kind == FirConstKind.UnsignedIntegerLiteral)
+                val integerLiteralType =
+                    ConeIntegerLiteralTypeImpl(constExpression.value as Long, isUnsigned = kind == FirConstKind.UnsignedIntegerLiteral)
                 val expectedType = data.expectedType?.coneTypeSafe<ConeKotlinType>()
                 if (expectedType != null) {
                     val approximatedType = integerLiteralType.getApproximatedType(expectedType)
@@ -633,15 +635,13 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         var result = delegatedConstructorCall
         try {
             val lastDispatchReceiver = implicitReceiverStack.lastDispatchReceiver()
-            val name = lastDispatchReceiver?.boundSymbol?.classId?.shortClassName
-            if (lastDispatchReceiver != null) {
-                context.implicitReceiverStack.pop(name)
-                context.implicitReceiverStack.add(name, lastDispatchReceiver.copyForDelegated())
-            }
-            delegatedConstructorCall.transformChildren(transformer, ResolutionMode.ContextDependent)
-            if (lastDispatchReceiver != null) {
-                context.implicitReceiverStack.pop(name)
-                context.implicitReceiverStack.add(name, lastDispatchReceiver)
+            context.withTowerDataCleanup {
+                if ((context.containerIfAny as? FirConstructor)?.isPrimary == true) {
+                    context.replaceTowerDataContext(context.getTowerDataContextForConstructorResolution())
+                    context.getPrimaryConstructorParametersScope()?.let(context::addLocalScope)
+                }
+
+                delegatedConstructorCall.transformChildren(transformer, ResolutionMode.ContextDependent)
             }
             val typeArguments: List<FirTypeProjection>
             val reference = delegatedConstructorCall.calleeReference
@@ -700,7 +700,10 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    override fun transformAugmentedArraySetCall(augmentedArraySetCall: FirAugmentedArraySetCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformAugmentedArraySetCall(
+        augmentedArraySetCall: FirAugmentedArraySetCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         assert(augmentedArraySetCall.operation in FirOperation.ASSIGNMENTS)
         assert(augmentedArraySetCall.operation != FirOperation.ASSIGN)
 

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.scopes.impl
 
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 
@@ -18,12 +17,4 @@ class FirTypeResolveScopeForBodyResolve(
 }
 
 fun BodyResolveComponents.createCurrentScopeList(): List<FirScope> =
-    mutableListOf<FirScope>().apply {
-        addAll(localScopes.asReversed())
-        implicitReceiverStack.receiversAsReversed().mapNotNullTo(this) {
-            (it as? ImplicitDispatchReceiverValue)?.implicitScope
-        }
-
-        addAll(typeParametersScopes.asReversed())
-        addAll(fileImportsScope.asReversed())
-    }
+    towerDataElements.asReversed().mapNotNull { it.scope }