[NI] Refactor adding expected type constraint

Add comments, make computation needed for expected type more clear

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/KotlinCallCompleter.kt

@@ -28,18 +28,22 @@ class KotlinCallCompleter(
         resolutionCallbacks: KotlinResolutionCallbacks
     ): CallResolutionResult {
         val diagnosticHolder = KotlinDiagnosticsHolder.SimpleHolder()
-        if (candidates.isEmpty()) {
-            diagnosticHolder.addDiagnostic(NoneCandidatesCallDiagnostic(factory.kotlinCall))
-        }
-        if (candidates.size > 1) {
-            diagnosticHolder.addDiagnostic(ManyCandidatesCallDiagnostic(factory.kotlinCall, candidates))
+        when {
+            candidates.isEmpty() -> diagnosticHolder.addDiagnostic(NoneCandidatesCallDiagnostic(factory.kotlinCall))
+            candidates.size > 1 -> diagnosticHolder.addDiagnostic(ManyCandidatesCallDiagnostic(factory.kotlinCall, candidates))
         }
 
         val candidate = prepareCandidateForCompletion(factory, candidates, resolutionCallbacks)
-        val completionType = candidate.prepareForCompletion(expectedType, resolutionCallbacks)
+        val returnType = candidate.returnTypeWithSmartCastInfo(resolutionCallbacks)
+
+        candidate.addExpectedTypeConstraint(returnType, expectedType, resolutionCallbacks)
 
         return if (resolutionCallbacks.inferenceSession.shouldRunCompletion(candidate))
-            candidate.runCompletion(completionType, diagnosticHolder, resolutionCallbacks)
+            candidate.runCompletion(
+                candidate.computeCompletionMode(expectedType, returnType),
+                diagnosticHolder,
+                resolutionCallbacks
+            )
         else
             candidate.asCallResolutionResult(ConstraintSystemCompletionMode.PARTIAL, diagnosticHolder)
     }
@@ -51,7 +55,10 @@ class KotlinCallCompleter(
     ): CallResolutionResult {
         val diagnosticsHolder = KotlinDiagnosticsHolder.SimpleHolder()
         for (candidate in candidates) {
-            candidate.prepareForCompletion(expectedType, resolutionCallbacks)
+            candidate.addExpectedTypeConstraint(
+                candidate.returnTypeWithSmartCastInfo(resolutionCallbacks), expectedType, resolutionCallbacks
+            )
+
             runCompletion(
                 candidate.resolvedCall,
                 ConstraintSystemCompletionMode.FULL,
@@ -122,27 +129,40 @@ class KotlinCallCompleter(
         return candidate ?: factory.createErrorCandidate().forceResolution()
     }
 
-    // true if we should complete this call
-    private fun KotlinResolutionCandidate.prepareForCompletion(
+    private fun KotlinResolutionCandidate.returnTypeWithSmartCastInfo(resolutionCallbacks: KotlinResolutionCallbacks): UnwrappedType? {
+        val returnType = resolvedCall.candidateDescriptor.returnType?.unwrap() ?: return null
+        val returnTypeWithSmartCastInfo = computeReturnTypeWithSmartCastInfo(returnType, resolutionCallbacks)
+        return resolvedCall.substitutor.substituteKeepAnnotations(returnTypeWithSmartCastInfo)
+    }
+
+    private fun KotlinResolutionCandidate.addExpectedTypeConstraint(
+        returnType: UnwrappedType?,
         expectedType: UnwrappedType?,
         resolutionCallbacks: KotlinResolutionCallbacks
-    ): ConstraintSystemCompletionMode {
-        if (expectedType != null && TypeUtils.noExpectedType(expectedType)) return ConstraintSystemCompletionMode.FULL
+    ) {
+        if (returnType == null) return
+        if (expectedType == null || TypeUtils.noExpectedType(expectedType)) return
 
-        val returnType = resolvedCall.candidateDescriptor.returnType?.unwrap() ?: return ConstraintSystemCompletionMode.PARTIAL
-        val substitutedType: UnwrappedType
-        if (expectedType != null) {
-            val returnTypeWithSmartCastInfo = computeReturnTypeWithSmartCastInfo(returnType, resolutionCallbacks)
-            substitutedType = resolvedCall.substitutor.substituteKeepAnnotations(returnTypeWithSmartCastInfo)
-
-            if (!resolutionCallbacks.isCompileTimeConstant(resolvedCall, expectedType)) {
-                csBuilder.addSubtypeConstraint(substitutedType, expectedType, ExpectedTypeConstraintPosition(resolvedCall.atom))
-            }
-        } else {
-            substitutedType = resolvedCall.substitutor.substituteKeepAnnotations(returnType)
+        // We don't add expected type constraint for constant expression like "1 + 1" because of type coercion for numbers:
+        // val a: Long = 1 + 1, note that result type of "1 + 1" will be Int and adding constraint with Long will produce type mismatch
+        if (!resolutionCallbacks.isCompileTimeConstant(resolvedCall, expectedType)) {
+            csBuilder.addSubtypeConstraint(returnType, expectedType, ExpectedTypeConstraintPosition(resolvedCall.atom))
         }
+    }
 
-        return if (expectedType != null || csBuilder.isProperType(substitutedType))
+    private fun KotlinResolutionCandidate.computeCompletionMode(
+        expectedType: UnwrappedType?,
+        currentReturnType: UnwrappedType?
+    ): ConstraintSystemCompletionMode {
+        // Presence of expected type means that we trying to complete outermost call => completion mode should be full
+        if (expectedType != null) return ConstraintSystemCompletionMode.FULL
+
+        // This is questionable as null return type can be only for error call
+        if (currentReturnType == null) return ConstraintSystemCompletionMode.PARTIAL
+
+        // Consider call foo(bar(x)), if return type of bar is a proper one, then we can complete resolve for bar => full completion mode
+        // Otherwise, we shouldn't complete bar until we process call foo
+        return if (csBuilder.isProperType(currentReturnType))
             ConstraintSystemCompletionMode.FULL
         else
             ConstraintSystemCompletionMode.PARTIAL