From ca15df781536bd3d9e072947233162d1ace66aba Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 16 Mar 2021 19:15:37 +0300
Subject: [PATCH] FIR: introduce separate containingClass in BodyResolveContext

---
 .../body/resolve/BodyResolveContext.kt        | 21 +++++++++++++++++--
 .../FirDeclarationsResolveTransformer.kt      |  5 +----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
index 75412810e2c..62d4b4f7e31 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
+++ b/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)
             }
         }
     }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
index b6a8fbd1d75..279ca575788 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
+++ b/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(