FIR: introduce BodyResolveContext.forFunctionBody

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

@@ -379,7 +379,6 @@ class BodyResolveContext(
 
     inline fun <T> withSimpleFunction(
         simpleFunction: FirSimpleFunction,
-        holder: SessionHolder,
         crossinline f: () -> T
     ): T {
         if (containerIfAny !is FirClass<*>) {
@@ -387,12 +386,45 @@ class BodyResolveContext(
         }
 
         return withTypeParametersOf(simpleFunction) {
+            withContainer(simpleFunction, f)
+        }
+    }
+
+    inline fun <T> forFunctionBody(
+        function: FirFunction<*>,
+        holder: SessionHolder,
+        crossinline f: () -> T
+    ): T {
+        return withTowerDataCleanup {
+            addLocalScope(FirLocalScope())
+            if (function is FirSimpleFunction) {
+                val receiverTypeRef = function.receiverTypeRef
+                withLabelAndReceiverType(function.name, function, receiverTypeRef?.coneType, holder, f)
+            } else {
+                f()
+            }
+        }
+    }
+
+    inline fun <T> forConstructorBody(
+        constructor: FirConstructor,
+        parametersScope: FirLocalScope?,
+        crossinline f: () -> T
+    ): T {
+        return if (constructor.isPrimary) {
+            /*
+             * Primary constructor may have body only if class delegates implementation to some property
+             *   In it's body we don't have this receiver for building class, so we need to use
+             *   special towerDataContext
+             */
+            withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
+                parametersScope?.let { addLocalScope(it) }
+                f()
+            }
+        } else {
             withTowerDataCleanup {
-                addLocalScope(FirLocalScope())
-                val receiverTypeRef = simpleFunction.receiverTypeRef
-                withContainer(simpleFunction) {
-                    withLabelAndReceiverType(simpleFunction.name, simpleFunction, receiverTypeRef?.coneType, holder, f)
-                }
+                parametersScope?.let { addLocalScope(it) }
+                f()
             }
         }
     }
@@ -584,8 +616,7 @@ class BodyResolveContext(
     }
 
     inline fun <T> forDelegatedConstructor(
-        constructor: FirConstructor,
-        parametersScope: FirLocalScope? = buildConstructorParametersScope(constructor),
+        parametersScope: FirLocalScope?,
         crossinline f: () -> T
     ): T {
         if (parametersScope == null) return f()
@@ -595,29 +626,6 @@ class BodyResolveContext(
         }
     }
 
-    inline fun <T> forConstructorBody(
-        constructor: FirConstructor,
-        parametersScope: FirLocalScope? = buildConstructorParametersScope(constructor),
-        crossinline f: () -> T
-    ): T {
-        return if (constructor.isPrimary) {
-            /*
-             * Primary constructor may have body only if class delegates implementation to some property
-             *   In it's body we don't have this receiver for building class, so we need to use
-             *   special towerDataContext
-             */
-            withTowerDataMode(FirTowerDataMode.CONSTRUCTOR_HEADER) {
-                parametersScope?.let { addLocalScope(it) }
-                f()
-            }
-        } else {
-            withTowerDataCleanup {
-                parametersScope?.let { addLocalScope(it) }
-                f()
-            }
-        }
-    }
-
     fun buildConstructorParametersScope(constructor: FirConstructor): FirLocalScope? =
         if (constructor.isPrimary) getPrimaryConstructorAllParametersScope()
         else constructor.valueParameters.fold(FirLocalScope()) { acc, param -> acc.storeVariable(param) }

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

@@ -456,19 +456,19 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         }
 
         doTransformTypeParameters(simpleFunction)
-        // TODO: I think it worth creating something like runAllPhasesForLocalFunction
-        context.withTypeParametersOf(simpleFunction) {
-            val containingDeclaration = context.containerIfAny
+
+        val containingDeclaration = context.containerIfAny
+        return context.withSimpleFunction(simpleFunction) {
+            // TODO: I think it worth creating something like runAllPhasesForLocalFunction
             if (containingDeclaration != null && containingDeclaration !is FirClass<*>) {
                 // For class members everything should be already prepared
                 prepareSignatureForBodyResolve(simpleFunction)
                 simpleFunction.transformStatus(this, simpleFunction.resolveStatus().mode())
             }
-        }
-
-        return withFullBodyResolve {
-            context.withSimpleFunction(simpleFunction, components) {
-                transformFunctionWithGivenSignature(simpleFunction, ResolutionMode.ContextIndependent)
+            context.forFunctionBody(simpleFunction, components) {
+                withFullBodyResolve {
+                    transformFunctionWithGivenSignature(simpleFunction, ResolutionMode.ContextIndependent)
+                }
             }
         }
     }
@@ -553,7 +553,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 constructor.transformValueParameters(transformer, data)
             }
             val parametersScope = context.buildConstructorParametersScope(constructor)
-            context.forDelegatedConstructor(constructor, parametersScope) {
+            context.forDelegatedConstructor(parametersScope) {
                 constructor.transformDelegatedConstructor(transformer, data)
             }
             context.forConstructorBody(constructor, parametersScope) {
@@ -587,8 +587,9 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             valueParameter.replaceReturnTypeRef(
                 valueParameter.returnTypeRef.errorTypeFromPrototype(ConeSimpleDiagnostic("Unresolved value parameter type"))
             )
-            context.storeVariable(valueParameter)
-            return valueParameter.compose()
+            return context.withValueParameter(valueParameter) {
+                valueParameter.compose()
+            }
         }
 
         dataFlowAnalyzer.enterValueParameter(valueParameter)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -66,8 +66,10 @@ open class FirContractResolveTransformer(
                 return simpleFunction.compose()
             }
             @Suppress("UNCHECKED_CAST")
-            return context.withSimpleFunction(simpleFunction, components) {
-                transformContractDescriptionOwner(simpleFunction)
+            return context.withSimpleFunction(simpleFunction) {
+                context.forFunctionBody(simpleFunction, components) {
+                    transformContractDescriptionOwner(simpleFunction)
+                }
             }
         }