FIR: Do not run completion for lambda's explicit return too early

See the test case
Completion for synthetic call: x ?: return@myRun materialize()
makes `materialize()` while it's obviously too early for that

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -2223,6 +2223,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdaInElvis.kt");
         }
 
+        @TestMetadata("lambdasReturns.kt")
+        public void testLambdasReturns() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdasReturns.kt");
+        }
+
         @TestMetadata("nestedExtensionFunctionType.kt")
         public void testNestedExtensionFunctionType() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/nestedExtensionFunctionType.kt");

compiler/fir/analysis-tests/testData/resolve/inference/lambdasReturns.fir.txt

+FILE: lambdasReturns.kt
+    public final fun <R> myRun(b: R|() -> R|): R|R| {
+        ^myRun R|<local>/b|.R|SubstitutionOverride<kotlin/Function0.invoke: R|R|>|()
+    }
+    public final fun <T> materialize(): R|T| {
+        ^materialize R|kotlin/TODO|()
+    }
+    public final fun foo(x: R|kotlin/String?|): R|kotlin/Unit| {
+        lval r: R|kotlin/Int| = R|/myRun|<R|kotlin/Int|>(<L> = myRun@fun <anonymous>(): R|kotlin/Int| {
+            lval y: R|kotlin/String| = R|<local>/x| ?: ^@myRun R?C|/materialize|()
+            ^ R|<local>/y|.R|kotlin/String.length|
+        }
+        )
+        R|<local>/r|.R|kotlin/Int.minus|(Int(1))
+    }

compiler/fir/analysis-tests/testData/resolve/inference/lambdasReturns.kt

+// !DIAGNOSTICS: -UNUSED_VARIABLE
+// !WITH_NEW_INFERENCE
+
+fun <R> myRun(b: () -> R): R = b()
+
+fun <T> materialize(): T = TODO()
+
+fun foo(x: String?) {
+    val r = myRun {
+        val y = x ?: return@myRun materialize()
+        y.length
+    }
+
+    r.minus(1)
+}

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -2547,6 +2547,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdaInElvis.kt");
             }
 
+            @Test
+            @TestMetadata("lambdasReturns.kt")
+            public void testLambdasReturns() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdasReturns.kt");
+            }
+
             @Test
             @TestMetadata("nestedExtensionFunctionType.kt")
             public void testNestedExtensionFunctionType() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -2568,6 +2568,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdaInElvis.kt");
         }
 
+        @Test
+        @TestMetadata("lambdasReturns.kt")
+        public void testLambdasReturns() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/inference/lambdasReturns.kt");
+        }
+
         @Test
         @TestMetadata("nestedExtensionFunctionType.kt")
         public void testNestedExtensionFunctionType() throws Exception {

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

@@ -83,7 +83,8 @@ class FirCallCompleter(
                     val completedCall = call.transformSingle(
                         FirCallCompletionResultsWriterTransformer(
                             session, finalSubstitutor, components.returnTypeCalculator,
-                            session.inferenceComponents.approximator
+                            session.inferenceComponents.approximator,
+                            components.dataFlowAnalyzer,
                         ),
                         null
                     )
@@ -153,6 +154,7 @@ class FirCallCompleter(
         return FirCallCompletionResultsWriterTransformer(
             session, substitutor, components.returnTypeCalculator,
             session.inferenceComponents.approximator,
+            components.dataFlowAnalyzer,
             mode
         )
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -13,18 +13,16 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirErrorReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.varargElementType
-import org.jetbrains.kotlin.fir.resolve.constructFunctionalTypeRef
-import org.jetbrains.kotlin.fir.resolve.createFunctionalType
-import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.inference.returnType
-import org.jetbrains.kotlin.fir.resolve.propagateTypeFromQualifiedAccessAfterNullCheck
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirArrayOfCallTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.remapArgumentsWithVararg
@@ -48,6 +46,7 @@ class FirCallCompletionResultsWriterTransformer(
     private val finalSubstitutor: ConeSubstitutor,
     private val typeCalculator: ReturnTypeCalculator,
     private val typeApproximator: AbstractTypeApproximator,
+    private val dataFlowAnalyzer: FirDataFlowAnalyzer<*>,
     private val mode: Mode = Mode.Normal
 ) : FirAbstractTreeTransformer<ExpectedArgumentType?>(phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
 
@@ -472,6 +471,7 @@ class FirCallCompletionResultsWriterTransformer(
         }
 
         val result = transformElement(anonymousFunction, null)
+
         val resultFunction = result.single
         if (resultFunction.returnTypeRef.coneTypeSafe<ConeIntegerLiteralType>() != null) {
             val blockType = resultFunction.body?.typeRef?.coneTypeSafe<ConeKotlinType>()
@@ -480,9 +480,26 @@ class FirCallCompletionResultsWriterTransformer(
                 resultFunction.constructFunctionalTypeRef(isSuspend = expectedType?.isSuspendFunctionType(session) == true)
             )
         }
+
+        for (expression in dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(anonymousFunction)) {
+            expression.transform<FirElement, ExpectedArgumentType?>(this, null)
+        }
+
         return result
     }
 
+    override fun transformReturnExpression(
+        returnExpression: FirReturnExpression,
+        data: ExpectedArgumentType?
+    ): CompositeTransformResult<FirStatement> {
+        val labeledElement = returnExpression.target.labeledElement
+        if (labeledElement is FirAnonymousFunction) {
+            return returnExpression.compose()
+        }
+
+        return super.transformReturnExpression(returnExpression, data)
+    }
+
     override fun transformBlock(block: FirBlock, data: ExpectedArgumentType?): CompositeTransformResult<FirStatement> {
         val initialType = block.resultType.coneTypeSafe<ConeKotlinType>()
         if (initialType != null) {

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

@@ -759,6 +759,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     ConeSubstitutor.Empty,
                     components.returnTypeCalculator,
                     inferenceComponents.approximator,
+                    dataFlowAnalyzer,
                 )
                 lambda.transformSingle(writer, expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType())
                 val returnTypes = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda)