FIR: coerce to Unit when a lambda has early returns

^KT-39075 Fixed

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

@@ -11,7 +11,7 @@ FILE: coercionToUnitWithEarlyReturn.kt
     public final fun foo(x: R|() -> kotlin/Unit|): R|kotlin/Unit| {
     }
     public final fun main(x: R|A?|): R|kotlin/Unit| {
-        lval lambda: R|() -> kotlin/Unit?| = l@fun <anonymous>(): R|kotlin/Unit?| {
+        lval lambda: R|() -> kotlin/Unit| = l@fun <anonymous>(): R|kotlin/Unit| {
             when () {
                 ==(R|<local>/x|?.{ $subj$.R|kotlin/Any.hashCode|() }, Int(0)) ->  {
                     ^@l Unit
@@ -21,5 +21,5 @@ FILE: coercionToUnitWithEarlyReturn.kt
             ^ R|<local>/x|?.{ $subj$.R|/A.unit|() }
         }
 
-        <Inapplicable(INAPPLICABLE): /foo>#(R|<local>/lambda|)
+        R|/foo|(R|<local>/lambda|)
     }

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

@@ -14,6 +14,5 @@ fun main(x: A?) {
         x?.unit()
     }
 
-    // lambda has a type (() -> Unit?)
+    foo(lambda)
-    <!INAPPLICABLE_CANDIDATE{LT}!><!INAPPLICABLE_CANDIDATE{PSI}!>foo<!>(lambda)<!>
 }

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

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
@@ -780,13 +781,28 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     dataFlowAnalyzer,
                 )
                 lambda.transformSingle(writer, expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType())
-                val returnTypes = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda)
+
-                    .mapNotNull { (it as? FirExpression)?.resultType?.coneType }
+                val returnStatements = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda)
-                lambda.replaceReturnTypeRef(
+                val returnExpressionsExceptLast =
-                    lambda.returnTypeRef.resolvedTypeFromPrototype(
+                    if (returnStatements.size > 1)
-                        inferenceComponents.ctx.commonSuperTypeOrNull(returnTypes) ?: session.builtinTypes.unitType.type
+                        returnStatements - lambda.body?.statements?.lastOrNull()
-                    )
+                    else
-                )
+                        returnStatements
+                val implicitReturns = returnExpressionsExceptLast.filter {
+                    (it as? FirExpression)?.typeRef is FirImplicitUnitTypeRef
+                }
+                val returnType =
+                    if (implicitReturns.isNotEmpty()) {
+                        // i.e., early return, e.g., l@{ ... return@l ... }
+                        // Note that the last statement will be coerced to Unit if needed.
+                        session.builtinTypes.unitType.type
+                    } else {
+                        // Otherwise, compute the common super type of all possible return expressions
+                        inferenceComponents.ctx.commonSuperTypeOrNull(
+                            returnStatements.mapNotNull { (it as? FirExpression)?.resultType?.coneType }
+                        ) ?: session.builtinTypes.unitType.type
+                    }
+                lambda.replaceReturnTypeRef(lambda.returnTypeRef.resolvedTypeFromPrototype(returnType))
                 lambda.replaceTypeRef(
                     lambda.constructFunctionalTypeRef(
                         isSuspend = expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true