[FIR] Consider explicit returns during computing return type of anonymous function

Introduce FirAnonymousFunctionReturnExpressionInfo

^KT-59386

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

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedErrorReference
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.dfa.PropertyStability
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.FirAnonymousFunctionReturnExpressionInfo
 import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType

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

@@ -136,10 +136,10 @@ abstract class FirDataFlowAnalyzer(
         return variable.stability to types.toMutableList()
     }
 
-    fun returnExpressionsOfAnonymousFunctionOrNull(function: FirAnonymousFunction): Collection<FirExpression>? =
+    fun returnExpressionsOfAnonymousFunctionOrNull(function: FirAnonymousFunction): Collection<FirAnonymousFunctionReturnExpressionInfo>? =
         graphBuilder.returnExpressionsOfAnonymousFunction(function)
 
-    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirExpression> =
+    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirAnonymousFunctionReturnExpressionInfo> =
         returnExpressionsOfAnonymousFunctionOrNull(function)
             ?: error("anonymous function ${function.render()} not analyzed")
 

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

@@ -25,6 +25,8 @@ import org.jetbrains.kotlin.fir.util.listMultimapOf
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.utils.getOrPutNullable
 
+data class FirAnonymousFunctionReturnExpressionInfo(val expression: FirExpression, val isExplicit: Boolean)
+
 @OptIn(CfgInternals::class)
 class ControlFlowGraphBuilder {
     private val graphs: Stack<ControlFlowGraph> = stackOf()
@@ -82,10 +84,10 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Public API -----------------------------------
 
-    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirExpression>? {
+    fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirAnonymousFunctionReturnExpressionInfo>? {
         val exitNode = function.controlFlowGraphReference?.controlFlowGraph?.exitNode ?: return null
 
-        fun CFGNode<*>.returnExpression(): FirExpression? = when (this) {
+        fun CFGNode<*>.returnExpression(): FirAnonymousFunctionReturnExpressionInfo? = when (this) {
             is BlockExitNode -> when {
                 // lambda@{ x } -> x
                 // lambda@{ class C } -> Unit-returning stub
@@ -106,16 +108,21 @@ class ControlFlowGraphBuilder {
                                 lastStatement.source?.kind != KtFakeSourceElementKind.ImplicitReturn.FromLastStatement ->
                             null
                         else ->
-                            lastStatement as? FirExpression
-                                ?: buildUnitExpression { source = fir.statements.lastOrNull()?.source ?: fir.source }
+                            (lastStatement as? FirExpression
+                                ?: buildUnitExpression { source = fir.statements.lastOrNull()?.source ?: fir.source }).let {
+                                FirAnonymousFunctionReturnExpressionInfo(it, isExplicit = false)
+                            }
                     }
                 }
                 // fun() { terminatingExpression } -> nothing (checker will emit an error if return type is not Unit)
                 // fun() { throw } or fun() { returnsNothing() } -> Nothing-returning stub
                 else -> FirStub.takeIf { _ -> previousNodes.all { it is StubNode } }
+                    ?.let { FirAnonymousFunctionReturnExpressionInfo(it, isExplicit = false) }
             }
             // lambda@{ return@lambda x } -> x
-            is JumpNode -> (fir as? FirReturnExpression)?.takeIf { it.target.labeledElement.symbol == function.symbol }?.result
+            is JumpNode -> (fir as? FirReturnExpression)?.takeIf { it.target.labeledElement.symbol == function.symbol }?.result?.let {
+                FirAnonymousFunctionReturnExpressionInfo(it, isExplicit = true)
+            }
             else -> null // shouldn't happen? expression bodies are implicitly wrapped in `FirBlock`s
         }
 

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

@@ -33,7 +33,6 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.visitors.transformSingle
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.addEqualityConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
@@ -341,7 +340,7 @@ class FirCallCompleter(
             }
             transformer.context.dropContextForAnonymousFunction(lambdaArgument)
 
-            val returnArguments = components.dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambdaArgument)
+            val returnArguments = components.dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambdaArgument).map { it.expression }
 
             return ReturnArgumentsAnalysisResult(returnArguments, builderInferenceSession)
         }

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

@@ -552,7 +552,7 @@ class FirCallCompletionResultsWriterTransformer(
         // The case where we can't find any return expressions 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. Example: second lambda in the call like list.filter({}, {})
-        val returnExpressions = dataFlowAnalyzer.returnExpressionsOfAnonymousFunctionOrNull(anonymousFunction)
+        val returnExpressions = dataFlowAnalyzer.returnExpressionsOfAnonymousFunctionOrNull(anonymousFunction)?.map { it.expression }
             ?: return transformImplicitTypeRefInAnonymousFunction(anonymousFunction)
 
         val expectedType = data?.getExpectedType(anonymousFunction)?.let { expectedArgumentType ->

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

@@ -865,16 +865,16 @@ open class FirDeclarationsResolveTransformer(
     }
 
     private fun FirAnonymousFunction.computeReturnTypeRef(expected: FirResolvedTypeRef?): FirResolvedTypeRef {
-        // Any lambda expression assigned to `(...) -> Unit` returns Unit
-        if (isLambda && expected?.type?.isUnit == true) return expected
-        // `lambda@ { return@lambda }` always returns Unit
         val returnExpressions = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(this)
-        if (shouldReturnUnit(returnExpressions)) return session.builtinTypes.unitType
+        // Any lambda expression assigned to `(...) -> Unit` returns Unit if all return expressions are implicit
+        // `lambda@ { return@lambda }` always returns Unit
+        if (isLambda && expected?.type?.isUnit == true && returnExpressions.all { !it.isExplicit }) return expected
+        if (shouldReturnUnit(returnExpressions.map { it.expression })) return session.builtinTypes.unitType
         // Here is a questionable moment where we could prefer the expected type over an inferred one.
         // In correct code this doesn't matter, as all return expression types should be subtypes of the expected type.
         // In incorrect code, this would change diagnostics: we can get errors either on the entire lambda, or only on its
         // return statements. The former kind of makes more sense, but the latter is more readable.
-        val inferredFromReturnExpressions = session.typeContext.commonSuperTypeOrNull(returnExpressions.map { it.resultType.coneType })
+        val inferredFromReturnExpressions = session.typeContext.commonSuperTypeOrNull(returnExpressions.map { it.expression.resultType.coneType })
         return inferredFromReturnExpressions?.let { returnTypeRef.resolvedTypeFromPrototype(it) }
             ?: session.builtinTypes.unitType // Empty lambda returns Unit
     }

compiler/testData/diagnostics/tests/functionLiterals/return/IfWithoutElseWithExplicitType.fir.kt

@@ -4,9 +4,9 @@ val a: () -> Int = <!INITIALIZER_TYPE_MISMATCH!>l@ {
     if (flag) return@l 4
 }<!>
 
-val b: () -> Unit = l@ {
+val b: () -> Unit = <!INITIALIZER_TYPE_MISMATCH!>l@ {
     if (flag) return@l 4
-}
+}<!>
 
 val c: () -> Any = l@ {
     if (flag) return@l 4

compiler/testData/diagnostics/tests/inference/coercionToUnit/coercionWithExpectedType.fir.kt

@@ -10,13 +10,13 @@ val a: () -> Unit = l@{
     if (true) 42
 }
 
-val b: () -> Unit = l@{
+val b: () -> Unit = <!INITIALIZER_TYPE_MISMATCH!>l@{
     // Error, coercion can't be applied at this position!
     if (true) return@l "hello"
 
     // However, this is OK, because here coercion is applied
     "hello"
-}
+}<!>
 
 val c: () -> Unit = {
     // Interesting enough, for such expessions we use expected type Unit