[FIR] Safe implicit receiver stack in lambda atom for resolve of inner lambdas

#KT-36887 Fixed

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir
 
-import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildExpressionStub
@@ -46,7 +45,6 @@ class FirCallResolver(
     components: BodyResolveComponents,
     topLevelScopes: List<FirScope>,
     localScopes: List<FirLocalScope>,
-    override val implicitReceiverStack: ImplicitReceiverStack,
     private val qualifiedResolver: FirQualifiedNameResolver,
 ) : BodyResolveComponents by components {
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -39,7 +39,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         fun createFirDataFlowAnalyzer(
             components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents
         ): FirDataFlowAnalyzer<*> = object : FirDataFlowAnalyzer<PersistentFlow>(components) {
-            private val receiverStack: ImplicitReceiverStackImpl = components.implicitReceiverStack as ImplicitReceiverStackImpl
+            private val receiverStack: ImplicitReceiverStackImpl
+                get() = components.implicitReceiverStack as ImplicitReceiverStackImpl
 
             override val logicSystem: PersistentLogicSystem = object : PersistentLogicSystem(components.inferenceComponents.ctx) {
                 override fun processUpdatedReceiverVariable(flow: PersistentFlow, variable: RealVariable) {

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

@@ -140,14 +140,14 @@ class FirCallCompleter(
 
     private inner class LambdaAnalyzerImpl : LambdaAnalyzer {
         override fun analyzeAndGetLambdaReturnArguments(
-            lambdaArgument: FirAnonymousFunction,
+            lambdaAtom: ResolvedLambdaAtom,
-            isSuspend: Boolean,
             receiverType: ConeKotlinType?,
             parameters: List<ConeKotlinType>,
             expectedReturnType: ConeKotlinType?,
             rawReturnType: ConeKotlinType,
             stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>
         ): ReturnArgumentsAnalysisResult {
+            val lambdaArgument: FirAnonymousFunction = lambdaAtom.atom
             val needItParam = lambdaArgument.valueParameters.isEmpty() && parameters.size == 1
 
             val itParam = when {
@@ -185,7 +185,9 @@ class FirCallCompleter(
             lambdaArgument.replaceValueParameters(lambdaArgument.valueParameters + listOfNotNull(itParam))
             lambdaArgument.replaceReturnTypeRef(expectedReturnTypeRef ?: noExpectedType)
 
-            lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+            transformer.components.withImplicitReceiverStack(lambdaAtom.implicitReceiverStack) {
+                lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+            }
 
             val returnArguments = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambdaArgument)
 

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

@@ -9,12 +9,10 @@ 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.DoubleColonLHS
+import org.jetbrains.kotlin.fir.resolve.*
 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.*
@@ -36,8 +34,8 @@ fun Candidate.preprocessLambdaArgument(
     }
 
     val resolvedArgument =
-        extractLambdaInfoFromFunctionalType(expectedType, expectedTypeRef, argument, bodyResolveComponents.session)
+        extractLambdaInfoFromFunctionalType(expectedType, expectedTypeRef, argument, bodyResolveComponents.session, bodyResolveComponents)
-            ?: extraLambdaInfo(expectedType, argument, csBuilder, bodyResolveComponents.session)
+            ?: extraLambdaInfo(expectedType, argument, csBuilder, bodyResolveComponents.session, bodyResolveComponents)
 
     if (expectedType != null) {
         // TODO: add SAM conversion processing
@@ -106,7 +104,8 @@ private fun extraLambdaInfo(
     expectedType: ConeKotlinType?,
     argument: FirAnonymousFunction,
     csBuilder: ConstraintSystemBuilder,
-    session: FirSession
+    session: FirSession,
+    components: BodyResolveComponents
 ): ResolvedLambdaAtom {
     val isSuspend = expectedType?.isSuspendFunctionType(session) ?: false
 
@@ -130,18 +129,27 @@ private fun extraLambdaInfo(
     val newTypeVariableUsed = returnType == typeVariable.defaultType
     if (newTypeVariableUsed) csBuilder.registerVariable(typeVariable)
 
-    return ResolvedLambdaAtom(argument, isSuspend, receiverType, parameters, returnType, typeVariable.takeIf { newTypeVariableUsed })
+    return ResolvedLambdaAtom(
+        argument,
+        isSuspend,
+        receiverType,
+        parameters,
+        returnType,
+        typeVariable.takeIf { newTypeVariableUsed },
+        components.implicitReceiverStack.snapshot()
+    )
 }
 
 internal fun extractLambdaInfoFromFunctionalType(
     expectedType: ConeKotlinType?,
     expectedTypeRef: FirTypeRef,
     argument: FirAnonymousFunction,
-    session: FirSession
+    session: FirSession,
+    components: BodyResolveComponents
 ): ResolvedLambdaAtom? {
     if (expectedType == null) return null
     if (expectedType is ConeFlexibleType) {
-        return extractLambdaInfoFromFunctionalType(expectedType.lowerBound, expectedTypeRef, argument, session)
+        return extractLambdaInfoFromFunctionalType(expectedType.lowerBound, expectedTypeRef, argument, session, components)
     }
     if (!expectedType.isBuiltinFunctionalType(session)) return null
 
@@ -155,7 +163,8 @@ internal fun extractLambdaInfoFromFunctionalType(
         receiverType,
         parameters,
         returnType,
-        typeVariableForLambdaReturnType = null
+        typeVariableForLambdaReturnType = null,
+        components.implicitReceiverStack.snapshot()
     )
 }
 
@@ -199,7 +208,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

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.resolve.inference
 
 import org.jetbrains.kotlin.fir.FirCallResolver
-import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
@@ -32,8 +31,7 @@ data class ReturnArgumentsAnalysisResult(
 
 interface LambdaAnalyzer {
     fun analyzeAndGetLambdaReturnArguments(
-        lambdaArgument: FirAnonymousFunction,
+        lambdaAtom: ResolvedLambdaAtom,
-        isSuspend: Boolean,
         receiverType: ConeKotlinType?,
         parameters: List<ConeKotlinType>,
         expectedReturnType: ConeKotlinType?, // null means, that return type is not proper i.e. it depends on some type variables
@@ -128,8 +126,7 @@ class PostponedArgumentsAnalyzer(
         }
 
         val (returnArguments, inferenceSession) = lambdaAnalyzer.analyzeAndGetLambdaReturnArguments(
-            lambda.atom,
+            lambda,
-            lambda.isSuspend,
             receiver,
             parameters,
             expectedTypeForReturnArguments,

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

@@ -106,7 +106,10 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         override val symbolProvider: FirSymbolProvider = session.firSymbolProvider
 
         override val returnTypeCalculator: ReturnTypeCalculator = transformer.returnTypeCalculator
-        override val implicitReceiverStack: ImplicitReceiverStack = ImplicitReceiverStackImpl()
+
+        @set:PrivateForInline
+        override var implicitReceiverStack: ImplicitReceiverStack = ImplicitReceiverStackImpl()
+
         override val inferenceComponents: InferenceComponents = inferenceComponents(session, returnTypeCalculator, scopeSession)
         override val resolutionStageRunner: ResolutionStageRunner = ResolutionStageRunner(inferenceComponents)
         override val samResolver: FirSamResolver = FirSamResolverImpl(session, scopeSession)
@@ -116,7 +119,6 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             this,
             topLevelScopes,
             localScopes,
-            implicitReceiverStack,
             qualifiedResolver
         )
         val typeResolverTransformer = FirSpecificTypeResolverTransformer(
@@ -147,6 +149,17 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             containerIfAny = prevContainer
             return result
         }
+
+        @UseExperimental(PrivateForInline::class)
+        inline fun <T> withImplicitReceiverStack(implicitReceiverStack: ImplicitReceiverStack, f: () -> T): T {
+            val existedStack = this.implicitReceiverStack
+            this.implicitReceiverStack = implicitReceiverStack
+            return try {
+                f()
+            } finally {
+                this.implicitReceiverStack = existedStack
+            }
+        }
     }
 }
 

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

@@ -489,7 +489,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
                 val expectedTypeRef = (data as? ResolutionMode.WithExpectedType)?.expectedTypeRef ?: buildImplicitTypeRef()
                 val resolvedLambdaAtom = (expectedTypeRef as? FirResolvedTypeRef)?.let {
                     extractLambdaInfoFromFunctionalType(
-                        it.type, it, anonymousFunction, session
+                        it.type, it, anonymousFunction, session, components
                     )
                 }
                 var af = anonymousFunction

compiler/fir/resolve/testData/resolveWithStdlib/withInInitializer.kt

@@ -7,7 +7,7 @@ class Second {
 
     val test = with(data) {
         list.filterIsInstance<Int>().filter {
-            it == <!UNRESOLVED_REFERENCE!>member<!>
+            it == member
         }
     }
 }

compiler/fir/resolve/testData/resolveWithStdlib/withInInitializer.txt

@@ -21,7 +21,7 @@ FILE: withInInitializer.kt
 
         public final val test: R|kotlin/collections/List<kotlin/Int>| = R|kotlin/with|<R|First|, R|kotlin/collections/List<kotlin/Int>|>(this@R|/Second|.R|/Second.data|, <L> = with@fun R|First|.<anonymous>(): R|kotlin/collections/List<kotlin/Int>| <kind=EXACTLY_ONCE>  {
             ^ this@R|/Second|.R|/Second.list|.R|kotlin/collections/filterIsInstance|<R|kotlin/Int|>().R|kotlin/collections/filter|<R|kotlin/Int|>(<L> = filter@fun <anonymous>(it: R|kotlin/Int|): R|kotlin/Boolean| <kind=UNKNOWN>  {
-                ^ ==(R|<local>/it|, <Unresolved name: member>#)
+                ^ ==(R|<local>/it|, this@R|special/anonymous|.R|/First.member|)
             }
             )
         }

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

@@ -288,6 +288,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         runTest("compiler/fir/resolve/testData/resolveWithStdlib/whenAsLambdaReturnStatement.kt");
     }
 
+    @TestMetadata("withInInitializer.kt")
+    public void testWithInInitializer() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolveWithStdlib/withInInitializer.kt");
+    }
+
     @TestMetadata("compiler/fir/resolve/testData/resolveWithStdlib/callableReferences")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
@@ -714,11 +719,6 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         public void testWeakHashMap() throws Exception {
             runTest("compiler/fir/resolve/testData/resolveWithStdlib/problems/weakHashMap.kt");
         }
-
-        @TestMetadata("withInInitializer.kt")
-        public void testWithInInitializer() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolveWithStdlib/problems/withInInitializer.kt");
-        }
     }
 
     @TestMetadata("compiler/fir/resolve/testData/resolveWithStdlib/smartcasts")