[FIR] Cleanup storing context for lambdas during body resolution

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

@@ -649,11 +649,9 @@ class BodyResolveContext(
         mode: ResolutionMode,
         f: () -> T
     ): T {
+        require(mode !is ResolutionMode.ContextDependent && mode !is ResolutionMode.ContextDependentDelegate)
         if (mode !is ResolutionMode.LambdaResolution) {
-            specialTowerDataContexts.storeAnonymousFunctionContext(anonymousFunction.symbol, towerDataContext)
-        }
-        if (mode is ResolutionMode.ContextDependent || mode is ResolutionMode.ContextDependentDelegate) {
-            return f()
+            storeContextForAnonymousFunction(anonymousFunction)
         }
         return withTowerDataCleanup {
             addLocalScope(FirLocalScope(holder.session))
@@ -671,6 +669,11 @@ class BodyResolveContext(
         }
     }
 
+    @OptIn(PrivateForInline::class)
+    fun storeContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+        specialTowerDataContexts.storeAnonymousFunctionContext(anonymousFunction.symbol, towerDataContext)
+    }
+
     @OptIn(PrivateForInline::class)
     inline fun <T> withField(
         field: FirField,

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

@@ -718,9 +718,8 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         }
         return when (data) {
             is ResolutionMode.ContextDependent, is ResolutionMode.ContextDependentDelegate -> {
-                context.withAnonymousFunction(anonymousFunction, components, data) {
-                    anonymousFunction
-                }
+                context.storeContextForAnonymousFunction(anonymousFunction)
+                anonymousFunction
             }
             is ResolutionMode.LambdaResolution -> {
                 val expectedReturnTypeRef =