FIR: extract BodyResolveContext.withAnonymousInitializer, withParameter

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.resolve.*
@@ -281,7 +282,9 @@ class BodyResolveContext(
                 }
             }
 
-            withScopesForClass(regularClass, holder, f)
+            withScopesForClass(regularClass, holder) {
+                withContainer(regularClass, f)
+            }
         }
     }
 
@@ -290,7 +293,9 @@ class BodyResolveContext(
         holder: SessionHolder,
         crossinline f: () -> T
     ): T {
-        return withScopesForClass(anonymousObject, holder, f)
+        return withScopesForClass(anonymousObject, holder) {
+            withContainer(anonymousObject, f)
+        }
     }
 
     inline fun <T> withScopesForClass(
@@ -383,7 +388,9 @@ class BodyResolveContext(
             withTowerDataCleanup {
                 addLocalScope(FirLocalScope())
                 val receiverTypeRef = simpleFunction.receiverTypeRef
-                withLabelAndReceiverType(simpleFunction.name, simpleFunction, receiverTypeRef?.coneType, holder, f)
+                withContainer(simpleFunction) {
+                    withLabelAndReceiverType(simpleFunction.name, simpleFunction, receiverTypeRef?.coneType, holder, f)
+                }
             }
         }
     }
@@ -398,11 +405,13 @@ class BodyResolveContext(
             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()
+            withContainer(anonymousFunction) {
+                withLabelAndReceiverType(labelName, anonymousFunction, receiverTypeRef?.coneType, holder) {
+                    if (isInDependentContext) {
+                        withLambdaBeingAnalyzedInDependentContext(anonymousFunction.symbol, f)
+                    } else {
+                        f()
+                    }
                 }
             }
         }
@@ -422,6 +431,25 @@ class BodyResolveContext(
         }
     }
 
+    inline fun <T> withAnonymousInitializer(
+        anonymousInitializer: FirAnonymousInitializer,
+        crossinline f: () -> T
+    ): T {
+        return withTowerDataCleanup {
+            getPrimaryConstructorPureParametersScope()?.let { addLocalScope(it) }
+            addLocalScope(FirLocalScope())
+            withContainer(anonymousInitializer, f)
+        }
+    }
+
+    inline fun <T> withValueParameter(
+        valueParameter: FirValueParameter,
+        crossinline f: () -> T
+    ): T {
+        storeVariable(valueParameter)
+        return withContainer(valueParameter, f)
+    }
+
     inline fun <T> withProperty(
         property: FirProperty,
         crossinline f: () -> T
@@ -433,9 +461,13 @@ class BodyResolveContext(
 
     inline fun <T> withPropertyAccessor(
         property: FirProperty,
+        accessor: FirPropertyAccessor,
         holder: SessionHolder,
         crossinline f: () -> T
     ): T {
+        if (accessor is FirDefaultPropertyAccessor || accessor.body == null) {
+            return withContainer(accessor, f)
+        }
         return withTowerDataCleanup {
             val receiverTypeRef = property.receiverTypeRef
             addLocalScope(FirLocalScope())
@@ -444,10 +476,8 @@ class BodyResolveContext(
             ) {
                 storeBackingField(property)
             }
-            if (receiverTypeRef != null) {
-                withLabelAndReceiverType(property.name, property, receiverTypeRef.coneType, holder, f)
-            } else {
-                f()
+            withContainer(accessor) {
+                withLabelAndReceiverType(property.name, property, receiverTypeRef?.coneType, holder, f)
             }
         }
     }

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

@@ -69,10 +69,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         declaration: FirDeclaration, data: ResolutionMode
     ): CompositeTransformResult<FirDeclaration> {
         transformer.onBeforeDeclarationContentResolve(declaration)
-        return context.withContainer(declaration) {
-            transformer.replaceDeclarationResolvePhaseIfNeeded(declaration, transformerPhase)
-            transformer.transformDeclarationContent(declaration, data)
-        }
+        transformer.replaceDeclarationResolvePhaseIfNeeded(declaration, transformerPhase)
+        return transformer.transformDeclarationContent(declaration, data)
     }
 
     override fun transformDeclarationStatus(
@@ -327,24 +325,24 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         enhancedTypeRef: FirTypeRef,
         owner: FirProperty
     ) {
-        if (accessor is FirDefaultPropertyAccessor || accessor.body == null) {
-            transformFunction(accessor, withExpectedType(enhancedTypeRef))
-            return
-        }
-        val returnTypeRef = accessor.returnTypeRef
-        val expectedReturnTypeRef = if (enhancedTypeRef is FirResolvedTypeRef && returnTypeRef !is FirResolvedTypeRef) {
-            enhancedTypeRef
-        } else {
-            returnTypeRef
-        }
-        val resolutionMode = if (expectedReturnTypeRef.coneTypeSafe<ConeKotlinType>() == session.builtinTypes.unitType.type) {
-            ResolutionMode.ContextIndependent
-        } else {
-            withExpectedType(expectedReturnTypeRef)
-        }
+        context.withPropertyAccessor(owner, accessor, components) {
+            if (accessor is FirDefaultPropertyAccessor || accessor.body == null) {
+                transformFunction(accessor, withExpectedType(enhancedTypeRef))
+            } else {
+                val returnTypeRef = accessor.returnTypeRef
+                val expectedReturnTypeRef = if (enhancedTypeRef is FirResolvedTypeRef && returnTypeRef !is FirResolvedTypeRef) {
+                    enhancedTypeRef
+                } else {
+                    returnTypeRef
+                }
+                val resolutionMode = if (expectedReturnTypeRef.coneTypeSafe<ConeKotlinType>() == session.builtinTypes.unitType.type) {
+                    ResolutionMode.ContextIndependent
+                } else {
+                    withExpectedType(expectedReturnTypeRef)
+                }
 
-        context.withPropertyAccessor(owner, components) {
-            transformFunctionWithGivenSignature(accessor, resolutionMode)
+                transformFunctionWithGivenSignature(accessor, resolutionMode)
+            }
         }
     }
 
@@ -367,9 +365,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         }
 
         doTransformTypeParameters(regularClass)
-        return context.withRegularClass(regularClass, components) {
-            doTransformRegularClass(regularClass, data)
-        }
+        return doTransformRegularClass(regularClass, data)
     }
 
     private fun doTransformRegularClass(
@@ -382,7 +378,9 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             dataFlowAnalyzer.enterClass()
         }
 
-        val result = transformDeclarationContent(regularClass, data).single as FirRegularClass
+        val result = context.withRegularClass(regularClass, components) {
+            transformDeclarationContent(regularClass, data).single as FirRegularClass
+        }
 
         if (notAnalyzed) {
             if (!implicitTypeOnly) {
@@ -393,8 +391,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             }
         }
 
-        return (@Suppress("UNCHECKED_CAST")
-        result.compose())
+        (@Suppress("UNCHECKED_CAST")
+        return result.compose())
     }
 
     override fun transformAnonymousObject(
@@ -588,9 +586,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         data: ResolutionMode
     ): CompositeTransformResult<FirDeclaration> {
         if (implicitTypeOnly) return anonymousInitializer.compose()
-        return withPrimaryConstructorParameters(includeProperties = false) {
-            dataFlowAnalyzer.enterInitBlock(anonymousInitializer)
-            addNewLocalScope()
+        dataFlowAnalyzer.enterInitBlock(anonymousInitializer)
+        return context.withAnonymousInitializer(anonymousInitializer) {
             val result =
                 transformDeclarationContent(anonymousInitializer, ResolutionMode.ContextIndependent).single as FirAnonymousInitializer
             val graph = dataFlowAnalyzer.exitInitBlock(result)
@@ -600,20 +597,22 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     }
 
     override fun transformValueParameter(valueParameter: FirValueParameter, data: ResolutionMode): CompositeTransformResult<FirStatement> {
-        context.storeVariable(valueParameter)
         if (valueParameter.returnTypeRef is FirImplicitTypeRef) {
             transformer.replaceDeclarationResolvePhaseIfNeeded(valueParameter, transformerPhase)
             valueParameter.replaceReturnTypeRef(
                 valueParameter.returnTypeRef.errorTypeFromPrototype(ConeSimpleDiagnostic("Unresolved value parameter type"))
             )
+            context.storeVariable(valueParameter)
             return valueParameter.compose()
         }
 
         dataFlowAnalyzer.enterValueParameter(valueParameter)
-        val result = transformDeclarationContent(
-            valueParameter,
-            withExpectedType(valueParameter.returnTypeRef)
-        ).single as FirValueParameter
+        val result = context.withValueParameter(valueParameter) {
+            transformDeclarationContent(
+                valueParameter,
+                withExpectedType(valueParameter.returnTypeRef)
+            ).single as FirValueParameter
+        }
 
         dataFlowAnalyzer.exitValueParameter(result)?.let { graph ->
             result.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(graph))