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

#KT-4915 Fixed

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/BasicCompletionSession.kt

@@ -35,7 +35,10 @@ import org.jetbrains.kotlin.idea.stubindex.PackageIndexUtil
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.*
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.isAncestor
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindExclude
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.utils.addToStdlib.check
@@ -194,7 +197,7 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
                 // if "this" is parsed correctly in the current context - insert it and all this@xxx items
                     "this" -> {
                         if (expression != null) {
-                            collector.addElements(thisExpressionItems(bindingContext, expression, prefix).map { it.factory() })
+                            collector.addElements(thisExpressionItems(bindingContext, expression, prefix).map { it.createLookupElement() })
                         }
                         else {
                             // for completion in secondary constructor delegation call
@@ -281,14 +284,13 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
             sorter = sorter.weighBefore(DeprecatedWeigher.toString(), ParameterNameAndTypeCompletion.Weigher)
         }
 
-        val expectedInfos = nameExpression
+        if (expression != null && completionKind == CompletionKind.ALL) {
-                ?.check { completionKind == CompletionKind.ALL }
+            val expectedInfos = ExpectedInfos(bindingContext, resolutionFacade, moduleDescriptor).calculate(expression)
-                ?.let { it.getQualifiedElement() as? JetExpression }
-                ?.let { ExpectedInfos(bindingContext, resolutionFacade, moduleDescriptor).calculate(it) }
 
-        if (expectedInfos != null && expectedInfos.isNotEmpty()) {
+            if (expectedInfos != null && expectedInfos.isNotEmpty()) {
-            val smartCastCalculator = SmartCastCalculator(bindingContext, moduleDescriptor, nameExpression!!)
+                val smartCastCalculator = SmartCastCalculator(bindingContext, moduleDescriptor, expression)
-            sorter = sorter.weighBefore(KindWeigher.toString(), ExpectedInfoMatchWeigher(expectedInfos, smartCastCalculator))
+                sorter = sorter.weighBefore(KindWeigher.toString(), ExpectedInfoMatchWeigher(expectedInfos, smartCastCalculator))
+            }
         }
 
         return sorter

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionUtils.kt

@@ -30,7 +30,10 @@ import org.jetbrains.kotlin.idea.JetIcons
 import org.jetbrains.kotlin.idea.caches.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.completion.handlers.CastReceiverInsertHandler
 import org.jetbrains.kotlin.idea.completion.handlers.WithTailInsertHandler
-import org.jetbrains.kotlin.idea.util.*
+import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
+import org.jetbrains.kotlin.idea.util.ShortenReferences
+import org.jetbrains.kotlin.idea.util.findLabelAndCall
+import org.jetbrains.kotlin.idea.util.getImplicitReceiversWithInstanceToExpression
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
@@ -159,26 +162,22 @@ fun shouldCompleteThisItems(prefixMatcher: PrefixMatcher): Boolean {
     return prefix.startsWith(s) || s.startsWith(prefix)
 }
 
-data class ThisItemInfo(val factory: () -> LookupElement, val type: FuzzyType)
+class ThisItemLookupObject(val receiverParameter: ReceiverParameterDescriptor, val labelName: Name?) : KeywordLookupObject()
 
-fun thisExpressionItems(bindingContext: BindingContext, position: JetExpression, prefix: String): Collection<ThisItemInfo> {
+fun ThisItemLookupObject.createLookupElement() = createKeywordWithLabelElement("this", labelName, lookupObject = this)
+        .withTypeText(DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(receiverParameter.type))
+
+fun thisExpressionItems(bindingContext: BindingContext, position: JetExpression, prefix: String): Collection<ThisItemLookupObject> {
     val scope = bindingContext[BindingContext.RESOLUTION_SCOPE, position] ?: return listOf()
 
     val psiFactory = JetPsiFactory(position)
 
-    val result = ArrayList<ThisItemInfo>()
+    val result = ArrayList<ThisItemLookupObject>()
     for ((receiver, expressionFactory) in scope.getImplicitReceiversWithInstanceToExpression()) {
         if (expressionFactory == null) continue
         // if prefix does not start with "this@" do not include immediate this in the form with label
         val expression = expressionFactory.createExpression(psiFactory, shortThis = !prefix.startsWith("this@")) as? JetThisExpression ?: continue
-
+        result.add(ThisItemLookupObject(receiver, expression.getLabelNameAsName()))
-        val thisType = receiver.getType()
-        val fuzzyType = FuzzyType(thisType, listOf())
-
-        fun createLookupElement() = createKeywordWithLabelElement("this", expression.getLabelNameAsName())
-                .withTypeText(DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(thisType))
-
-        result.add(ThisItemInfo(::createLookupElement, fuzzyType))
     }
     return result
 }
@@ -223,8 +222,8 @@ private fun JetDeclarationWithBody.returnType(bindingContext: BindingContext): J
     return callable.getReturnType()
 }
 
-private fun createKeywordWithLabelElement(keyword: String, label: Name?, addSpace: Boolean): LookupElement {
+private fun createKeywordWithLabelElement(keyword: String, label: Name?, addSpace: Boolean, lookupObject: KeywordLookupObject = KeywordLookupObject()): LookupElement {
-    val element = createKeywordWithLabelElement(keyword, label)
+    val element = createKeywordWithLabelElement(keyword, label, lookupObject)
     return if (addSpace) {
         object: LookupElementDecorator<LookupElement>(element) {
             override fun handleInsert(context: InsertionContext) {
@@ -237,9 +236,9 @@ private fun createKeywordWithLabelElement(keyword: String, label: Name?, addSpac
     }
 }
 
-private fun createKeywordWithLabelElement(keyword: String, label: Name?): LookupElementBuilder {
+private fun createKeywordWithLabelElement(keyword: String, label: Name?, lookupObject: KeywordLookupObject = KeywordLookupObject()): LookupElementBuilder {
     val labelInCode = label?.render()
-    var element = LookupElementBuilder.create(KeywordLookupObject, if (label == null) keyword else "$keyword@$labelInCode")
+    var element = LookupElementBuilder.create(lookupObject, if (label == null) keyword else "$keyword@$labelInCode")
     element = element.withPresentableText(keyword)
     element = element.withBoldness(true)
     if (label != null) {

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KeywordCompletion.kt

@@ -37,7 +37,7 @@ import org.jetbrains.kotlin.psi.psiUtil.nextLeaf
 import org.jetbrains.kotlin.psi.psiUtil.prevLeaf
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 
-object KeywordLookupObject
+open class KeywordLookupObject
 
 object KeywordCompletion {
     private val NON_ACTUAL_KEYWORDS = setOf(CAPITALIZED_THIS_KEYWORD,
@@ -72,7 +72,7 @@ object KeywordCompletion {
             }
 
             if (keyword.startsWith(prefix)/* use simple matching by prefix, not prefix matcher from completion*/ && parserFilter(keywordToken)) {
-                val element = LookupElementBuilder.create(KeywordLookupObject, keyword)
+                val element = LookupElementBuilder.create(KeywordLookupObject(), keyword)
                         .bold()
                         .withInsertHandler(if (keywordToken !in FUNCTION_KEYWORDS)
                                                KotlinKeywordInsertHandler

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/MultipleArgumentsItemProvider.kt

@@ -91,7 +91,7 @@ class MultipleArgumentsItemProvider(val bindingContext: BindingContext,
         val name = parameter.getName()
         //TODO: there can be more than one property with such name in scope and we should be able to select one (but we need API for this)
         val variable = scope.getLocalVariable(name) ?: scope.getProperties(name).singleOrNull() ?: return null
-        return if (smartCastCalculator(variable).any { JetTypeChecker.DEFAULT.isSubtypeOf(it, parameter.getType()) })
+        return if (smartCastCalculator.types(variable).any { JetTypeChecker.DEFAULT.isSubtypeOf(it, parameter.getType()) })
             variable
         else
             null

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/SmartCompletion.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.completion.*
 import org.jetbrains.kotlin.idea.core.SmartCastCalculator
+import org.jetbrains.kotlin.idea.util.FuzzyType
 import org.jetbrains.kotlin.idea.util.isAlmostEverything
 import org.jetbrains.kotlin.idea.util.makeNullable
 import org.jetbrains.kotlin.lexer.JetTokens
@@ -125,14 +126,13 @@ class SmartCompletion(
         else
             originalExpectedInfos
 
-        val smartCastCalculator = (expression as? JetSimpleNameExpression)?.let { SmartCastCalculator(bindingContext, moduleDescriptor, it) }
+        val smartCastCalculator = SmartCastCalculator(bindingContext, moduleDescriptor, expression)
 
         val itemsToSkip = calcItemsToSkip(expressionWithType)
 
         val functionExpectedInfos = expectedInfos.filter { it.fuzzyType != null && KotlinBuiltIns.isExactFunctionOrExtensionFunctionType(it.fuzzyType!!.type) }
 
         fun filterDeclaration(descriptor: DeclarationDescriptor): Collection<LookupElement> {
-            if (smartCastCalculator == null) return emptyList() // only happens for this@ completion
             if (descriptor in itemsToSkip) return emptyList()
 
             val result = SmartList<LookupElement>()
@@ -166,15 +166,13 @@ class SmartCompletion(
                 StaticMembers(bindingContext, lookupElementFactory).addToCollection(additionalItems, expectedInfos, expression, itemsToSkip)
             }
 
-            additionalItems.addThisItems(expression, expectedInfos)
+            additionalItems.addThisItems(expression, expectedInfos, smartCastCalculator)
 
             LambdaItems.addToCollection(additionalItems, functionExpectedInfos)
 
             KeywordValues.addToCollection(additionalItems, originalExpectedInfos/* use originalExpectedInfos to not include null after == */, expression)
 
-            if (smartCastCalculator != null) {
+            MultipleArgumentsItemProvider(bindingContext, smartCastCalculator).addToCollection(additionalItems, expectedInfos, expression)
-                MultipleArgumentsItemProvider(bindingContext, smartCastCalculator).addToCollection(additionalItems, expectedInfos, expression)
-            }
         }
 
         val inheritanceSearcher = if (inheritanceSearchers.isNotEmpty())
@@ -188,13 +186,14 @@ class SmartCompletion(
         return Result(::filterDeclaration, additionalItems, inheritanceSearcher)
     }
 
-    private fun MutableCollection<LookupElement>.addThisItems(place: JetExpression, expectedInfos: Collection<ExpectedInfo>) {
+    private fun MutableCollection<LookupElement>.addThisItems(place: JetExpression, expectedInfos: Collection<ExpectedInfo>, smartCastCalculator: SmartCastCalculator) {
         if (shouldCompleteThisItems(prefixMatcher)) {
             val items = thisExpressionItems(bindingContext, place, prefixMatcher.getPrefix())
-            for ((factory, type) in items) {
+            for (item in items) {
-                val classifier = { expectedInfo: ExpectedInfo -> type.classifyExpectedInfo(expectedInfo) }
+                val types = smartCastCalculator.types(item.receiverParameter).map { FuzzyType(it, emptyList()) }
+                val classifier = { expectedInfo: ExpectedInfo -> types.classifyExpectedInfo(expectedInfo) }
                 addLookupElements(null, expectedInfos, classifier) {
-                    factory().assignSmartCompletionPriority(SmartCompletionItemPriority.THIS)
+                    item.createLookupElement().assignSmartCompletionPriority(SmartCompletionItemPriority.THIS)
                 }
             }
         }

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/Utils.kt

@@ -326,7 +326,7 @@ fun DeclarationDescriptor.fuzzyTypesForSmartCompletion(smartCastCalculator: Smar
         if (returnType.type.isNothing() || returnType.isAlmostEverything()) return emptyList()
 
         if (this is VariableDescriptor) { //TODO: generic properties!
-            return smartCastCalculator(this).map { FuzzyType(it, emptyList()) }
+            return smartCastCalculator.types(this).map { FuzzyType(it, emptyList()) }
         }
         else {
             return listOf(returnType)

idea/idea-completion/testData/smart/smartCasts/SmartCastThisType.kt.todo

@@ -1,12 +0,0 @@
-open class Foo{
-    fun f() {
-        if (this is Bar){
-            var a : Bar = <caret>
-        }
-    }
-}
-
-class Bar : Foo
-
-
-// EXIST: { itemText:"this" }

idea/idea-completion/testData/smart/smartCasts/SmartCastThisType1.kt

@@ -0,0 +1,13 @@
+open class Foo{
+    fun Any.f() {
+        if (this@Foo is Bar && this is Bar){
+            var a: Bar = <caret>
+        }
+    }
+}
+
+class Bar : Foo
+
+
+// EXIST: { lookupString: "this" }
+// EXIST: { lookupString: "this@Foo" }

idea/idea-completion/testData/smart/smartCasts/SmartCastThisType2.kt

@@ -0,0 +1,16 @@
+open class Foo{
+    fun Bar.f() {
+        fun Any.g() {
+            if (this is Bar){
+                var a: Bar = this@<caret>
+            }
+        }
+    }
+}
+
+class Bar : Foo
+
+
+// EXIST: { lookupString: "this@g" }
+// EXIST: { lookupString: "this@f" }
+// ABSENT: { lookupString: "this@Foo" }

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmSmartCompletionTestGenerated.java

@@ -1240,6 +1240,18 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
             doTest(fileName);
         }
 
+        @TestMetadata("SmartCastThisType1.kt")
+        public void testSmartCastThisType1() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/smartCasts/SmartCastThisType1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("SmartCastThisType2.kt")
+        public void testSmartCastThisType2() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/smartCasts/SmartCastThisType2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("SmartCastType.kt")
         public void testSmartCastType() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/smart/smartCasts/SmartCastType.kt");

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
+        expression: JetExpression
-): (VariableDescriptor) -> Collection<JetType> {
+) {
+    // 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> {
+    fun types(descriptor: VariableDescriptor): Collection<JetType> {
-        return function(descriptor)
+        val type = descriptor.returnType ?: return emptyList()
+        return entityType(descriptor, type)
     }
 
-    private val function: (VariableDescriptor) -> Collection<JetType> = run {
+    fun types(thisReceiverParameter: ReceiverParameterDescriptor): Collection<JetType> {
-        val receiver = nameExpression.getReceiverExpression()
+        val type = thisReceiverParameter.type
-        val dataFlowInfo = bindingContext.getDataFlowInfo(nameExpression)
+        val thisReceiver = thisReceiverParameter.value as? ThisReceiver ?: return listOf(type)
-        val (variableToTypes, notNullVariables) = processDataFlowInfo(dataFlowInfo, receiver)
+        return entityType(thisReceiver, type)
+    }
 
-        fun typesOf(descriptor: VariableDescriptor): Collection<JetType> {
+    private fun entityType(entity: Any, ownType: JetType): Collection<JetType> {
-            var type = descriptor.getReturnType() ?: return listOf()
+        val smartCastInfo = entityToSmartCastInfo[entity] ?: return listOf(ownType)
-            if (notNullVariables.contains(descriptor)) {
-                type = type.makeNotNullable()
-            }
 
-            val smartCastTypes = variableToTypes[descriptor]
+        var types = smartCastInfo.types + ownType
-            if (smartCastTypes == null || smartCastTypes.isEmpty()) return type.singletonOrEmptyList()
+
-            return smartCastTypes + type.singletonOrEmptyList()
+        if (smartCastInfo.notNull) {
+            types = types.map { it.makeNotNullable() }
         }
 
-        ::typesOf
+        return types
     }
 
-    private data class ProcessDataFlowInfoResult(
+    private data class SmartCastInfo(var types: Collection<JetType>, var notNull: Boolean) {
-            val variableToTypes: Map<VariableDescriptor, Collection<JetType>> = Collections.emptyMap(),
+        constructor() : this(emptyList(), false)
-            val notNullVariables: Set<VariableDescriptor> = Collections.emptySet()
+    }
-    )
 
-    private fun processDataFlowInfo(dataFlowInfo: DataFlowInfo, receiver: JetExpression?): ProcessDataFlowInfoResult {
+    private fun processDataFlowInfo(dataFlowInfo: DataFlowInfo, receiver: JetExpression?): Map<Any, SmartCastInfo> {
-        if (dataFlowInfo == DataFlowInfo.EMPTY) return ProcessDataFlowInfoResult()
+        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 receiverType = bindingContext.getType(receiver) ?: return emptyMap()
-            val receiverId = DataFlowValueFactory.createDataFlowValue(receiver, receiverType, bindingContext, containingDeclarationOrModule).getId()
+            val receiverId = DataFlowValueFactory.createDataFlowValue(receiver, receiverType, bindingContext, containingDeclarationOrModule).id
-            dataFlowValueToVariable = { value ->
+            dataFlowValueToEntity = { value ->
-                val id = value.getId()
+                val id = value.id
                 if (id is Pair<*, *> && id.first == receiverId) id.second as? VariableDescriptor else null
             }
         }
         else {
-            dataFlowValueToVariable = { value ->
+            dataFlowValueToEntity = { value ->
-                val id = value.getId()
+                val id = value.id
-                when {
+                when(id) {
-                    id is VariableDescriptor -> id
+                    is VariableDescriptor, is ThisReceiver -> id
-                    id is Pair<*, *> && id.first is ThisReceiver -> id.second as? VariableDescriptor
+                    is Pair<*, *> -> if (id.first is ThisReceiver) id.second as? VariableDescriptor else null
                     else -> null
                 }
             }
         }
 
-        val variableToType = HashMap<VariableDescriptor, Collection<JetType>>()
+        val entityToInfo = HashMap<Any, SmartCastInfo>()
-        val typeInfo: SetMultimap<DataFlowValue, JetType> = dataFlowInfo.getCompleteTypeInfo()
+
-        for ((dataFlowValue, types) in typeInfo.asMap().entrySet()) {
+        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo.asMap().entrySet()) {
-            val variable = dataFlowValueToVariable.invoke(dataFlowValue)
+            val entity = dataFlowValueToEntity.invoke(dataFlowValue)
-            if (variable != null) {
+            if (entity != null) {
-                variableToType[variable] = types
+                entityToInfo[entity] = SmartCastInfo(types, false)
             }
         }
 
-        val nullabilityInfo: Map<DataFlowValue, Nullability> = dataFlowInfo.getCompleteNullabilityInfo()
+        for ((dataFlowValue, nullability) in dataFlowInfo.completeNullabilityInfo) {
-        val notNullVariables = nullabilityInfo
+            if (nullability == Nullability.NOT_NULL) {
-                .filter { it.getValue() == Nullability.NOT_NULL }
+                val entity = dataFlowValueToEntity(dataFlowValue) ?: continue
-                .map { dataFlowValueToVariable(it.getKey()) }
+                entityToInfo.getOrPut(entity, { SmartCastInfo() }).notNull = true
-                .filterNotNullTo(HashSet<VariableDescriptor>())
+            }
+        }
 
-        return ProcessDataFlowInfoResult(variableToType, notNullVariables)
+        return entityToInfo
     }
 }