FIR: enhance (using mode inside) BodyResolveContext.withAnonymousFunction

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

@@ -214,10 +214,6 @@ class BodyResolveContext(
         updateLastScope { storeBackingField(property) }
     }
 
-    fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-        towerDataContextForAnonymousFunctions[anonymousFunction.symbol] = towerDataContext
-    }
-
     fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
         towerDataContextForAnonymousFunctions.remove(anonymousFunction.symbol)
     }
@@ -404,16 +400,22 @@ class BodyResolveContext(
     inline fun <T> withAnonymousFunction(
         anonymousFunction: FirAnonymousFunction,
         holder: SessionHolder,
-        isInDependentContext: Boolean = false,
+        mode: ResolutionMode,
         crossinline f: () -> T
     ): T {
+        if (mode !is ResolutionMode.LambdaResolution) {
+            towerDataContextForAnonymousFunctions[anonymousFunction.symbol] = towerDataContext
+        }
+        if (mode is ResolutionMode.ContextDependent || mode is ResolutionMode.ContextDependentDelegate) {
+            return f()
+        }
         return withTowerDataCleanup {
             addLocalScope(FirLocalScope())
             val receiverTypeRef = anonymousFunction.receiverTypeRef
             val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }
             withContainer(anonymousFunction) {
                 withLabelAndReceiverType(labelName, anonymousFunction, receiverTypeRef?.coneType, holder) {
-                    if (isInDependentContext) {
+                    if (mode is ResolutionMode.LambdaResolution) {
                         withLambdaBeingAnalyzedInDependentContext(anonymousFunction.symbol, f)
                     } else {
                         f()

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

@@ -104,14 +104,6 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         }
     }
 
-    protected inline fun <T> withTypeParametersOf(declaration: FirMemberDeclaration, crossinline l: () -> T): T {
-        val scope = createTypeParameterScope(declaration)
-        return context.withTowerDataCleanup {
-            scope?.let { context.addNonLocalTowerDataElement(it.asTowerDataElement(isLocal = false)) }
-            l()
-        }
-    }
-
     override fun transformEnumEntry(enumEntry: FirEnumEntry, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
         if (enumEntry.resolvePhase == transformerPhase) return enumEntry.compose()
         transformer.replaceDeclarationResolvePhaseIfNeeded(enumEntry, transformerPhase)
@@ -424,38 +416,31 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     private fun transformAnonymousFunctionWithLambdaResolution(
         anonymousFunction: FirAnonymousFunction, lambdaResolution: ResolutionMode.LambdaResolution
     ): FirAnonymousFunction {
-        fun transform(): FirAnonymousFunction {
-            val expectedReturnType =
-                lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
-
-            val result = transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
-
-            val body = result.body
-            if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
-                // TODO: This part seems unnecessary because for lambdas in dependent context will be completed and their type
-                //  should be replaced there properly
-                val returnType =
-                    dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(result)
-                        .firstNotNullResult { (it as? FirExpression)?.resultType?.coneTypeSafe() }
-
-                if (returnType != null) {
-                    result.transformReturnTypeRef(transformer, withExpectedType(returnType))
-                } else {
-                    result.transformReturnTypeRef(
-                        transformer,
-                        withExpectedType(buildErrorTypeRef {
-                            diagnostic =
-                                ConeSimpleDiagnostic("Unresolved lambda return type", DiagnosticKind.InferenceError)
-                        })
-                    )
-                }
+        val expectedReturnType =
+            lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
+        val result = transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
+        val body = result.body
+        if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
+            // TODO: This part seems unnecessary because for lambdas in dependent context will be completed and their type
+            //  should be replaced there properly
+            val returnType =
+                dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(result)
+                    .firstNotNullResult { (it as? FirExpression)?.resultType?.coneTypeSafe() }
+
+            if (returnType != null) {
+                result.transformReturnTypeRef(transformer, withExpectedType(returnType))
+            } else {
+                result.transformReturnTypeRef(
+                    transformer,
+                    withExpectedType(buildErrorTypeRef {
+                        diagnostic =
+                            ConeSimpleDiagnostic("Unresolved lambda return type", DiagnosticKind.InferenceError)
+                    })
+                )
             }
-            return result
         }
 
-        return context.withAnonymousFunction(anonymousFunction, components, isInDependentContext = true) {
-            transform()
-        }
+        return result
     }
 
     override fun transformSimpleFunction(
@@ -630,15 +615,18 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             anonymousFunction.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent)
             anonymousFunction.transformReceiverTypeRef(transformer, ResolutionMode.ContextIndependent)
             anonymousFunction.valueParameters.forEach { it.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent) }
-            context.saveContextForAnonymousFunction(anonymousFunction)
         }
         return when (data) {
             is ResolutionMode.ContextDependent, is ResolutionMode.ContextDependentDelegate -> {
-                dataFlowAnalyzer.visitPostponedAnonymousFunction(anonymousFunction)
-                anonymousFunction.addReturn().compose()
+                context.withAnonymousFunction(anonymousFunction, components, data) {
+                    dataFlowAnalyzer.visitPostponedAnonymousFunction(anonymousFunction)
+                    anonymousFunction.addReturn().compose()
+                }
             }
             is ResolutionMode.LambdaResolution -> {
-                transformAnonymousFunctionWithLambdaResolution(anonymousFunction, data).addReturn().compose()
+                context.withAnonymousFunction(anonymousFunction, components, data) {
+                    transformAnonymousFunctionWithLambdaResolution(anonymousFunction, data).addReturn().compose()
+                }
             }
             is ResolutionMode.WithExpectedType, is ResolutionMode.ContextIndependent -> {
                 val expectedTypeRef = (data as? ResolutionMode.WithExpectedType)?.expectedTypeRef ?: buildImplicitTypeRef()
@@ -699,7 +687,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 )
                 lambda = lambda.transformValueParameters(ImplicitToErrorTypeTransformer, null)
                 val bodyExpectedType = returnTypeRefFromResolvedAtom ?: expectedTypeRef
-                context.withAnonymousFunction(lambda, components) {
+                context.withAnonymousFunction(lambda, components, data) {
                     lambda = transformFunction(lambda, withExpectedType(bodyExpectedType)).single as FirAnonymousFunction
                 }
                 // To separate function and separate commit