From c25a5bc9c1e8bf7024f294f6147e203005c8b54b Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Mon, 6 Feb 2023 10:29:51 +0200
Subject: [PATCH] [FIR] Cleanup storing context for lambdas during body
 resolution

---
 .../transformers/body/resolve/BodyResolveContext.kt   | 11 +++++++----
 .../body/resolve/FirDeclarationsResolveTransformer.kt |  5 ++---
 2 files changed, 9 insertions(+), 7 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 fd4e2e19aa6..0ac15418b54 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
@@ -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,
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 dd449ebfdf9..12229a40bd6 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
@@ -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 =