[NI] Capture types from smart-cast types.

Sometimes we have something like if (a is Foo<*>) a.bar()
where bar declared: fun <T> Foo<T>.bar().
For such case we should create receiver with possible types Capture(*).

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

@@ -26,13 +26,12 @@ import org.jetbrains.kotlin.resolve.StatementFilter
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCall
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.resolve.calls.model.LambdaKotlinCallArgument
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.prepareReceiverRegardingCaptureTypes
 import org.jetbrains.kotlin.types.UnwrappedType
-import org.jetbrains.kotlin.types.checker.prepareArgumentTypeRegardingCaptureTypes
 import org.jetbrains.kotlin.types.expressions.KotlinTypeInfo
 
 class SimpleTypeArgumentImpl(
@@ -172,16 +171,13 @@ internal fun createSimplePSICallArgument(
         bindingContext.get(BindingContext.ONLY_RESOLVED_CALL, it)
     }
     val baseType = onlyResolvedCall?.currentReturnType ?: typeInfoForArgument.type?.unwrap() ?: return null
-    val preparedType = prepareArgumentTypeRegardingCaptureTypes(baseType) ?: baseType
 
     // we should use DFI after this argument, because there can be some useful smartcast. Popular case: if branches.
     val receiverToCast = transformToReceiverWithSmartCastInfo(
             ownerDescriptor, bindingContext,
             typeInfoForArgument.dataFlowInfo,
             ExpressionReceiver.create(ktExpression, baseType, bindingContext)
-    ).let {
+    ).prepareReceiverRegardingCaptureTypes()
-        ReceiverValueWithSmartCastInfo(it.receiverValue.replaceType(preparedType), it.possibleTypes, it.isStable)
-    }
 
     return if (onlyResolvedCall == null) {
         ExpressionKotlinCallArgumentImpl(valueArgument, dataFlowInfoBeforeThisArgument, typeInfoForArgument.dataFlowInfo, receiverToCast)

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

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
-import org.jetbrains.kotlin.types.checker.prepareArgumentTypeRegardingCaptureTypes
+import org.jetbrains.kotlin.resolve.scopes.receivers.prepareReceiverRegardingCaptureTypes
 
 
 class FakeKotlinCallArgumentForCallableReference(
@@ -45,14 +45,7 @@ class ReceiverExpressionKotlinCallArgument private constructor(
                 receiver: ReceiverValueWithSmartCastInfo,
                    isSafeCall: Boolean = false,
                    isVariableReceiverForInvoke: Boolean = false
-        ): ReceiverExpressionKotlinCallArgument {
+        ) = ReceiverExpressionKotlinCallArgument(receiver.prepareReceiverRegardingCaptureTypes(), isSafeCall, isVariableReceiverForInvoke)
-            val newType = prepareArgumentTypeRegardingCaptureTypes(receiver.receiverValue.type.unwrap())
-            val newReceiver = if (newType != null) {
-                ReceiverValueWithSmartCastInfo(receiver.receiverValue.replaceType(newType), receiver.possibleTypes, receiver.isStable)
-            } else receiver
-
-            return ReceiverExpressionKotlinCallArgument(newReceiver, isSafeCall, isVariableReceiverForInvoke)
-        }
     }
 }
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/scopes/receivers/QualifierReceiver.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.resolve.scopes.receivers
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.checker.prepareArgumentTypeRegardingCaptureTypes
 
 // this receiver used only for resolution. see subtypes
 interface DetailedReceiver
@@ -42,3 +43,13 @@ interface QualifierReceiver : Receiver, DetailedReceiver {
     val classValueReceiverWithSmartCastInfo: ReceiverValueWithSmartCastInfo?
         get() = classValueReceiver?.let { ReceiverValueWithSmartCastInfo(it, emptySet(), true) }
 }
+
+fun ReceiverValueWithSmartCastInfo.prepareReceiverRegardingCaptureTypes(): ReceiverValueWithSmartCastInfo {
+    val preparedBaseType = prepareArgumentTypeRegardingCaptureTypes(receiverValue.type.unwrap())
+    if (preparedBaseType == null && possibleTypes.isEmpty()) return this
+
+    val newPossibleTypes = possibleTypes.mapTo(HashSet<KotlinType>()) { prepareArgumentTypeRegardingCaptureTypes(it.unwrap()) ?: it }
+    val newReceiver = if (preparedBaseType != null) receiverValue.replaceType(preparedBaseType) else receiverValue
+
+    return ReceiverValueWithSmartCastInfo(newReceiver, newPossibleTypes, isStable)
+}