FIR: introduce separate containingClass in BodyResolveContext

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

@@ -75,6 +75,9 @@ class BodyResolveContext(
     @set:PrivateForInline
     var containers: PersistentList<FirDeclaration> = persistentListOf()
 
+    @set:PrivateForInline
+    var containingClass: FirRegularClass? = null
+
     val containerIfAny: FirDeclaration?
         get() = containers.lastOrNull()
 
@@ -114,6 +117,20 @@ class BodyResolveContext(
         }
     }
 
+    @PrivateForInline
+    private inline fun <T> withContainerClass(declaration: FirRegularClass, f: () -> T): T {
+        val oldContainers = containers
+        val oldContainingClass = containingClass
+        containers = containers.add(declaration)
+        containingClass = declaration
+        return try {
+            f()
+        } finally {
+            containers = oldContainers
+            containingClass = oldContainingClass
+        }
+    }
+
     @PrivateForInline
     inline fun <R> withTowerDataCleanup(l: () -> R): R {
         val initialContext = towerDataContext
@@ -340,7 +357,7 @@ class BodyResolveContext(
         storeClassIfNotNested(regularClass)
         if (forContracts) {
             return withTypeParametersOf(regularClass) {
-                withContainer(regularClass, f)
+                withContainerClass(regularClass, f)
             }
         }
         return withTowerModeCleanup {
@@ -353,7 +370,7 @@ class BodyResolveContext(
             }
 
             withScopesForClass(regularClass, holder) {
-                withContainer(regularClass, f)
+                withContainerClass(regularClass, f)
             }
         }
     }

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

@@ -39,9 +39,6 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
-    private val containingClass: FirRegularClass?
-        get() = context.containers.asReversed().firstOrNull { it is FirRegularClass } as? FirRegularClass
-
     private val statusResolver: FirStatusResolver = FirStatusResolver(session, scopeSession)
 
     private fun FirDeclaration.visibilityForApproximation(): Visibility {
@@ -187,7 +184,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         if (typeRef is FirResolvedTypeRef && property.returnTypeRef is FirResolvedTypeRef) {
             val typeArguments = (typeRef.type as ConeClassLikeType).typeArguments
             val extensionType = property.receiverTypeRef?.coneType
-            val dispatchType = containingClass?.let { containingClass ->
+            val dispatchType = context.containingClass?.let { containingClass ->
                 containingClass.symbol.constructStarProjectedType(containingClass.typeParameters.size)
             }
             propertyReferenceAccess.replaceTypeRef(