FIR: Analyze return statements for implicit lambda with independent context

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.sure
 
@@ -56,6 +57,8 @@ class BodyResolveContext(
     @set:PrivateForInline
     var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
 
+    val anonymousFunctionsAnalyzedInDependentContext: MutableSet<FirFunctionSymbol<*>> = mutableSetOf()
+
     @OptIn(PrivateForInline::class)
     inline fun <R> withInferenceSession(inferenceSession: FirInferenceSession, block: () -> R): R {
         val oldSession = this.inferenceSession
@@ -125,6 +128,16 @@ class BodyResolveContext(
         }
     }
 
+    @OptIn(PrivateForInline::class)
+    inline fun <R> withLambdaBeingAnalyzedInDependentContext(lambda: FirAnonymousFunctionSymbol, l: () -> R): R {
+        anonymousFunctionsAnalyzedInDependentContext.add(lambda)
+        return try {
+            l()
+        } finally {
+            anonymousFunctionsAnalyzedInDependentContext.remove(lambda)
+        }
+    }
+
     @OptIn(PrivateForInline::class)
     fun replaceTowerDataContext(newContext: FirTowerDataContext) {
         towerDataContext = newContext
@@ -189,5 +202,6 @@ class BodyResolveContext(
         towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
         containers = this@BodyResolveContext.containers
         towerDataContext = this@BodyResolveContext.towerDataContext
+        anonymousFunctionsAnalyzedInDependentContext.addAll(this@BodyResolveContext.anonymousFunctionsAnalyzedInDependentContext)
     }
 }

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

@@ -342,6 +342,13 @@ open class FirBodyResolveTransformer(
         return controlFlowStatementsTransformer.transformJump(jump, data)
     }
 
+    override fun transformReturnExpression(
+        returnExpression: FirReturnExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return controlFlowStatementsTransformer.transformReturnExpression(returnExpression, data)
+    }
+
     override fun transformThrowExpression(
         throwExpression: FirThrowExpression,
         data: ResolutionMode

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

@@ -187,18 +187,30 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     // ------------------------------- Jumps -------------------------------
 
     override fun <E : FirTargetElement> transformJump(jump: FirJump<E>, data: ResolutionMode): CompositeTransformResult<FirStatement> {
-        val expectedTypeRef = (jump as? FirReturnExpression)?.target?.labeledElement?.returnTypeRef
-
-        val mode = if (expectedTypeRef != null) {
-            ResolutionMode.WithExpectedType(expectedTypeRef)
-        } else {
-            ResolutionMode.ContextIndependent
-        }
-        val result = transformer.transformExpression(jump, mode).single
+        val result = transformer.transformExpression(jump, data).single
         dataFlowAnalyzer.exitJump(jump)
         return result.compose()
     }
 
+    override fun transformReturnExpression(
+        returnExpression: FirReturnExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        val labeledElement = returnExpression.target.labeledElement
+        val expectedTypeRef = labeledElement.returnTypeRef
+        @Suppress("IntroduceWhenSubject")
+        val mode = when {
+            labeledElement.symbol in context.anonymousFunctionsAnalyzedInDependentContext -> {
+                ResolutionMode.ContextDependent
+            }
+            else -> {
+                ResolutionMode.WithExpectedType(expectedTypeRef)
+            }
+        }
+
+        return transformJump(returnExpression, mode)
+    }
+
     override fun transformThrowExpression(
         throwExpression: FirThrowExpression,
         data: ResolutionMode

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

@@ -435,7 +435,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         fun transform(): FirAnonymousFunction {
             val expectedReturnType =
                 lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
-            val result = transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
+
+            val result = context.withLambdaBeingAnalyzedInDependentContext(anonymousFunction.symbol) {
+                transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
+            }
+
             val body = result.body
             if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
                 result.transformReturnTypeRef(transformer, withExpectedType(body.resultType))

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

@@ -3,7 +3,7 @@
 fun testLambda() {
     val basicTest: (Int) -> Int = myRun {
         val x: Any? = null
-        if (x is String) return@myRun { it -> x.length <!AMBIGUITY!>+<!> it }
+        if (x is String) return@myRun { it -> x.length + it }
         if (x !is Int) return@myRun { it -> it }
 
         { it -> x + it }