diff --git a/compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.fir.txt b/compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.fir.txt index a819ba09632..e04a089fcf3 100644 --- a/compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.fir.txt +++ b/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 (): R|kotlin/Unit?| { + lval lambda: R|() -> kotlin/Unit| = l@fun (): R|kotlin/Unit| { when () { ==(R|/x|?.{ $subj$.R|kotlin/Any.hashCode|() }, Int(0)) -> { ^@l Unit @@ -21,5 +21,5 @@ FILE: coercionToUnitWithEarlyReturn.kt ^ R|/x|?.{ $subj$.R|/A.unit|() } } - #(R|/lambda|) + R|/foo|(R|/lambda|) } diff --git a/compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.kt b/compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.kt index 575ef2a0fe7..17e399120e2 100644 --- a/compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.kt +++ b/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) + foo(lambda) } diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt index 6fc9b4397a9..826950c76fc 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt +++ b/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()?.toExpectedType()) - val returnTypes = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda) - .mapNotNull { (it as? FirExpression)?.resultType?.coneType } - lambda.replaceReturnTypeRef( - lambda.returnTypeRef.resolvedTypeFromPrototype( - inferenceComponents.ctx.commonSuperTypeOrNull(returnTypes) ?: session.builtinTypes.unitType.type - ) - ) + + val returnStatements = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda) + val returnExpressionsExceptLast = + if (returnStatements.size > 1) + 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()?.isSuspendFunctionType(session) == true