[FIR] Add saving of local context in anonymous functions for further resolve

#KT-37066 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SessionHolder.kt

@@ -9,6 +9,7 @@ import kotlinx.collections.immutable.PersistentList
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSymbolOwner
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -20,10 +21,18 @@ import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 
 typealias FirLocalScopes = PersistentList<FirLocalScope>
 
+class FirLocalContext(
+    val localScopes: FirLocalScopes,
+    val implicitReceiverStack: ImplicitReceiverStack
+)
+
+typealias LocalContextForAnonymousFunctions = Map<FirAnonymousFunctionSymbol, FirLocalContext>
+
 interface SessionHolder {
     val session: FirSession
 }
@@ -33,6 +42,7 @@ interface BodyResolveComponents : SessionHolder {
     val implicitReceiverStack: ImplicitReceiverStack
     val topLevelScopes: List<FirScope>
     val localScopes: FirLocalScopes
+    val localContextForAnonymousFunctions: LocalContextForAnonymousFunctions
     val noExpectedType: FirTypeRef
     val symbolProvider: FirSymbolProvider
     val file: FirFile
@@ -51,4 +61,7 @@ interface BodyResolveComponents : SessionHolder {
 
     val <D> AbstractFirBasedSymbol<D>.phasedFir: D where D : FirDeclaration, D : FirSymbolOwner<D>
         get() = phasedFir(FirResolvePhase.DECLARATIONS)
-}
+
+    fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction)
+    fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction)
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -185,9 +185,11 @@ class FirCallCompleter(
             lambdaArgument.replaceValueParameters(lambdaArgument.valueParameters + listOfNotNull(itParam))
             lambdaArgument.replaceReturnTypeRef(expectedReturnTypeRef ?: noExpectedType)
 
-            transformer.components.withImplicitReceiverStack(lambdaAtom.implicitReceiverStack) {
+            val localContext = localContextForAnonymousFunctions.getValue(lambdaArgument.symbol)
+            transformer.components.withLocalContext(localContext) {
                 lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
             }
+            dropContextForAnonymousFunction(lambdaArgument)
 
             val returnArguments = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambdaArgument)
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt

@@ -9,10 +9,13 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.DoubleColonLHS
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.CandidateApplicability
 import org.jetbrains.kotlin.fir.resolve.calls.isExtensionFunctionType
+import org.jetbrains.kotlin.fir.resolve.createFunctionalType
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.*
@@ -135,8 +138,7 @@ private fun extraLambdaInfo(
         receiverType,
         parameters,
         returnType,
-        typeVariable.takeIf { newTypeVariableUsed },
+        typeVariable.takeIf { newTypeVariableUsed }
-        components.implicitReceiverStack.snapshot()
     )
 }
 
@@ -163,8 +165,7 @@ internal fun extractLambdaInfoFromFunctionalType(
         receiverType,
         parameters,
         returnType,
-        typeVariableForLambdaReturnType = null,
+        typeVariableForLambdaReturnType = null
-        components.implicitReceiverStack.snapshot()
     )
 }
 
@@ -208,10 +209,8 @@ class ResolvedLambdaAtom(
     val receiver: ConeKotlinType?,
     val parameters: List<ConeKotlinType>,
     val returnType: ConeKotlinType,
-    val typeVariableForLambdaReturnType: TypeVariableForLambdaReturnType?,
+    val typeVariableForLambdaReturnType: TypeVariableForLambdaReturnType?
-    val implicitReceiverStack: ImplicitReceiverStack
 ) : PostponedResolvedAtomMarker {
-
     override var analyzed: Boolean = false
     lateinit var returnStatements: List<FirStatement>
 

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirTypeResolveScopeForBodyResolve
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 
@@ -110,6 +111,8 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
 
         override val noExpectedType: FirTypeRef = buildImplicitTypeRef()
 
+        override val localContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirLocalContext> = mutableMapOf()
+
         @set:PrivateForInline
         override lateinit var file: FirFile
             internal set
@@ -160,24 +163,18 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         }
 
         @UseExperimental(PrivateForInline::class)
-        inline fun <T> withImplicitReceiverStack(implicitReceiverStack: ImplicitReceiverStack, f: () -> T): T {
+        inline fun <T> withLocalContext(localContext: FirLocalContext, f: () -> T): T {
             val existedStack = this.implicitReceiverStack
-            this.implicitReceiverStack = implicitReceiverStack
+            val existedLocalScopes = this.localScopes
+
+            implicitReceiverStack = localContext.implicitReceiverStack
+            localScopes = localContext.localScopes
+
             return try {
                 f()
             } finally {
-                this.implicitReceiverStack = existedStack
+                implicitReceiverStack = existedStack
-            }
+                localScopes = existedLocalScopes
-        }
-
-        @UseExperimental(PrivateForInline::class)
-        inline fun <R> withLocalScopes(localScopes: FirLocalScopes, l: () -> R): R {
-            val initialLocalScopes = this.localScopes
-            this.localScopes = localScopes
-            return try {
-                l()
-            } finally {
-                this.localScopes = initialLocalScopes
             }
         }
 
@@ -212,6 +209,18 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             updateLastScope { storeBackingField(property) }
         }
 
+        override fun saveContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+            localContextForAnonymousFunctions[anonymousFunction.symbol] = FirLocalContext(localScopes, implicitReceiverStack.snapshot())
+        }
+
+        override fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+            localContextForAnonymousFunctions.remove(anonymousFunction.symbol)
+        }
+
+        fun cleanContextForAnonymousFunction() {
+            localContextForAnonymousFunctions.clear()
+        }
+
         @UseExperimental(PrivateForInline::class)
         private inline fun updateLastScope(transform: FirLocalScope.() -> FirLocalScope) {
             val lastScope = localScopes.lastOrNull() ?: return

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

@@ -37,6 +37,7 @@ open class FirBodyResolveTransformer(
     private val controlFlowStatementsTransformer = FirControlFlowStatementsResolveTransformer(this)
 
     override fun transformFile(file: FirFile, data: ResolutionMode): CompositeTransformResult<FirFile> {
+        components.cleanContextForAnonymousFunction()
         @UseExperimental(PrivateForInline::class)
         components.file = file
         packageFqName = file.packageFqName

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

@@ -477,6 +477,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             anonymousFunction.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent)
             anonymousFunction.transformReceiverTypeRef(transformer, ResolutionMode.ContextIndependent)
             anonymousFunction.valueParameters.forEach { it.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent) }
+            components.saveContextForAnonymousFunction(anonymousFunction)
         }
         return when (data) {
             ResolutionMode.ContextDependent -> {

compiler/fir/resolve/testData/resolve/inference/lambdaAsReturnStatementOfLambda.kt

@@ -0,0 +1,15 @@
+// ISSUE: KT-37070
+
+class A
+
+fun test(a: A) {
+
+    val lambda = a.let {
+        { it }
+    }
+
+    val alsoA = lambda()
+    takeA(alsoA)
+}
+
+fun takeA(a: A) {}

compiler/fir/resolve/testData/resolve/inference/lambdaAsReturnStatementOfLambda.txt

@@ -0,0 +1,20 @@
+FILE: lambdaAsReturnStatementOfLambda.kt
+    public final class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final fun test(a: R|A|): R|kotlin/Unit| {
+        lval lambda: R|() -> A| = R|<local>/a|.R|kotlin/let|<R|A|, R|() -> A|>(<L> = let@fun <anonymous>(it: R|A|): R|() -> A| <kind=EXACTLY_ONCE>  {
+            ^ let@fun <anonymous>(): R|A| {
+                ^ R|<local>/it|
+            }
+
+        }
+        )
+        lval alsoA: R|A| = R|<local>/lambda|.R|FakeOverride<kotlin/Function0.invoke: R|A|>|()
+        R|/takeA|(R|<local>/alsoA|)
+    }
+    public final fun takeA(a: R|A|): R|kotlin/Unit| {
+    }

compiler/fir/resolve/testData/resolve/lambdaArgInScopeFunction.kt

@@ -18,7 +18,7 @@ fun case1(kotlinClass: KotlinClass?) {
     value.checkType { _<KotlinClass?>() }
 
     val lambda = kotlinClass?.let {
-        {<!UNRESOLVED_REFERENCE!>it<!>}
+        {it}
     }
 
     lambda.checkType { <!INAPPLICABLE_CANDIDATE!>_<!><Function1<Unit, KotlinClass?>>() }
@@ -33,7 +33,7 @@ fun case2(kotlinClass: KotlinClass) {
     value.checkType { _<KotlinClass>() }
 
     val lambda = kotlinClass.let {
-        {<!UNRESOLVED_REFERENCE!>it<!>}
+        {it}
     }
 
     lambda.checkType { <!INAPPLICABLE_CANDIDATE!>_<!><Function1<Unit, KotlinClass?>>() }

compiler/fir/resolve/testData/resolve/lambdaArgInScopeFunction.txt

@@ -21,14 +21,14 @@ FILE: lambdaArgInScopeFunction.kt
             this@R|special/anonymous|.R|tests/_checkType/_|<R|KotlinClass?|>()
         }
         )
-        lval lambda: R|kotlin/Nothing?| = R|<local>/kotlinClass|?.R|kotlin/let|<R|KotlinClass|, R|kotlin/Nothing|>(<L> = let@fun <anonymous>(it: R|KotlinClass|): R|kotlin/Nothing| <kind=EXACTLY_ONCE>  {
+        lval lambda: R|() -> KotlinClass| = R|<local>/kotlinClass|?.R|kotlin/let|<R|KotlinClass|, R|() -> KotlinClass|>(<L> = let@fun <anonymous>(it: R|KotlinClass|): R|() -> KotlinClass| <kind=EXACTLY_ONCE>  {
-            ^ let@fun <anonymous>(): <ERROR TYPE REF: Unresolved name: it> {
+            ^ let@fun <anonymous>(): R|KotlinClass| {
-                ^ <Unresolved name: it>#
+                ^ R|<local>/it|
             }
 
         }
         )
-        R|<local>/lambda|.R|tests/_checkType/checkType|<R|kotlin/Nothing?|>(<L> = checkType@fun R|tests/_checkType/Inv<kotlin/Nothing?>|.<anonymous>(): R|kotlin/Unit| {
+        R|<local>/lambda|.R|tests/_checkType/checkType|<R|() -> KotlinClass|>(<L> = checkType@fun R|tests/_checkType/Inv<kotlin/Function0<KotlinClass>?>|.<anonymous>(): R|kotlin/Unit| {
             ^ <Inapplicable(WRONG_RECEIVER): [tests/_checkType/_]>#<R|(kotlin/Unit) -> KotlinClass?|>()
         }
         )
@@ -42,14 +42,14 @@ FILE: lambdaArgInScopeFunction.kt
             this@R|special/anonymous|.R|tests/_checkType/_|<R|KotlinClass|>()
         }
         )
-        lval lambda: R|kotlin/Nothing| = R|<local>/kotlinClass|.R|kotlin/let|<R|KotlinClass|, R|kotlin/Nothing|>(<L> = let@fun <anonymous>(it: R|KotlinClass|): R|kotlin/Nothing| <kind=EXACTLY_ONCE>  {
+        lval lambda: R|() -> KotlinClass| = R|<local>/kotlinClass|.R|kotlin/let|<R|KotlinClass|, R|() -> KotlinClass|>(<L> = let@fun <anonymous>(it: R|KotlinClass|): R|() -> KotlinClass| <kind=EXACTLY_ONCE>  {
-            ^ let@fun <anonymous>(): <ERROR TYPE REF: Unresolved name: it> {
+            ^ let@fun <anonymous>(): R|KotlinClass| {
-                ^ <Unresolved name: it>#
+                ^ R|<local>/it|
             }
 
         }
         )
-        R|<local>/lambda|.R|tests/_checkType/checkType|<R|kotlin/Nothing|>(<L> = checkType@fun R|tests/_checkType/Inv<kotlin/Nothing>|.<anonymous>(): R|kotlin/Unit| {
+        R|<local>/lambda|.R|tests/_checkType/checkType|<R|() -> KotlinClass|>(<L> = checkType@fun R|tests/_checkType/Inv<kotlin/Function0<KotlinClass>>|.<anonymous>(): R|kotlin/Unit| {
             ^ <Inapplicable(WRONG_RECEIVER): [tests/_checkType/_]>#<R|(kotlin/Unit) -> KotlinClass?|>()
         }
         )

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -1200,6 +1200,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/inference/definitelyNotNullIntersectionType.kt");
         }
 
+        @TestMetadata("lambdaAsReturnStatementOfLambda.kt")
+        public void testLambdaAsReturnStatementOfLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/inference/lambdaAsReturnStatementOfLambda.kt");
+        }
+
         @TestMetadata("nestedLambdas.kt")
         public void testNestedLambdas() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/inference/nestedLambdas.kt");

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -1200,6 +1200,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/inference/definitelyNotNullIntersectionType.kt");
         }
 
+        @TestMetadata("lambdaAsReturnStatementOfLambda.kt")
+        public void testLambdaAsReturnStatementOfLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/inference/lambdaAsReturnStatementOfLambda.kt");
+        }
+
         @TestMetadata("nestedLambdas.kt")
         public void testNestedLambdas() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/inference/nestedLambdas.kt");

compiler/testData/diagnostics/tests/inference/constraints/returnLambdaFromLambda.fir.kt

@@ -6,14 +6,14 @@ fun testLambda() {
         if (x is String) return@myRun { it -> x.length <!AMBIGUITY!>+<!> it }
         if (x !is Int) return@myRun { it -> it }
 
-        { it -> <!UNRESOLVED_REFERENCE!>x<!> + it }
+        { it -> x <!AMBIGUITY!>+<!> it }
     }
 
     val twoLambda: (Int) -> Int = myRun {
         val x: Int = 1
         run {
             val y: Int = 2
-            { x <!AMBIGUITY!>+<!> <!UNRESOLVED_REFERENCE!>y<!> }
+            { x + y }
         }
     }