FIR: Do not use return statement for type of a block expression

Type of a block is a kind of irrelevant for lambdas: their type is much
more complicated and defined via FirDataFlowAnalyzer#returnExpressionsOfAnonymousFunction
at at FirCallCompleter.LambdaAnalyzerImpl#analyzeAndGetLambdaReturnArguments

compiler/fir/analysis-tests/testData/resolve/arguments/integerLiteralTypes.fir.txt

@@ -30,11 +30,11 @@ FILE: integerLiteralTypes.kt
         }
         ))
         R|/takeByte|(R|kotlin/run|<R|kotlin/Byte|>(<L> = run@fun <anonymous>(): R|kotlin/Byte| <kind=EXACTLY_ONCE>  {
-            ^ Int(1)
+            ^ Byte(1)
         }
         ))
         R|/takeLong|(R|kotlin/run|<R|kotlin/Long|>(<L> = run@fun <anonymous>(): R|kotlin/Long| <kind=EXACTLY_ONCE>  {
-            ^ Int(1)
+            ^ Long(1)
         }
         ))
     }

compiler/fir/analysis-tests/testData/resolve/cfg/lambdaAsReturnOfLambda.dot

@@ -44,21 +44,14 @@ digraph lambdaAsReturnOfLambda_kt {
             7 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
         }
         18 [label="Postponed exit from lambda"];
-        subgraph cluster_6 {
-            color=blue
-            19 [label="Enter block"];
-            20 [label="Exit block"];
-        }
-        21 [label="Function call: R|/run|<R|(kotlin/String) -> kotlin/Unit|>(...)"];
-        22 [label="Exit property" style="filled" fillcolor=red];
+        19 [label="Function call: R|/run|<R|(kotlin/String) -> kotlin/Unit|>(...)"];
+        20 [label="Exit property" style="filled" fillcolor=red];
     }
     16 -> {17};
     17 -> {18 0};
     17 -> {0} [style=dashed];
     18 -> {19};
     19 -> {20};
-    20 -> {21};
-    21 -> {22};
     0 -> {1};
     1 -> {2};
     2 -> {3 8};
@@ -74,39 +67,39 @@ digraph lambdaAsReturnOfLambda_kt {
     11 -> {12};
     12 -> {13};
 
-    subgraph cluster_7 {
+    subgraph cluster_6 {
         color=red
-        23 [label="Enter function bar" style="filled" fillcolor=red];
-        subgraph cluster_8 {
+        21 [label="Enter function bar" style="filled" fillcolor=red];
+        subgraph cluster_7 {
             color=blue
-            24 [label="Enter block"];
-            25 [label="Exit block"];
+            22 [label="Enter block"];
+            23 [label="Exit block"];
         }
-        26 [label="Exit function bar" style="filled" fillcolor=red];
+        24 [label="Exit function bar" style="filled" fillcolor=red];
     }
+    21 -> {22};
+    22 -> {23};
     23 -> {24};
-    24 -> {25};
-    25 -> {26};
 
-    subgraph cluster_9 {
+    subgraph cluster_8 {
         color=red
-        27 [label="Enter function run" style="filled" fillcolor=red];
-        subgraph cluster_10 {
+        25 [label="Enter function run" style="filled" fillcolor=red];
+        subgraph cluster_9 {
             color=blue
-            28 [label="Enter block"];
-            29 [label="Function call: R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|R|>|()"];
-            30 [label="Jump: ^run R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|R|>|()"];
-            31 [label="Stub" style="filled" fillcolor=gray];
-            32 [label="Exit block" style="filled" fillcolor=gray];
+            26 [label="Enter block"];
+            27 [label="Function call: R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|R|>|()"];
+            28 [label="Jump: ^run R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|R|>|()"];
+            29 [label="Stub" style="filled" fillcolor=gray];
+            30 [label="Exit block" style="filled" fillcolor=gray];
         }
-        33 [label="Exit function run" style="filled" fillcolor=red];
+        31 [label="Exit function run" style="filled" fillcolor=red];
     }
+    25 -> {26};
+    26 -> {27};
     27 -> {28};
-    28 -> {29};
-    29 -> {30};
-    30 -> {33};
+    28 -> {31};
+    28 -> {29} [style=dotted];
+    29 -> {30} [style=dotted];
     30 -> {31} [style=dotted];
-    31 -> {32} [style=dotted];
-    32 -> {33} [style=dotted];
 
 }

compiler/fir/analysis-tests/testData/resolve/expresssions/access.fir.txt

@@ -36,7 +36,7 @@ FILE: access.kt
             ^plus String()
         }
 
-        public final fun R|Foo|.check(): R|ERROR CLASS: Ambiguity: plus, [kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus]| {
+        public final fun R|Foo|.check(): <ERROR TYPE REF: Ambiguity: plus, [kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus]> {
             ^check this@R|/Bar.check|.R|/Foo.abc|().<Ambiguity: plus, [kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus, kotlin/Int.plus]>#(this@R|/Bar|.R|/Bar.bar|())
         }
 

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

@@ -7,7 +7,7 @@ FILE: lambdasReturns.kt
     }
     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|()
+            lval y: R|kotlin/String| = R|<local>/x| ?: ^@myRun R|/materialize|<R|kotlin/Int|>()
             ^ R|<local>/y|.R|kotlin/String.length|
         }
         )

compiler/fir/analysis-tests/testData/resolve/samConstructors/genericSam.fir.txt

 FILE: main.kt
     public final fun main(): R|kotlin/Unit| {
-        R|/MyFunction|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|kotlin/String| {
+        R|/MyFunction|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Int.toInt|().R|kotlin/Any.toString|()
         }
         )
-        R|/MyFunction|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|kotlin/Int|): R|kotlin/String| {
+        R|/MyFunction|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|kotlin/Int|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Any.toString|()
         }
         )

compiler/fir/analysis-tests/testData/resolve/samConstructors/genericSamInferenceFromExpectType.fir.txt

@@ -6,23 +6,23 @@ FILE: main.kt
     public final fun <X, Y> foo3(f: R|MyFunction<X, Y>|, x: R|X|): R|kotlin/Unit| {
     }
     public final fun main(): R|kotlin/Unit| {
-        R|/foo1|(R|/MyFunction|<R|kotlin/Int|, R|kotlin/String|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|kotlin/String| {
+        R|/foo1|(R|/MyFunction|<R|kotlin/Int|, R|kotlin/String|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Int.toInt|().R|kotlin/Any.toString|()
         }
         ))
-        R|/foo2|(R|/MyFunction|<R|kotlin/Number|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Number, kotlin/Number?>!|): R|kotlin/String| {
+        R|/foo2|(R|/MyFunction|<R|kotlin/Number|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Number, kotlin/Number?>!|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Number.toInt|().R|kotlin/Any.toString|()
         }
         ))
-        <Inapplicable(INAPPLICABLE): /foo2>#(R|/MyFunction|<R|kotlin/Number|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|kotlin/Int|): R|kotlin/String| {
+        <Inapplicable(INAPPLICABLE): /foo2>#(R|/MyFunction|<R|kotlin/Number|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|kotlin/Int|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Any.toString|()
         }
         ))
-        R|/foo3|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(R|/MyFunction|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|kotlin/String| {
+        R|/foo3|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(R|/MyFunction|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Int.plus|(Int(1)).R|kotlin/Any.toString|()
         }
         ), Int(1))
-        R|/foo3|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(R|/MyFunction|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|kotlin/Number|): R|kotlin/String| {
+        R|/foo3|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(R|/MyFunction|<R|kotlin/Int|, R|ft<kotlin/String, kotlin/String?>!|>(<L> = MyFunction@fun <anonymous>(x: R|kotlin/Number|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Number.toInt|().R|kotlin/Any.toString|()
         }
         ), Int(2))

compiler/fir/analysis-tests/testData/resolve/samConversions/genericSam.fir.txt

@@ -12,11 +12,11 @@ FILE: main.kt
             ^ R|<local>/x|.R|kotlin/Any.toString|()
         }
         )
-        Q|JavaUsage|.R|/JavaUsage.foo3|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(foo3@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|kotlin/String| {
+        Q|JavaUsage|.R|/JavaUsage.foo3|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(foo3@fun <anonymous>(x: R|ft<kotlin/Int, kotlin/Int?>!|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Int.plus|(Int(1)).R|kotlin/Any.toString|()
         }
         , Int(1))
-        Q|JavaUsage|.R|/JavaUsage.foo3|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(foo3@fun <anonymous>(x: R|kotlin/Number|): R|kotlin/String| {
+        Q|JavaUsage|.R|/JavaUsage.foo3|<R|ft<kotlin/Int, kotlin/Int?>!|, R|ft<kotlin/String, kotlin/String?>!|>(foo3@fun <anonymous>(x: R|kotlin/Number|): R|ft<kotlin/String, kotlin/String?>!| {
             ^ R|<local>/x|.R|kotlin/Number.toInt|().R|kotlin/Any.toString|()
         }
         , Int(2))

compiler/fir/analysis-tests/testData/resolve/samConversions/kotlinSam.fir.txt

@@ -38,12 +38,12 @@ FILE: kotlinSam.kt
         }
         )
         R|/foo1|(R|<local>/f|)
-        <Inapplicable(INAPPLICABLE): /foo2>#(<L> = foo2@fun <anonymous>(x: R|kotlin/Nothing|): R|kotlin/Boolean| {
+        <Inapplicable(INAPPLICABLE): /foo2>#(<L> = foo2@fun <anonymous>(x: R|kotlin/Nothing|): R|kotlin/Nothing| {
             ^ CMP(>, R|<local>/x|.<Unresolved name: compareTo>#(Int(1)))
         }
         )
         <Inapplicable(INAPPLICABLE): /foo2>#(R|<local>/f|)
-        <Inapplicable(INAPPLICABLE): /foo3>#(<L> = foo3@fun <anonymous>(x: R|kotlin/Nothing|): R|kotlin/Boolean| {
+        <Inapplicable(INAPPLICABLE): /foo3>#(<L> = foo3@fun <anonymous>(x: R|kotlin/Nothing|): R|kotlin/Nothing| {
             ^ CMP(>, R|<local>/x|.<Unresolved name: compareTo>#(Int(1)))
         }
         )

compiler/fir/analysis-tests/testData/resolve/samConversions/notSamBecauseOfSupertype.fir.txt

@@ -2,7 +2,7 @@ FILE: main.kt
     public final fun foo(m: R|MyRunnable|): R|kotlin/Unit| {
     }
     public final fun main(): R|kotlin/Unit| {
-        Q|JavaUsage|.<Inapplicable(INAPPLICABLE): /JavaUsage.foo>#(<L> = foo@fun <anonymous>(x: R|kotlin/Nothing|): R|kotlin/Boolean| {
+        Q|JavaUsage|.<Inapplicable(INAPPLICABLE): /JavaUsage.foo>#(<L> = foo@fun <anonymous>(x: R|kotlin/Nothing|): R|kotlin/Nothing| {
             ^ CMP(>, R|<local>/x|.<Unresolved name: compareTo>#(Int(1)))
         }
         )

compiler/fir/analysis-tests/testData/resolveWithStdlib/callableReferences/implicitTypes.fir.txt

@@ -11,6 +11,6 @@ FILE: implicitTypes.kt
     public final fun loop1(): R|(kotlin/Any?) -> kotlin/Nothing| {
         ^loop1 <Inapplicable(INAPPLICABLE): /use>#<R|kotlin/Any?|, R|kotlin/Nothing|>(::<Unresolved reference: loop2>#)
     }
-    public final fun loop2(): R|ERROR CLASS: cycle| {
+    public final fun loop2(): <ERROR TYPE REF: cycle> {
         ^loop2 R|/loop1|()
     }

compiler/fir/analysis-tests/testData/resolveWithStdlib/j+k/MyMap.fir.txt

 FILE: test.kt
     public final fun test(map: R|MyMap|): R|kotlin/Unit| {
-        lval result: R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!| = R|<local>/map|.R|kotlin/collections/getOrPut|<R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!|, R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!|>(String(key), <L> = getOrPut@fun <anonymous>(): R|kotlin/String| <kind=UNKNOWN>  {
+        lval result: R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!| = R|<local>/map|.R|kotlin/collections/getOrPut|<R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!|, R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!|>(String(key), <L> = getOrPut@fun <anonymous>(): R|ft<@FlexibleNullability kotlin/String, kotlin/String?>!| <kind=UNKNOWN>  {
             ^ String(value)
         }
         )

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

@@ -143,6 +143,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         return graphBuilder.returnExpressionsOfAnonymousFunction(function)
     }
 
+    fun isThereControlFlowInfoForAnonymousFunction(function: FirAnonymousFunction): Boolean =
+        graphBuilder.isThereControlFlowInfoForAnonymousFunction(function)
+
     fun dropSubgraphFromCall(call: FirFunctionCall) {
         graphBuilder.dropSubgraphFromCall(call)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -120,6 +120,11 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Public API -----------------------------------
 
+    fun isThereControlFlowInfoForAnonymousFunction(function: FirAnonymousFunction): Boolean =
+        function.controlFlowGraphReference?.controlFlowGraph != null ||
+                exitsOfAnonymousFunctions.containsKey(function.symbol)
+
+    // This function might throw exception if !isThereControlFlowInfoForAnonymousFunction(function)
     fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirStatement> {
         fun FirElement.extractArgument(): FirElement = when {
             this is FirReturnExpression && target.labeledElement.symbol == function.symbol -> result.extractArgument()

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

@@ -19,10 +19,7 @@ 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.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.inference.*
 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
@@ -340,14 +337,22 @@ class FirCallCompletionResultsWriterTransformer(
             .let { finalSubstitutor.substituteOrSelf(it) }
 
     private fun Candidate.createArgumentsMapping(): ExpectedArgumentType? {
-        return argumentMapping?.map { (argument, valueParameter) ->
+        val lambdasReturnType = postponedAtoms.filterIsInstance<ResolvedLambdaAtom>().associate {
+            Pair(it.atom, finalSubstitutor.substituteOrSelf(substitutor.substituteOrSelf(it.returnType)))
+        }
+
+        val arguments = argumentMapping?.map { (argument, valueParameter) ->
             val expectedType = if (valueParameter.isVararg) {
                 valueParameter.returnTypeRef.substitute(this).varargElementType()
             } else {
                 valueParameter.returnTypeRef.substitute(this)
             }
             argument.unwrapArgument() to expectedType
-        }?.toMap()?.toExpectedType()
+        }?.toMap()
+
+
+        if (lambdasReturnType.isEmpty() && arguments.isNullOrEmpty()) return null
+        return ExpectedArgumentType.ArgumentsMap(arguments ?: emptyMap(), lambdasReturnType)
     }
 
     override fun transformDelegatedConstructorCall(
@@ -425,6 +430,13 @@ class FirCallCompletionResultsWriterTransformer(
         anonymousFunction: FirAnonymousFunction,
         data: ExpectedArgumentType?,
     ): CompositeTransformResult<FirStatement> {
+        // This case is not common, and happens when there are anonymous function arguments that aren't mapped to any parameter in the call
+        // So, we don't run body resolve transformation for them, thus there's no control flow info either
+        // Control flow info is necessary prerequisite because we collect return expressions in that function
+        //
+        // Example: second lambda in the call like list.filter({}, {})
+        if (!dataFlowAnalyzer.isThereControlFlowInfoForAnonymousFunction(anonymousFunction)) return anonymousFunction.compose()
+
         val expectedType = data?.getExpectedType(anonymousFunction)?.let { expectedArgumentType ->
             // From the argument mapping, the expected type of this anonymous function would be:
             when {
@@ -455,10 +467,13 @@ class FirCallCompletionResultsWriterTransformer(
             needUpdateLambdaType = true
         }
 
-        val expectedReturnType = expectedType?.returnType(session) as? ConeClassLikeType
         val initialType = anonymousFunction.returnTypeRef.coneTypeSafe<ConeKotlinType>()
-        if (initialType != null) {
-            val finalType = expectedReturnType ?: finalSubstitutor.substituteOrNull(initialType)
+        val finalType =
+            expectedType?.returnType(session) as? ConeClassLikeType
+                ?: (data as? ExpectedArgumentType.ArgumentsMap)?.lambdasReturnTypes?.get(anonymousFunction)
+                ?: initialType?.let(finalSubstitutor::substituteOrSelf)
+
+        if (finalType != null) {
             val resultType = anonymousFunction.returnTypeRef.withReplacedConeType(finalType)
             anonymousFunction.transformReturnTypeRef(StoreType, resultType)
             needUpdateLambdaType = true
@@ -472,19 +487,24 @@ class FirCallCompletionResultsWriterTransformer(
 
         val result = transformElement(anonymousFunction, null)
 
+        val returnExpressionsOfAnonymousFunction: Collection<FirStatement> =
+            dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(anonymousFunction)
+        for (expression in returnExpressionsOfAnonymousFunction) {
+            expression.transform<FirElement, ExpectedArgumentType?>(this, finalType?.toExpectedType())
+        }
+
         val resultFunction = result.single
         if (resultFunction.returnTypeRef.coneTypeSafe<ConeIntegerLiteralType>() != null) {
-            val blockType = resultFunction.body?.typeRef?.coneTypeSafe<ConeKotlinType>()
-            resultFunction.replaceReturnTypeRef(resultFunction.returnTypeRef.withReplacedConeType(blockType))
+            val lastExpressionType =
+                (returnExpressionsOfAnonymousFunction.lastOrNull() as? FirExpression)
+                    ?.typeRef?.coneTypeSafe<ConeKotlinType>()
+
+            resultFunction.replaceReturnTypeRef(resultFunction.returnTypeRef.withReplacedConeType(lastExpressionType))
             resultFunction.replaceTypeRef(
                 resultFunction.constructFunctionalTypeRef(isSuspend = expectedType?.isSuspendFunctionType(session) == true)
             )
         }
 
-        for (expression in dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(anonymousFunction)) {
-            expression.transform<FirElement, ExpectedArgumentType?>(this, null)
-        }
-
         return result
     }
 
@@ -622,7 +642,11 @@ class FirCallCompletionResultsWriterTransformer(
 }
 
 sealed class ExpectedArgumentType {
-    class ArgumentsMap(val map: Map<FirExpression, ConeKotlinType>) : ExpectedArgumentType()
+    class ArgumentsMap(
+        val map: Map<FirExpression, ConeKotlinType>,
+        val lambdasReturnTypes: Map<FirAnonymousFunction, ConeKotlinType>
+    ) : ExpectedArgumentType()
+
     class ExpectedType(val type: ConeKotlinType) : ExpectedArgumentType()
     object NoApproximation : ExpectedArgumentType()
 }
@@ -633,7 +657,6 @@ private fun ExpectedArgumentType.getExpectedType(argument: FirExpression): ConeK
     ExpectedArgumentType.NoApproximation -> null
 }
 
-private fun Map<FirExpression, ConeKotlinType>.toExpectedType(): ExpectedArgumentType = ExpectedArgumentType.ArgumentsMap(this)
 fun ConeKotlinType.toExpectedType(): ExpectedArgumentType = ExpectedArgumentType.ExpectedType(this)
 
 private fun FirExpression.unwrapArgument(): FirExpression = when (this) {

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

@@ -9,7 +9,9 @@ import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
-import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirNamedArgumentExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildVarargArgumentsExpression
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -83,7 +85,6 @@ internal fun remapArgumentsWithVararg(
 
 fun FirBlock.writeResultType(session: FirSession) {
     val resultExpression = when (val statement = statements.lastOrNull()) {
-        is FirReturnExpression -> statement.result
         is FirExpression -> statement
         else -> null
     }

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

@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
@@ -442,7 +443,23 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
             val body = result.body
             if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
-                result.transformReturnTypeRef(transformer, withExpectedType(body.resultType))
+                // TODO: This part seems unnecessary because for lambdas in dependent context will be completed and their type
+                //  should be replaced there properly
+                val returnType =
+                    dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(result)
+                        .firstNotNullResult { (it as? FirExpression)?.resultType?.coneTypeSafe() }
+
+                if (returnType != null) {
+                    result.transformReturnTypeRef(transformer, withExpectedType(returnType))
+                } else {
+                    result.transformReturnTypeRef(
+                        transformer,
+                        withExpectedType(buildErrorTypeRef {
+                            diagnostic =
+                                ConeSimpleDiagnostic("Unresolved lambda return type", DiagnosticKind.InferenceError)
+                        })
+                    )
+                }
             }
             return result
         }
@@ -506,11 +523,12 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         val body = result.body
         if (result.returnTypeRef is FirImplicitTypeRef) {
             val simpleFunction = function as? FirSimpleFunction
-            if (body != null) {
+            val returnExpression = (body?.statements?.single() as? FirReturnExpression)?.result
+            if (returnExpression != null && returnExpression.typeRef is FirResolvedTypeRef) {
                 result.transformReturnTypeRef(
                     transformer,
                     withExpectedType(
-                        body.resultType.approximatedIfNeededOrSelf(
+                        returnExpression.resultType.approximatedIfNeededOrSelf(
                             inferenceComponents.approximator, simpleFunction?.visibility, simpleFunction?.isInline == true
                         )
                     )

compiler/testData/ir/irText/expressions/sam/samConversionToGeneric.fir.kt.txt

@@ -6,14 +6,14 @@ fun test1(): J<String?> {
 }
 
 fun test2(): J<String?> {
-  return local fun <anonymous>(x: String): String {
+  return local fun <anonymous>(x: String): String? {
     return x
   }
  /*-> J<String?> */
 }
 
 fun test3() {
-  return bar<String?>(j = local fun <anonymous>(x: String): String {
+  return bar<String?>(j = local fun <anonymous>(x: String): String? {
     return x
   }
  /*-> J<X?>? */)

compiler/testData/ir/irText/expressions/sam/samConversionToGeneric.fir.txt

@@ -13,11 +13,11 @@ FILE fqName:<root> fileName:/samConversionToGeneric.kt
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun test2 (): <root>.J<kotlin.String?> declared in <root>'
         TYPE_OP type=<root>.J<kotlin.String?> origin=SAM_CONVERSION typeOperand=<root>.J<kotlin.String?>
-          FUN_EXPR type=kotlin.Function1<kotlin.String, kotlin.String> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (x:kotlin.String) returnType:kotlin.String
+          FUN_EXPR type=kotlin.Function1<kotlin.String, kotlin.String?> origin=LAMBDA
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (x:kotlin.String) returnType:kotlin.String?
               VALUE_PARAMETER name:x index:0 type:kotlin.String
               BLOCK_BODY
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (x: kotlin.String): kotlin.String declared in <root>.test2'
+                RETURN type=kotlin.Nothing from='local final fun <anonymous> (x: kotlin.String): kotlin.String? declared in <root>.test2'
                   GET_VAR 'x: kotlin.String declared in <root>.test2.<anonymous>' type=kotlin.String origin=null
   FUN name:test3 visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
@@ -25,11 +25,11 @@ FILE fqName:<root> fileName:/samConversionToGeneric.kt
         CALL 'public open fun bar <X> (j: <root>.J<X of <root>.H.bar?>?): kotlin.Unit declared in <root>.H' type=kotlin.Unit origin=null
           <X>: kotlin.String?
           j: TYPE_OP type=<root>.J<X of <root>.H.bar?>? origin=SAM_CONVERSION typeOperand=<root>.J<X of <root>.H.bar?>?
-            FUN_EXPR type=kotlin.Function1<kotlin.String, kotlin.String> origin=LAMBDA
-              FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (x:kotlin.String) returnType:kotlin.String
+            FUN_EXPR type=kotlin.Function1<kotlin.String, kotlin.String?> origin=LAMBDA
+              FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (x:kotlin.String) returnType:kotlin.String?
                 VALUE_PARAMETER name:x index:0 type:kotlin.String
                 BLOCK_BODY
-                  RETURN type=kotlin.Nothing from='local final fun <anonymous> (x: kotlin.String): kotlin.String declared in <root>.test3'
+                  RETURN type=kotlin.Nothing from='local final fun <anonymous> (x: kotlin.String): kotlin.String? declared in <root>.test3'
                     GET_VAR 'x: kotlin.String declared in <root>.test3.<anonymous>' type=kotlin.String origin=null
   FUN name:test4 visibility:public modality:FINAL <> (a:kotlin.Any) returnType:kotlin.Unit
     VALUE_PARAMETER name:a index:0 type:kotlin.Any

compiler/testData/ir/irText/types/nullChecks/nullCheckOnLambdaResult/stringVsTXArray.fir.kt.txt

@@ -3,8 +3,9 @@ fun <T : Any?> useTX(x: T, fn: Function0<T>): T {
 }
 
 fun testNoNullCheck(xs: Array<String>) {
-  useTX<Serializable?>(x = xs, fn = local fun <anonymous>(): String? {
+  useTX<Serializable?>(x = xs, fn = local fun <anonymous>(): Serializable? {
     return string()
   }
 ) /*~> Unit */
 }
+

compiler/testData/ir/irText/types/nullChecks/nullCheckOnLambdaResult/stringVsTXArray.fir.txt

@@ -14,8 +14,8 @@ FILE fqName:<root> fileName:/stringVsTXArray.kt
         CALL 'public final fun useTX <T> (x: T of <root>.useTX, fn: kotlin.Function0<T of <root>.useTX>): T of <root>.useTX declared in <root>' type=java.io.Serializable? origin=null
           <T>: java.io.Serializable?
           x: GET_VAR 'xs: kotlin.Array<kotlin.String> declared in <root>.testNoNullCheck' type=kotlin.Array<kotlin.String> origin=null
-          fn: FUN_EXPR type=kotlin.Function0<kotlin.String?> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.String?
+          fn: FUN_EXPR type=kotlin.Function0<java.io.Serializable?> origin=LAMBDA
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:java.io.Serializable?
               BLOCK_BODY
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.String? declared in <root>.testNoNullCheck'
+                RETURN type=kotlin.Nothing from='local final fun <anonymous> (): java.io.Serializable? declared in <root>.testNoNullCheck'
                   CALL 'public open fun string (): kotlin.String? declared in <root>.J' type=kotlin.String? origin=null