From 4235075f355ac9f538dfcdc82f5fa4f70ca887f7 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 9 Mar 2021 17:09:10 +0300
Subject: [PATCH] FIR: extract BodyResolveContext.withAnonymousFunction

---
 .../body/resolve/BodyResolveContext.kt        | 31 +++++++++++---
 .../FirDeclarationsResolveTransformer.kt      | 40 +++++++------------
 2 files changed, 39 insertions(+), 32 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 8335a94d7b8..67743f4b9ed 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
@@ -380,11 +380,30 @@ class BodyResolveContext(
         }
 
         return withTypeParametersOf(simpleFunction) {
-            val receiverTypeRef = simpleFunction.receiverTypeRef
-            if (receiverTypeRef != null) {
-                withLabelAndReceiverType(simpleFunction.name, simpleFunction, receiverTypeRef.coneType, holder, f)
-            } else {
-                f()
+            withTowerDataCleanup {
+                addLocalScope(FirLocalScope())
+                val receiverTypeRef = simpleFunction.receiverTypeRef
+                withLabelAndReceiverType(simpleFunction.name, simpleFunction, receiverTypeRef?.coneType, holder, f)
+            }
+        }
+    }
+
+    inline fun <T> withAnonymousFunction(
+        anonymousFunction: FirAnonymousFunction,
+        holder: SessionHolder,
+        isInDependentContext: Boolean = false,
+        crossinline f: () -> T
+    ): T {
+        return withTowerDataCleanup {
+            addLocalScope(FirLocalScope())
+            val receiverTypeRef = anonymousFunction.receiverTypeRef
+            val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }
+            withLabelAndReceiverType(labelName, anonymousFunction, receiverTypeRef?.coneType, holder) {
+                if (isInDependentContext) {
+                    withLambdaBeingAnalyzedInDependentContext(anonymousFunction.symbol, f)
+                } else {
+                    f()
+                }
             }
         }
     }
@@ -419,10 +438,10 @@ class BodyResolveContext(
     ): T {
         return withTowerDataCleanup {
             val receiverTypeRef = property.receiverTypeRef
+            addLocalScope(FirLocalScope())
             if (receiverTypeRef == null && property.returnTypeRef !is FirImplicitTypeRef &&
                 !property.isLocal && property.delegate == null
             ) {
-                addLocalScope(FirLocalScope())
                 storeBackingField(property)
             }
             if (receiverTypeRef != null) {
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 bd69f6778a5..ec38fa43830 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
@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.fir.resolve.inference.extractLambdaInfoFromFunctiona
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.*
-import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirMemberTypeParameterScope
 import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
@@ -427,14 +426,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     private fun transformAnonymousFunctionWithLambdaResolution(
         anonymousFunction: FirAnonymousFunction, lambdaResolution: ResolutionMode.LambdaResolution
     ): FirAnonymousFunction {
-        val receiverTypeRef = anonymousFunction.receiverTypeRef
         fun transform(): FirAnonymousFunction {
             val expectedReturnType =
                 lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
 
-            val result = context.withLambdaBeingAnalyzedInDependentContext(anonymousFunction.symbol) {
-                transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
-            }
+            val result = transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
 
             val body = result.body
             if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
@@ -459,12 +455,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             return result
         }
 
-        val label = anonymousFunction.label
-        return if (label != null || receiverTypeRef is FirResolvedTypeRef) {
-            withLabelAndReceiverType(label?.name?.let { Name.identifier(it) }, anonymousFunction, receiverTypeRef?.coneTypeSafe()) {
-                transform()
-            }
-        } else {
+        return context.withAnonymousFunction(anonymousFunction, components, isInDependentContext = true) {
             transform()
         }
     }
@@ -536,20 +527,18 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         function: FirFunction<F>,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
-        return withNewLocalScope {
-            val functionIsNotAnalyzed = transformerPhase != function.resolvePhase
-            if (functionIsNotAnalyzed) {
-                dataFlowAnalyzer.enterFunction(function)
-            }
-            @Suppress("UNCHECKED_CAST")
-            transformDeclarationContent(function, data).also {
-                if (functionIsNotAnalyzed) {
-                    val result = it.single as FirFunction<*>
-                    val controlFlowGraphReference = dataFlowAnalyzer.exitFunction(result)
-                    result.replaceControlFlowGraphReference(controlFlowGraphReference)
-                }
-            } as CompositeTransformResult<FirStatement>
+        val functionIsNotAnalyzed = transformerPhase != function.resolvePhase
+        if (functionIsNotAnalyzed) {
+            dataFlowAnalyzer.enterFunction(function)
         }
+        @Suppress("UNCHECKED_CAST")
+        return transformDeclarationContent(function, data).also {
+            if (functionIsNotAnalyzed) {
+                val result = it.single as FirFunction<*>
+                val controlFlowGraphReference = dataFlowAnalyzer.exitFunction(result)
+                result.replaceControlFlowGraphReference(controlFlowGraphReference)
+            }
+        } as CompositeTransformResult<FirStatement>
     }
 
     override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
@@ -711,8 +700,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 )
                 lambda = lambda.transformValueParameters(ImplicitToErrorTypeTransformer, null)
                 val bodyExpectedType = returnTypeRefFromResolvedAtom ?: expectedTypeRef
-                val labelName = lambda.label?.name?.let { Name.identifier(it) }
-                withLabelAndReceiverType(labelName, lambda, lambda.receiverTypeRef?.coneType) {
+                context.withAnonymousFunction(lambda, components) {
                     lambda = transformFunction(lambda, withExpectedType(bodyExpectedType)).single as FirAnonymousFunction
                 }
                 // To separate function and separate commit