FIR: Replicate coercion-to-unit behavior from FE 1.0

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

@@ -2411,6 +2411,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/definitelyNotNullIntersectionType.kt");
         }
 
+        @TestMetadata("elvisAndUnit.kt")
+        public void testElvisAndUnit() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/inference/elvisAndUnit.kt");
+        }
+
         @TestMetadata("extensionCallableReferences.kt")
         public void testExtensionCallableReferences() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/inference/extensionCallableReferences.kt");

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

@@ -0,0 +1,18 @@
+FILE: elvisAndUnit.kt
+    public final fun foo(x: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit| {
+    }
+    public final class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun bar(): R|kotlin/Unit| {
+        }
+
+    }
+    public final fun main(a: R|A?|, y: R|kotlin/String|): R|kotlin/Unit| {
+        R|/foo|(<L> = foo@fun <anonymous>(it: R|kotlin/Int|): R|kotlin/Unit| <inline=NoInline>  {
+            R|<local>/a|?.{ $subj$.R|/A.bar|() } ?: R|<local>/y|.R|kotlin/String.get|(Int(0))
+        }
+        )
+    }

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

@@ -0,0 +1,11 @@
+fun foo(x: (Int) -> Unit) {}
+
+class A {
+    fun bar() {}
+}
+
+fun main(a: A?, y: String) {
+    foo {
+        a?.bar() ?: y.get(0)
+    }
+}

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

@@ -2739,6 +2739,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/inference/definitelyNotNullIntersectionType.kt");
             }
 
+            @Test
+            @TestMetadata("elvisAndUnit.kt")
+            public void testElvisAndUnit() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/inference/elvisAndUnit.kt");
+            }
+
             @Test
             @TestMetadata("extensionCallableReferences.kt")
             public void testExtensionCallableReferences() throws Exception {

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

@@ -2739,6 +2739,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolve/inference/definitelyNotNullIntersectionType.kt");
             }
 
+            @Test
+            @TestMetadata("elvisAndUnit.kt")
+            public void testElvisAndUnit() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/inference/elvisAndUnit.kt");
+            }
+
             @Test
             @TestMetadata("extensionCallableReferences.kt")
             public void testExtensionCallableReferences() throws Exception {

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

@@ -16,26 +16,32 @@ sealed class ResolutionMode {
     object ContextDependentDelegate : ResolutionMode()
     object ContextIndependent : ResolutionMode()
     // TODO: it's better not to use WithExpectedType(FirImplicitTypeRef)
-    class WithExpectedType(val expectedTypeRef: FirTypeRef) : ResolutionMode()
+    class WithExpectedType(val expectedTypeRef: FirTypeRef, val mayBeCoercionToUnitApplied: Boolean = false) : ResolutionMode()
 
     class WithStatus(val status: FirDeclarationStatus) : ResolutionMode()
 
     class LambdaResolution(val expectedReturnTypeRef: FirResolvedTypeRef?) : ResolutionMode()
 }
 
+fun ResolutionMode.expectedType(components: BodyResolveComponents): FirTypeRef? = when (this) {
+    is ResolutionMode.WithExpectedType -> expectedTypeRef
+    is ResolutionMode.ContextIndependent -> components.noExpectedType
+    else -> null
+}
+
 fun withExpectedType(expectedTypeRef: FirTypeRef?): ResolutionMode =
     expectedTypeRef?.let { ResolutionMode.WithExpectedType(it) } ?: ResolutionMode.ContextDependent
 
 @JvmName("withExpectedTypeNullable")
-fun withExpectedType(coneType: ConeKotlinType?): ResolutionMode {
+fun withExpectedType(coneType: ConeKotlinType?, mayBeCoercionToUnitApplied: Boolean = false): ResolutionMode {
-    return coneType?.let { withExpectedType(it) } ?: ResolutionMode.ContextDependent
+    return coneType?.let { withExpectedType(it, mayBeCoercionToUnitApplied) } ?: ResolutionMode.ContextDependent
 }
 
-fun withExpectedType(coneType: ConeKotlinType): ResolutionMode {
+fun withExpectedType(coneType: ConeKotlinType, mayBeCoercionToUnitApplied: Boolean = false): ResolutionMode {
     val typeRef = buildResolvedTypeRef {
         type = coneType
     }
-    return ResolutionMode.WithExpectedType(typeRef)
+    return ResolutionMode.WithExpectedType(typeRef, mayBeCoercionToUnitApplied)
 }
 
 fun FirDeclarationStatus.mode(): ResolutionMode =

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

@@ -17,7 +17,10 @@ import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
+import org.jetbrains.kotlin.fir.resolve.calls.isUnitOrFlexibleUnit
+import org.jetbrains.kotlin.fir.resolve.expectedType
 import org.jetbrains.kotlin.fir.resolve.inference.model.ConeArgumentConstraintPosition
+import org.jetbrains.kotlin.fir.resolve.inference.model.ConeExpectedTypeConstraintPosition
 import org.jetbrains.kotlin.fir.resolve.initialTypeOfCandidate
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
@@ -33,9 +36,9 @@ 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.addSubtypeConstraintIfCompatible
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.ConstraintSystemCompletionMode
-import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.model.StubTypeMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
@@ -53,7 +56,13 @@ class FirCallCompleter(
 
     data class CompletionResult<T>(val result: T, val callCompleted: Boolean)
 
-    fun <T> completeCall(call: T, expectedTypeRef: FirTypeRef?): CompletionResult<T>
+    fun <T> completeCall(call: T, expectedTypeRef: FirTypeRef?): CompletionResult<T> where T : FirResolvable, T : FirStatement =
+        completeCall(call, expectedTypeRef, mayBeCoercionToUnitApplied = false)
+
+    fun <T> completeCall(call: T, data: ResolutionMode): CompletionResult<T> where T : FirResolvable, T : FirStatement =
+        completeCall(call, data.expectedType(components), (data as? ResolutionMode.WithExpectedType)?.mayBeCoercionToUnitApplied == true)
+
+    fun <T> completeCall(call: T, expectedTypeRef: FirTypeRef?, mayBeCoercionToUnitApplied: Boolean): CompletionResult<T>
             where T : FirResolvable, T : FirStatement {
         val typeRef = components.typeFromCallee(call)
 
@@ -73,7 +82,15 @@ class FirCallCompleter(
         }
 
         if (expectedTypeRef is FirResolvedTypeRef) {
-            candidate.system.addSubtypeConstraint(initialType, expectedTypeRef.type, SimpleConstraintSystemConstraintPosition)
+            if (expectedTypeRef.coneType.isUnitOrFlexibleUnit && mayBeCoercionToUnitApplied) {
+                if (candidate.system.notFixedTypeVariables.isNotEmpty()) {
+                    candidate.system.addSubtypeConstraintIfCompatible(
+                        initialType, expectedTypeRef.type, ConeExpectedTypeConstraintPosition()
+                    )
+                }
+            } else {
+                candidate.system.addSubtypeConstraint(initialType, expectedTypeRef.type, ConeExpectedTypeConstraintPosition())
+            }
         }
 
         val completionMode = candidate.computeCompletionMode(session.inferenceComponents, expectedTypeRef, initialType)
@@ -149,7 +166,7 @@ class FirCallCompleter(
             isNullable = functionalType.isNullable,
             functionalType.attributes
         )
-        csBuilder.addSubtypeConstraint(expectedType, functionalType, ConeArgumentConstraintPosition())
+        csBuilder.addSubtypeConstraint(expectedType, functionalType, ConeArgumentConstraintPosition(atom.atom))
         atom.replaceExpectedType(expectedType, returnVariable.defaultType)
         atom.replaceTypeVariableForLambdaReturnType(returnVariable)
     }

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

@@ -84,11 +84,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
     protected inline val file: FirFile get() = components.file
 
     val ResolutionMode.expectedType: FirTypeRef?
-        get() = when (this) {
+        get() = expectedType(components)
-            is ResolutionMode.WithExpectedType -> expectedTypeRef
-            is ResolutionMode.ContextIndependent -> noExpectedType
-            else -> null
-        }
 
     class BodyResolveTransformerComponents(
         override val session: FirSession,

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.calls.isUnitOrFlexibleUnit
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.fir.resolve.withExpectedType
@@ -83,8 +84,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
                         return@with whenExpression
                     }
 
-                    val expectedTypeRef = data.expectedType
+                    val completionResult = callCompleter.completeCall(whenExpression, data)
-                    val completionResult = callCompleter.completeCall(whenExpression, expectedTypeRef)
                     whenExpression = completionResult.result
                 }
             }
@@ -147,8 +147,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
 
         @Suppress("NAME_SHADOWING")
         var result = syntheticCallGenerator.generateCalleeForTryExpression(tryExpression, resolutionContext).let {
-            val expectedTypeRef = data.expectedType
+            val completionResult = callCompleter.completeCall(it, data)
-            val completionResult = callCompleter.completeCall(it, expectedTypeRef)
             callCompleted = completionResult.callCompleted
             completionResult.result
         }
@@ -203,7 +202,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         throwExpression: FirThrowExpression,
         data: ResolutionMode
     ): FirStatement {
-        return transformer.transformExpression(throwExpression, data).also {
+        return transformer.transformExpression(throwExpression, ResolutionMode.ContextIndependent).also {
             dataFlowAnalyzer.exitThrowExceptionNode(it as FirThrowExpression)
         }
     }
@@ -218,15 +217,24 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         elvisExpression.transformAnnotations(transformer, data)
 
         val expectedType = data.expectedType?.coneTypeSafe<ConeKotlinType>()
-        val resolutionModeForLhs = withExpectedType(expectedType?.withNullability(ConeNullability.NULLABLE, session.typeContext))
+        val mayBeCoercionToUnitApplied = (data as? ResolutionMode.WithExpectedType)?.mayBeCoercionToUnitApplied == true
+
+        val resolutionModeForLhs =
+            if (mayBeCoercionToUnitApplied && expectedType?.isUnitOrFlexibleUnit == true)
+                withExpectedType(expectedType, mayBeCoercionToUnitApplied = true)
+            else
+                withExpectedType(expectedType?.withNullability(ConeNullability.NULLABLE, session.typeContext))
         elvisExpression.transformLhs(transformer, resolutionModeForLhs)
         dataFlowAnalyzer.exitElvisLhs(elvisExpression)
 
-        val resolutionModeForRhs = withExpectedType(expectedType)
+        val resolutionModeForRhs = withExpectedType(
+            expectedType,
+            mayBeCoercionToUnitApplied = mayBeCoercionToUnitApplied
+        )
         elvisExpression.transformRhs(transformer, resolutionModeForRhs)
 
         val result = syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression, resolutionContext)?.let {
-            callCompleter.completeCall(it, data.expectedType).result
+            callCompleter.completeCall(it, data).result
         } ?: elvisExpression.also {
             it.resultType = buildErrorTypeRef {
                 diagnostic = ConeSimpleDiagnostic("Can't resolve ?: operator call", DiagnosticKind.InferenceError)

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

@@ -26,9 +26,9 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.*
 import org.jetbrains.kotlin.fir.resolve.inference.FirStubInferenceSession
 import org.jetbrains.kotlin.fir.resolve.inference.inferenceComponents
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
-import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
+import org.jetbrains.kotlin.fir.resolve.substitution.substituteOrNull
-import org.jetbrains.kotlin.fir.resolve.transformers.StoreReceiver
+import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
-import org.jetbrains.kotlin.fir.resolve.transformers.firClassLike
+import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
@@ -125,7 +125,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                     if (!transformedCallee.isAcceptableResolvedQualifiedAccess()) {
                         return qualifiedAccessExpression
                     }
-                    callCompleter.completeCall(transformedCallee, data.expectedType).result
+                    callCompleter.completeCall(transformedCallee, data).result
                 } else {
                     transformedCallee
                 }
@@ -280,7 +280,6 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         functionCall.transformAnnotations(transformer, data)
         functionCall.transformSingle(InvocationKindTransformer, null)
         functionCall.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
-        val expectedTypeRef = data.expectedType
         val (completeInference, callCompleted) =
             try {
                 val initialExplicitReceiver = functionCall.explicitReceiver
@@ -290,7 +289,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                     // name.invoke() case
                     callCompleter.completeCall(resultExplicitReceiver, noExpectedType)
                 }
-                callCompleter.completeCall(resultExpression, expectedTypeRef)
+                callCompleter.completeCall(resultExpression, data)
             } catch (e: Throwable) {
                 throw RuntimeException("While resolving call ${functionCall.render()}", e)
             }
@@ -318,7 +317,14 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val numberOfStatements = block.statements.size
 
         block.transformStatementsIndexed(transformer) { index ->
-            val value = if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent
+            val value =
+                if (index == numberOfStatements - 1)
+                    if (data is ResolutionMode.WithExpectedType)
+                        ResolutionMode.WithExpectedType(data.expectedTypeRef, mayBeCoercionToUnitApplied = true)
+                    else
+                        data
+                else
+                    ResolutionMode.ContextIndependent
             transformer.firTowerDataContextCollector?.addStatementContext(block.statements[index], context.towerDataContext)
             TransformData.Data(value)
         }
@@ -632,7 +638,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
 
         var callCompleted = false
         val result = components.syntheticCallGenerator.generateCalleeForCheckNotNullCall(checkNotNullCall, resolutionContext)?.let {
-            val completionResult = callCompleter.completeCall(it, data.expectedType)
+            val completionResult = callCompleter.completeCall(it, data)
             callCompleted = completionResult.callCompleted
             completionResult.result
         } ?: run {