[NI] Use refined expected type by smartcasts in CS

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt

@@ -30,12 +30,15 @@ import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.components.KotlinResolutionCallbacks
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
+import org.jetbrains.kotlin.resolve.calls.model.KotlinCall
 import org.jetbrains.kotlin.resolve.calls.model.LambdaKotlinCallArgument
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallAtom
 import org.jetbrains.kotlin.resolve.calls.model.SimpleKotlinCallArgument
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
 import org.jetbrains.kotlin.types.TypeApproximator
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.TypeUtils
@@ -44,6 +47,7 @@ import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.KotlinTypeInfo
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class KotlinResolutionCallbacksImpl(
         val topLevelCallContext: BasicCallResolutionContext,
@@ -144,4 +148,16 @@ class KotlinResolutionCallbacksImpl(
         kotlinToResolvedCallTransformer.createStubResolvedCallAndWriteItToTrace<CallableDescriptor>(candidate, trace)
     }
 
+    override fun createReceiverWithSmartCastInfo(resolvedAtom: ResolvedCallAtom): ReceiverValueWithSmartCastInfo? {
+        val returnType = resolvedAtom.candidateDescriptor.returnType ?: return null
+        val psiKotlinCall = resolvedAtom.atom.psiKotlinCall
+        val expression = psiKotlinCall.psiCall.callElement.safeAs<KtExpression>() ?: return null
+
+        return transformToReceiverWithSmartCastInfo(
+                resolvedAtom.candidateDescriptor,
+                trace.bindingContext,
+                psiKotlinCall.resultDataFlowInfo,
+                ExpressionReceiver.create(expression, returnType, trace.bindingContext)
+        )
+    }
 }

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitScopeTower
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
 import org.jetbrains.kotlin.types.UnwrappedType
 
 // stateless component
@@ -45,4 +46,6 @@ interface KotlinResolutionCallbacks {
     ): List<SimpleKotlinCallArgument>
 
     fun bindStubResolvedCallForCandidate(candidate: ResolvedCallAtom)
+
+    fun createReceiverWithSmartCastInfo(resolvedAtom: ResolvedCallAtom): ReceiverValueWithSmartCastInfo?
 }

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

@@ -51,14 +51,14 @@ class KotlinCallCompleter(
 
         if (candidate == null || candidate.csBuilder.hasContradiction) {
             val candidateForCompletion = candidate ?: factory.createErrorCandidate().forceResolution()
-            candidateForCompletion.prepareForCompletion(expectedType)
+            candidateForCompletion.prepareForCompletion(expectedType, resolutionCallbacks)
             runCompletion(candidateForCompletion.resolvedCall, ConstraintSystemCompletionMode.FULL, diagnosticHolder, candidateForCompletion.getSystem(), resolutionCallbacks)
 
             val systemStorage = candidate?.getSystem()?.asReadOnlyStorage() ?: ConstraintStorage.Empty
             return CallResolutionResult(CallResolutionResult.Type.ERROR, candidate?.resolvedCall, diagnosticHolder.getDiagnostics(), systemStorage)
         }
 
-        val completionType = candidate.prepareForCompletion(expectedType)
+        val completionType = candidate.prepareForCompletion(expectedType, resolutionCallbacks)
         val constraintSystem = candidate.getSystem()
         runCompletion(candidate.resolvedCall, completionType, diagnosticHolder, constraintSystem, resolutionCallbacks)
 
@@ -77,7 +77,7 @@ class KotlinCallCompleter(
     ): CallResolutionResult {
         val diagnosticsHolder = KotlinDiagnosticsHolder.SimpleHolder()
         for (candidate in candidates) {
-            candidate.prepareForCompletion(expectedType)
+            candidate.prepareForCompletion(expectedType, resolutionCallbacks)
             runCompletion(
                     candidate.resolvedCall,
                     ConstraintSystemCompletionMode.FULL,
@@ -109,9 +109,16 @@ class KotlinCallCompleter(
 
 
     // true if we should complete this call
-    private fun KotlinResolutionCandidate.prepareForCompletion(expectedType: UnwrappedType?): ConstraintSystemCompletionMode {
+    private fun KotlinResolutionCandidate.prepareForCompletion(
+            expectedType: UnwrappedType?,
+            resolutionCallbacks: KotlinResolutionCallbacks
+    ): ConstraintSystemCompletionMode {
         val unsubstitutedReturnType = resolvedCall.candidateDescriptor.returnType?.unwrap() ?: return ConstraintSystemCompletionMode.PARTIAL
-        val returnType = resolvedCall.substitutor.substituteKeepAnnotations(unsubstitutedReturnType)
+        val withSmartCastInfo = resolutionCallbacks.createReceiverWithSmartCastInfo(resolvedCall)
+
+        val actualType = withSmartCastInfo?.stableType ?: unsubstitutedReturnType
+
+        val returnType = resolvedCall.substitutor.substituteKeepAnnotations(actualType)
         if (expectedType != null && !TypeUtils.noExpectedType(expectedType)) {
             csBuilder.addSubtypeConstraint(returnType, expectedType, ExpectedTypeConstraintPosition(resolvedCall.atom))
         }