[NI] Support smartcast info in ResolvedCall on receivers

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/DiagnosticReporterByTrackingStrategy.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.diagnostics.Errors.BadNamedArgumentsTarget.INVOKE_ON
 import org.jetbrains.kotlin.diagnostics.Errors.BadNamedArgumentsTarget.NON_KOTLIN_FUNCTION
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
@@ -126,28 +127,40 @@ class DiagnosticReporterByTrackingStrategy(
 
     private fun reportSmartCast(smartCastDiagnostic: SmartCastDiagnostic) {
         val expressionArgument = smartCastDiagnostic.argument
-        if (expressionArgument is ExpressionKotlinCallArgumentImpl) {
-            val context = context.replaceDataFlowInfo(expressionArgument.dataFlowInfoBeforeThisArgument)
-            val argumentExpression = KtPsiUtil.getLastElementDeparenthesized(expressionArgument.valueArgument.getArgumentExpression (), context.statementFilter)
-            val dataFlowValue = DataFlowValueFactory.createDataFlowValue(expressionArgument.receiver.receiverValue, context)
-            SmartCastManager.checkAndRecordPossibleCast(
-                    dataFlowValue, smartCastDiagnostic.smartCastType, argumentExpression, context, call,
-                    recordExpressionType = true)
-            trace.markAsReported()
+        val smartCastResult = when (expressionArgument) {
+            is ExpressionKotlinCallArgumentImpl -> {
+                trace.markAsReported()
+                val context = context.replaceDataFlowInfo(expressionArgument.dataFlowInfoBeforeThisArgument)
+                val argumentExpression = KtPsiUtil.getLastElementDeparenthesized(expressionArgument.valueArgument.getArgumentExpression (), context.statementFilter)
+                val dataFlowValue = DataFlowValueFactory.createDataFlowValue(expressionArgument.receiver.receiverValue, context)
+                SmartCastManager.checkAndRecordPossibleCast(
+                        dataFlowValue, smartCastDiagnostic.smartCastType, argumentExpression, context, call,
+                        recordExpressionType = true)
+            }
+            is ReceiverExpressionKotlinCallArgument -> {
+                trace.markAsReported()
+                val receiverValue = expressionArgument.receiver.receiverValue
+                val dataFlowValue = DataFlowValueFactory.createDataFlowValue(receiverValue, context)
+                SmartCastManager.checkAndRecordPossibleCast(
+                        dataFlowValue, smartCastDiagnostic.smartCastType, (receiverValue as? ExpressionReceiver)?.expression, context, call,
+                        recordExpressionType = true)
+            }
+            else -> null
         }
-        else if(expressionArgument is ReceiverExpressionKotlinCallArgument) {
-            val receiverValue = expressionArgument.receiver.receiverValue
-            val dataFlowValue = DataFlowValueFactory.createDataFlowValue(receiverValue, context)
-            SmartCastManager.checkAndRecordPossibleCast(
-                    dataFlowValue, smartCastDiagnostic.smartCastType, (receiverValue as? ExpressionReceiver)?.expression, context, call,
-                    recordExpressionType = true)
-            trace.markAsReported()
+        val resolvedCall = smartCastDiagnostic.kotlinCall?.psiKotlinCall?.psiCall?.getResolvedCall(trace.bindingContext) as? NewResolvedCallImpl<*>
+        if (resolvedCall != null && smartCastResult != null) {
+            if (resolvedCall.extensionReceiver == expressionArgument.receiver.receiverValue) {
+                resolvedCall.updateExtensionReceiverWithSmartCastIfNeeded(smartCastResult.resultType)
+            }
+            if (resolvedCall.dispatchReceiver == expressionArgument.receiver.receiverValue) {
+                resolvedCall.setSmartCastDispatchReceiverType(smartCastResult.resultType)
+            }
         }
     }
 
     private fun reportUnstableSmartCast(unstableSmartCast: UnstableSmartCast) {
         // todo hack -- remove it after removing SmartCastManager
-        reportSmartCast(SmartCastDiagnostic(unstableSmartCast.argument, unstableSmartCast.targetType))
+        reportSmartCast(SmartCastDiagnostic(unstableSmartCast.argument, unstableSmartCast.targetType, null))
     }
 
     override fun constraintError(diagnostic: KotlinCallDiagnostic) {

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

@@ -45,6 +45,8 @@ import org.jetbrains.kotlin.resolve.calls.results.ResolutionStatus
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
+import org.jetbrains.kotlin.resolve.scopes.receivers.CastImplicitClassReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitClassReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.NewCapturedType
@@ -264,21 +266,22 @@ class KotlinToResolvedCallTransformer(
     }
 
     private fun bindAndReport(context: BasicCallResolutionContext, trace: BindingTrace, simpleResolvedCall: NewResolvedCallImpl<*>) {
-        reportCallDiagnostic(context, trace, simpleResolvedCall.resolvedCallAtom, simpleResolvedCall.resultingDescriptor)
         val tracing = simpleResolvedCall.resolvedCallAtom.atom.psiKotlinCall.tracingStrategy
 
         tracing.bindReference(trace, simpleResolvedCall)
         tracing.bindResolvedCall(trace, simpleResolvedCall)
+
+        reportCallDiagnostic(context, trace, simpleResolvedCall.resolvedCallAtom, simpleResolvedCall.resultingDescriptor)
     }
 
     private fun bindAndReport(context: BasicCallResolutionContext, trace: BindingTrace, variableAsFunction: NewVariableAsFunctionResolvedCallImpl) {
-        reportCallDiagnostic(context, trace, variableAsFunction.variableCall.resolvedCallAtom, variableAsFunction.variableCall.resultingDescriptor)
-        reportCallDiagnostic(context, trace, variableAsFunction.functionCall.resolvedCallAtom, variableAsFunction.functionCall.resultingDescriptor)
-
         val outerTracingStrategy = variableAsFunction.baseCall.tracingStrategy
         outerTracingStrategy.bindReference(trace, variableAsFunction.variableCall)
         outerTracingStrategy.bindResolvedCall(trace, variableAsFunction)
         variableAsFunction.functionCall.kotlinCall.psiKotlinCall.tracingStrategy.bindReference(trace, variableAsFunction.functionCall)
+
+        reportCallDiagnostic(context, trace, variableAsFunction.variableCall.resolvedCallAtom, variableAsFunction.variableCall.resultingDescriptor)
+        reportCallDiagnostic(context, trace, variableAsFunction.functionCall.resolvedCallAtom, variableAsFunction.functionCall.resultingDescriptor)
     }
 
     private fun reportCallDiagnostic(
@@ -431,6 +434,9 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
         TypeApproximator().approximateToSuperType(substituted, TypeApproximatorConfiguration.CapturedTypesApproximation) ?: substituted
     }
 
+    private var extensionReceiver = resolvedCallAtom.extensionReceiverArgument?.receiver?.receiverValue
+    private var smartCastDispatchReceiverType: KotlinType? = null
+
     override val kotlinCall: KotlinCall get() = resolvedCallAtom.atom
 
     override fun getStatus(): ResolutionStatus = getResultApplicability(resolvedCallAtom.diagnostics).toResolutionStatus()
@@ -440,7 +446,7 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
 
     override fun getCandidateDescriptor(): D = resolvedCallAtom.candidateDescriptor as D
     override fun getResultingDescriptor(): D = resultingDescriptor as D
-    override fun getExtensionReceiver(): ReceiverValue? = resolvedCallAtom.extensionReceiverArgument?.receiver?.receiverValue
+    override fun getExtensionReceiver(): ReceiverValue? = extensionReceiver
     override fun getDispatchReceiver(): ReceiverValue? = resolvedCallAtom.dispatchReceiverArgument?.receiver?.receiverValue
     override fun getExplicitReceiverKind(): ExplicitReceiverKind = resolvedCallAtom.explicitReceiverKind
 
@@ -449,7 +455,20 @@ class NewResolvedCallImpl<D : CallableDescriptor>(
         return typeParameters.zip(typeArguments).toMap()
     }
 
-    override fun getSmartCastDispatchReceiverType(): KotlinType? = null // todo
+    override fun getSmartCastDispatchReceiverType(): KotlinType? = smartCastDispatchReceiverType
+
+    fun updateExtensionReceiverWithSmartCastIfNeeded(smartCastExtensionReceiverType: KotlinType) {
+        if (extensionReceiver is ImplicitClassReceiver) {
+            extensionReceiver = CastImplicitClassReceiver(
+                    (extensionReceiver as ImplicitClassReceiver).classDescriptor,
+                    smartCastExtensionReceiverType
+            )
+        }
+    }
+
+    fun setSmartCastDispatchReceiverType(smartCastDispatchReceiverType: KotlinType) {
+        this.smartCastDispatchReceiverType = smartCastDispatchReceiverType
+    }
 }
 
 fun ResolutionCandidateApplicability.toResolutionStatus(): ResolutionStatus = when (this) {

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

@@ -35,10 +35,14 @@ class AdditionalDiagnosticReporter {
         reportSmartCasts(candidate, resultingDescriptor, kotlinDiagnosticsHolder)
     }
 
-    private fun createSmartCastDiagnostic(argument: KotlinCallArgument, expectedResultType: UnwrappedType): SmartCastDiagnostic? {
+    private fun createSmartCastDiagnostic(
+            candidate: ResolvedCallAtom,
+            argument: KotlinCallArgument,
+            expectedResultType: UnwrappedType
+    ): SmartCastDiagnostic? {
         if (argument !is ExpressionKotlinCallArgument) return null
         if (!KotlinTypeChecker.DEFAULT.isSubtypeOf(argument.receiver.receiverValue.type, expectedResultType)) {
-            return SmartCastDiagnostic(argument, expectedResultType.unwrap())
+            return SmartCastDiagnostic(argument, expectedResultType.unwrap(), candidate.atom)
         }
         return null
     }
@@ -51,7 +55,7 @@ class AdditionalDiagnosticReporter {
         if (receiver == null || parameter == null) return null
         val expectedType = parameter.type.unwrap().let { if (receiver.isSafeCall) it.makeNullableAsSpecified(true) else it }
 
-        val smartCastDiagnostic = createSmartCastDiagnostic(receiver, expectedType) ?: return null
+        val smartCastDiagnostic = createSmartCastDiagnostic(candidate, receiver, expectedType) ?: return null
 
         // todo may be we have smart cast to Int?
         return smartCastDiagnostic.takeIf {
@@ -75,7 +79,7 @@ class AdditionalDiagnosticReporter {
 
         for (parameter in resultingDescriptor.valueParameters) {
             for (argument in candidate.argumentMappingByOriginal[parameter.original]?.arguments ?: continue) {
-                val smartCastDiagnostic = createSmartCastDiagnostic(argument, argument.getExpectedType(parameter)) ?: continue
+                val smartCastDiagnostic = createSmartCastDiagnostic(candidate, argument, argument.getExpectedType(parameter)) ?: continue
 
                 val thereIsUnstableSmartCastError = candidate.diagnostics.filterIsInstance<UnstableSmartCast>().any {
                     it.argument == argument

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

@@ -58,7 +58,7 @@ internal object CheckVisibility : ResolutionPart() {
         if (dispatchReceiverArgument is ExpressionKotlinCallArgument) {
             val smartCastReceiver = getReceiverValueWithSmartCast(receiverValue, dispatchReceiverArgument.receiver.stableType)
             if (Visibilities.findInvisibleMember(smartCastReceiver, candidateDescriptor, containingDescriptor) == null) {
-                addDiagnostic(SmartCastDiagnostic(dispatchReceiverArgument, dispatchReceiverArgument.receiver.stableType))
+                addDiagnostic(SmartCastDiagnostic(dispatchReceiverArgument, dispatchReceiverArgument.receiver.stableType, resolvedCall.atom))
                 return
             }
         }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/KotlinCallDiagnistics.kt

@@ -128,7 +128,8 @@ class NotCallableExpectedType(
 // SmartCasts
 class SmartCastDiagnostic(
         val argument: ExpressionKotlinCallArgument,
-        val smartCastType: UnwrappedType
+        val smartCastType: UnwrappedType,
+        val kotlinCall: KotlinCall?
 ): KotlinCallDiagnostic(RESOLVED) {
     override fun report(reporter: DiagnosticReporter) = reporter.onCallArgument(argument, this)
 }