[FIR] Add equality constraint from expected type for some synthetic function calls

This fixes some cases where we infer some type variable inside one
of the branches to Nothing instead of the expected type because Nothing
appeared in some other branch.

Specifically, we add an equality instead of a subtype constraint during
completion of calls to synthetic functions for if/when, try and !!.
We don't do it when the call contains a (possibly nested) elvis or is
inside the RHS of an assignment.
Otherwise, we would prevent some smart-casts.

#KT-65882 Fixed

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeCannotInferValueParameterType
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 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
@@ -30,6 +31,9 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBod
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolve.transformers.replaceLambdaArgumentInvocationKinds
 import org.jetbrains.kotlin.fir.resolve.typeFromCallee
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.SyntheticCallableId
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
@@ -133,10 +137,17 @@ class FirCallCompleter(
         resolutionMode: ResolutionMode,
     ) {
         if (resolutionMode !is ResolutionMode.WithExpectedType) return
-        val expectedType = resolutionMode.expectedTypeRef.coneTypeSafe<ConeKotlinType>() ?: return
+        val expectedType = resolutionMode.expectedTypeRef.type.fullyExpandedType(session)
 
         val system = candidate.system
         when {
+            // Only add equality constraint in independent contexts (resolutionMode.forceFullCompletion) for K1 compatibility.
+            // Otherwise,
+            // we miss some constraints from incorporation which leads to NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER in cases like
+            // compiler/testData/diagnostics/tests/inference/nestedIfWithExpectedType.kt.
+            resolutionMode.forceFullCompletion && candidate.isSyntheticFunctionCallThatShouldUseEqualityConstraint(expectedType) ->
+                system.addEqualityConstraintIfCompatible(initialType, expectedType, ConeExpectedTypeConstraintPosition)
+
             // If type mismatch is assumed to be reported in the checker, we should not add a subtyping constraint that leads to error.
             // Because it might make resulting type correct while, it's hopefully would be more clear if we let the call be inferred without
             // the expected type, and then would report diagnostic in the checker.
@@ -167,6 +178,44 @@ class FirCallCompleter(
         }
     }
 
+    /**
+     * For synthetic functions (when, try, !!, but **not** elvis), we need to add an equality constraint for the expected type
+     * so that some type variables aren't inferred to `Nothing` that appears in one of the branches.
+     *
+     * @See org.jetbrains.kotlin.types.expressions.ControlStructureTypingUtils.createKnownTypeParameterSubstitutorForSpecialCall
+     */
+    private fun Candidate.isSyntheticFunctionCallThatShouldUseEqualityConstraint(expectedType: ConeKotlinType): Boolean {
+        // If we're inside an assignment's RHS, we mustn't add an equality constraint because it might prevent smartcasts.
+        // Example: val x: String? = null; x = if (foo) "" else throw Exception()
+        if (components.context.isInsideAssignmentRhs) return false
+
+        val symbol = symbol as? FirCallableSymbol ?: return false
+        if (symbol.origin != FirDeclarationOrigin.Synthetic.FakeFunction ||
+            expectedType.isUnitOrNullableUnit ||
+            expectedType.isAnyOrNullableAny ||
+            // We don't want to add an equality constraint to a nullable type to a !! call.
+            // See compiler/testData/diagnostics/tests/inference/checkNotNullWithNullableExpectedType.kt
+            (symbol.callableId == SyntheticCallableId.CHECK_NOT_NULL && expectedType.canBeNull(session))
+        ) {
+            return false
+        }
+
+        // If our expression contains any elvis, even nested, we mustn't add an equality constraint because it might influence the
+        // inferred type of the elvis RHS.
+        if (system.allTypeVariables.values.any {
+                it is ConeTypeParameterBasedTypeVariable && it.typeParameterSymbol.containingDeclarationSymbol.isSyntheticElvisFunction()
+            }
+        ) {
+            return false
+        }
+
+        return true
+    }
+
+    private fun FirBasedSymbol<*>.isSyntheticElvisFunction(): Boolean {
+        return origin == FirDeclarationOrigin.Synthetic.FakeFunction && (this as? FirCallableSymbol)?.callableId == SyntheticCallableId.ELVIS_NOT_NULL
+    }
+
     fun <T> runCompletionForCall(
         candidate: Candidate,
         completionMode: ConstraintSystemCompletionMode,

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildSamConversionExpression
 import org.jetbrains.kotlin.fir.references.FirNamedReference
+import org.jetbrains.kotlin.fir.references.FirResolvedErrorReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
@@ -936,9 +937,25 @@ class FirCallCompletionResultsWriterTransformer(
 
         runPCLARelatedTasksForCandidate(calleeReference.candidate)
 
-        return syntheticCall.apply {
-            replaceCalleeReference(calleeReference.toResolvedReference())
+        val resolvedCalleeReference = calleeReference.toResolvedReference()
+
+        // If we have a conflict between the expected type and the inferred type, we would like to set the inferred type on the expression,
+        // so that we report INITIALIZER_TYPE_MISMATCH/RETURN_TYPE_MISMATCH.
+        // This is required so that the IDE provides the correct quick fixes.
+        if (syntheticCall.resultType !is ConeErrorType && resolvedCalleeReference is FirResolvedErrorReference) {
+            val diagnostic = resolvedCalleeReference.diagnostic
+            if (diagnostic is ConeConstraintSystemHasContradiction) {
+                val candidate = diagnostic.candidate as Candidate
+                val newSyntheticCallType = session.typeContext.commonSuperTypeOrNull(candidate.argumentMapping!!.keys.map { it.resolvedType })
+                if (newSyntheticCallType != null && !newSyntheticCallType.hasError()) {
+                    syntheticCall.replaceConeTypeOrNull(newSyntheticCallType)
+                }
+            }
         }
+
+        syntheticCall.replaceCalleeReference(resolvedCalleeReference)
+
+        return syntheticCall
     }
 
     private inline fun <reified D> transformSyntheticCallChildren(

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

@@ -87,6 +87,20 @@ class BodyResolveContext(
     @set:PrivateForInline
     var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
 
+    @set:PrivateForInline
+    var isInsideAssignmentRhs: Boolean = false
+
+    @OptIn(PrivateForInline::class)
+    inline fun <R> withAssignmentRhs(block: () -> R): R {
+        val oldMode = this.isInsideAssignmentRhs
+        this.isInsideAssignmentRhs = true
+        return try {
+            block()
+        } finally {
+            this.isInsideAssignmentRhs = oldMode
+        }
+    }
+
     /**
      * This is required to avoid changing current mode into [FirTowerDataMode.CLASS_HEADER_ANNOTATIONS].
      * E.g., we can visit the same annotation in two ways – during a class visiting and outside of this class

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

@@ -1056,13 +1056,15 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
             }
         }
 
-        val result = variableAssignment.transformRValue(
-            transformer,
-            withExpectedType(
-                variableAssignment.lValue.resolvedType.toFirResolvedTypeRef(),
-                expectedTypeMismatchIsReportedInChecker = true
-            ),
-        )
+        val result = context.withAssignmentRhs {
+            variableAssignment.transformRValue(
+                transformer,
+                withExpectedType(
+                    variableAssignment.lValue.resolvedType.toFirResolvedTypeRef(),
+                    expectedTypeMismatchIsReportedInChecker = true,
+                ),
+            )
+        }
 
         // for cases like
         // buildSomething { tVar = "" // Should infer TV from String assignment }