KT-4915 Smart completion should work for auto-casted 'this'

#KT-4915 Fixed

idea/idea-core/src/org/jetbrains/kotlin/idea/core/SmartCastCalculator.kt

@@ -16,9 +16,9 @@
 
 package org.jetbrains.kotlin.idea.core
 
-import com.google.common.collect.SetMultimap
 import com.intellij.openapi.util.Pair
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.psi.JetExpression
 import org.jetbrains.kotlin.psi.JetSimpleNameExpression
@@ -32,83 +32,83 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.Nullability
 import org.jetbrains.kotlin.resolve.scopes.receivers.ThisReceiver
 import org.jetbrains.kotlin.types.JetType
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
-import org.jetbrains.kotlin.utils.singletonOrEmptyList
-import java.util.Collections
 import java.util.HashMap
-import java.util.HashSet
 
 class SmartCastCalculator(
         val bindingContext: BindingContext,
         val containingDeclarationOrModule: DeclarationDescriptor,
-        nameExpression: JetSimpleNameExpression
-): (VariableDescriptor) -> Collection<JetType> {
+        expression: JetExpression
+) {
+    // keys are VariableDescriptor's and ThisReceiver's
+    private val entityToSmartCastInfo: Map<Any, SmartCastInfo>
+            = processDataFlowInfo(bindingContext.getDataFlowInfo(expression), if (expression is JetSimpleNameExpression) expression.getReceiverExpression() else null)
 
-    override fun invoke(descriptor: VariableDescriptor): Collection<JetType> {
-        return function(descriptor)
+    fun types(descriptor: VariableDescriptor): Collection<JetType> {
+        val type = descriptor.returnType ?: return emptyList()
+        return entityType(descriptor, type)
     }
 
-    private val function: (VariableDescriptor) -> Collection<JetType> = run {
-        val receiver = nameExpression.getReceiverExpression()
-        val dataFlowInfo = bindingContext.getDataFlowInfo(nameExpression)
-        val (variableToTypes, notNullVariables) = processDataFlowInfo(dataFlowInfo, receiver)
+    fun types(thisReceiverParameter: ReceiverParameterDescriptor): Collection<JetType> {
+        val type = thisReceiverParameter.type
+        val thisReceiver = thisReceiverParameter.value as? ThisReceiver ?: return listOf(type)
+        return entityType(thisReceiver, type)
+    }
 
-        fun typesOf(descriptor: VariableDescriptor): Collection<JetType> {
-            var type = descriptor.getReturnType() ?: return listOf()
-            if (notNullVariables.contains(descriptor)) {
-                type = type.makeNotNullable()
-            }
+    private fun entityType(entity: Any, ownType: JetType): Collection<JetType> {
+        val smartCastInfo = entityToSmartCastInfo[entity] ?: return listOf(ownType)
 
-            val smartCastTypes = variableToTypes[descriptor]
-            if (smartCastTypes == null || smartCastTypes.isEmpty()) return type.singletonOrEmptyList()
-            return smartCastTypes + type.singletonOrEmptyList()
+        var types = smartCastInfo.types + ownType
+
+        if (smartCastInfo.notNull) {
+            types = types.map { it.makeNotNullable() }
         }
 
-        ::typesOf
+        return types
     }
 
-    private data class ProcessDataFlowInfoResult(
-            val variableToTypes: Map<VariableDescriptor, Collection<JetType>> = Collections.emptyMap(),
-            val notNullVariables: Set<VariableDescriptor> = Collections.emptySet()
-    )
+    private data class SmartCastInfo(var types: Collection<JetType>, var notNull: Boolean) {
+        constructor() : this(emptyList(), false)
+    }
 
-    private fun processDataFlowInfo(dataFlowInfo: DataFlowInfo, receiver: JetExpression?): ProcessDataFlowInfoResult {
-        if (dataFlowInfo == DataFlowInfo.EMPTY) return ProcessDataFlowInfoResult()
+    private fun processDataFlowInfo(dataFlowInfo: DataFlowInfo, receiver: JetExpression?): Map<Any, SmartCastInfo> {
+        if (dataFlowInfo == DataFlowInfo.EMPTY) return emptyMap()
 
-        val dataFlowValueToVariable: (DataFlowValue) -> VariableDescriptor?
+        val dataFlowValueToEntity: (DataFlowValue) -> Any?
         if (receiver != null) {
-            val receiverType = bindingContext.getType(receiver) ?: return ProcessDataFlowInfoResult()
-            val receiverId = DataFlowValueFactory.createDataFlowValue(receiver, receiverType, bindingContext, containingDeclarationOrModule).getId()
-            dataFlowValueToVariable = { value ->
-                val id = value.getId()
+            val receiverType = bindingContext.getType(receiver) ?: return emptyMap()
+            val receiverId = DataFlowValueFactory.createDataFlowValue(receiver, receiverType, bindingContext, containingDeclarationOrModule).id
+            dataFlowValueToEntity = { value ->
+                val id = value.id
                 if (id is Pair<*, *> && id.first == receiverId) id.second as? VariableDescriptor else null
             }
         }
         else {
-            dataFlowValueToVariable = { value ->
-                val id = value.getId()
-                when {
-                    id is VariableDescriptor -> id
-                    id is Pair<*, *> && id.first is ThisReceiver -> id.second as? VariableDescriptor
+            dataFlowValueToEntity = { value ->
+                val id = value.id
+                when(id) {
+                    is VariableDescriptor, is ThisReceiver -> id
+                    is Pair<*, *> -> if (id.first is ThisReceiver) id.second as? VariableDescriptor else null
                     else -> null
                 }
             }
         }
 
-        val variableToType = HashMap<VariableDescriptor, Collection<JetType>>()
-        val typeInfo: SetMultimap<DataFlowValue, JetType> = dataFlowInfo.getCompleteTypeInfo()
-        for ((dataFlowValue, types) in typeInfo.asMap().entrySet()) {
-            val variable = dataFlowValueToVariable.invoke(dataFlowValue)
-            if (variable != null) {
-                variableToType[variable] = types
+        val entityToInfo = HashMap<Any, SmartCastInfo>()
+
+        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo.asMap().entrySet()) {
+            val entity = dataFlowValueToEntity.invoke(dataFlowValue)
+            if (entity != null) {
+                entityToInfo[entity] = SmartCastInfo(types, false)
             }
         }
 
-        val nullabilityInfo: Map<DataFlowValue, Nullability> = dataFlowInfo.getCompleteNullabilityInfo()
-        val notNullVariables = nullabilityInfo
-                .filter { it.getValue() == Nullability.NOT_NULL }
-                .map { dataFlowValueToVariable(it.getKey()) }
-                .filterNotNullTo(HashSet<VariableDescriptor>())
+        for ((dataFlowValue, nullability) in dataFlowInfo.completeNullabilityInfo) {
+            if (nullability == Nullability.NOT_NULL) {
+                val entity = dataFlowValueToEntity(dataFlowValue) ?: continue
+                entityToInfo.getOrPut(entity, { SmartCastInfo() }).notNull = true
+            }
+        }
 
-        return ProcessDataFlowInfoResult(variableToType, notNullVariables)
+        return entityToInfo
     }
 }