KT-5079 Support smart completion for when values

#KT-5079 Fixed

idea/src/org/jetbrains/jet/plugin/completion/ExpectedInfos.kt

@@ -53,6 +53,9 @@ import org.jetbrains.jet.lang.descriptors.Visibilities
 import org.jetbrains.jet.lang.psi.JetBlockExpression
 import org.jetbrains.jet.plugin.util.makeNullable
 import org.jetbrains.jet.plugin.util.makeNotNullable
+import org.jetbrains.jet.lang.psi.JetWhenConditionWithExpression
+import org.jetbrains.jet.lang.psi.JetWhenEntry
+import org.jetbrains.jet.lang.psi.JetWhenExpression
 
 enum class Tail {
     COMMA
@@ -70,6 +73,7 @@ class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: Mo
             ?: calculateForIf(expressionWithType)
             ?: calculateForElvis(expressionWithType)
             ?: calculateForBlockExpression(expressionWithType)
+            ?: calculateForWhenEntryValue(expressionWithType)
             ?: getFromBindingContext(expressionWithType)
     }
 
@@ -219,6 +223,20 @@ class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: Mo
         return calculate(block)?.map { ExpectedInfo(it.`type`, null) }
     }
 
+    private fun calculateForWhenEntryValue(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
+        val condition = expressionWithType.getParent() as? JetWhenConditionWithExpression ?: return null
+        val entry = condition.getParent() as JetWhenEntry
+        val whenExpression = entry.getParent() as JetWhenExpression
+        val subject = whenExpression.getSubjectExpression()
+        if (subject != null) {
+            val subjectType = bindingContext[BindingContext.EXPRESSION_TYPE, subject] ?: return null
+            return listOf(ExpectedInfo(subjectType, null))
+        }
+        else {
+            return listOf(ExpectedInfo(KotlinBuiltIns.getInstance().getBooleanType(), null))
+        }
+    }
+
     private fun getFromBindingContext(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
         val expectedType = bindingContext[BindingContext.EXPECTED_EXPRESSION_TYPE, expressionWithType] ?: return null
         return listOf(ExpectedInfo(expectedType, null))

idea/src/org/jetbrains/jet/plugin/completion/smart/KeywordValues.kt

@@ -20,14 +20,39 @@ import com.intellij.codeInsight.lookup.LookupElement
 import org.jetbrains.jet.plugin.completion.ExpectedInfo
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns
 import com.intellij.codeInsight.lookup.LookupElementBuilder
+import org.jetbrains.jet.lang.psi.*
+import com.intellij.codeInsight.lookup.LookupElementDecorator
+import com.intellij.codeInsight.completion.InsertionContext
+import org.jetbrains.jet.plugin.completion.handlers.WithTailInsertHandler
 
 object KeywordValues {
-    public fun addToCollection(collection: MutableCollection<LookupElement>, expectedInfos: Collection<ExpectedInfo>) {
-        val booleanInfoClassifier = { (info: ExpectedInfo) ->
-            if (info.`type` == KotlinBuiltIns.getInstance().getBooleanType()) ExpectedInfoClassification.MATCHES else ExpectedInfoClassification.NOT_MATCHES
+    public fun addToCollection(collection: MutableCollection<LookupElement>, expectedInfos: Collection<ExpectedInfo>, expressionWithType: JetExpression) {
+        var skipTrueFalse = false
+
+        val whenCondition = expressionWithType.getParent() as? JetWhenConditionWithExpression
+        if (whenCondition != null) {
+            val entry = whenCondition.getParent() as JetWhenEntry
+            val whenExpression = entry.getParent() as JetWhenExpression
+            if (whenExpression.getElseExpression() == null && entry == whenExpression.getEntries().last) {
+                val lookupElement = LookupElementBuilder.create("else").bold().withTailText(" ->")
+                collection.add(object: LookupElementDecorator<LookupElement>(lookupElement) {
+                    override fun handleInsert(context: InsertionContext) {
+                        WithTailInsertHandler("->", spaceBefore = true, spaceAfter = true).handleInsert(context, getDelegate())
+                    }
+                })
+            }
+            if (whenExpression.getSubjectExpression() == null) { // no sense in true or false entries for when with no subject
+                skipTrueFalse = true
+            }
+        }
+
+        if (!skipTrueFalse) {
+            val booleanInfoClassifier = { (info: ExpectedInfo) ->
+                if (info.`type` == KotlinBuiltIns.getInstance().getBooleanType()) ExpectedInfoClassification.MATCHES else ExpectedInfoClassification.NOT_MATCHES
+            }
+            collection.addLookupElements(expectedInfos, booleanInfoClassifier, { LookupElementBuilder.create("true").bold() })
+            collection.addLookupElements(expectedInfos, booleanInfoClassifier, { LookupElementBuilder.create("false").bold() })
         }
-        collection.addLookupElements(expectedInfos, booleanInfoClassifier, { LookupElementBuilder.create("true").bold() })
-        collection.addLookupElements(expectedInfos, booleanInfoClassifier, { LookupElementBuilder.create("false").bold() })
 
         collection.addLookupElements(expectedInfos,
                                      { info -> if (info.`type`.isNullable()) ExpectedInfoClassification.MATCHES else ExpectedInfoClassification.NOT_MATCHES },

idea/src/org/jetbrains/jet/plugin/completion/smart/SmartCompletion.kt

@@ -106,25 +106,25 @@ class SmartCompletion(val expression: JetSimpleNameExpression,
         if (receiver == null) {
             TypeInstantiationItems(bindingContext, resolveSession, visibilityFilter).addToCollection(result, expectedInfos)
 
-            StaticMembers(bindingContext, resolveSession).addToCollection(result, expectedInfos, expression)
+            StaticMembers(bindingContext, resolveSession).addToCollection(result, expectedInfos, expression, itemsToSkip)
 
             ThisItems(bindingContext).addToCollection(result, expressionWithType, expectedInfos)
 
             LambdaItems.addToCollection(result, functionExpectedInfos)
 
-            KeywordValues.addToCollection(result, expectedInfos)
+            KeywordValues.addToCollection(result, expectedInfos, expressionWithType)
         }
 
         return result
     }
 
-    private fun calcItemsToSkip(expression: JetExpression): Collection<DeclarationDescriptor> {
+    private fun calcItemsToSkip(expression: JetExpression): Set<DeclarationDescriptor> {
         val parent = expression.getParent()
         when(parent) {
             is JetProperty -> {
                 //TODO: this can be filtered out by ordinary completion
                 if (expression == parent.getInitializer()) {
-                    return resolveSession.resolveToElement(parent)[BindingContext.DECLARATION_TO_DESCRIPTOR, parent].toList()
+                    return resolveSession.resolveToElement(parent)[BindingContext.DECLARATION_TO_DESCRIPTOR, parent].toSet()
                 }
             }
 
@@ -134,13 +134,36 @@ class SmartCompletion(val expression: JetSimpleNameExpression,
                     if (operationToken == JetTokens.EQ || operationToken == JetTokens.EQEQ || operationToken == JetTokens.EXCLEQ) {
                         val left = parent.getLeft()
                         if (left is JetReferenceExpression) {
-                            return resolveSession.resolveToElement(left)[BindingContext.REFERENCE_TARGET, left].toList()
+                            return resolveSession.resolveToElement(left)[BindingContext.REFERENCE_TARGET, left].toSet()
                         }
                     }
                 }
             }
+
+            is JetWhenConditionWithExpression -> {
+                val entry = parent.getParent() as JetWhenEntry
+                val whenExpression = entry.getParent() as JetWhenExpression
+                val subject = whenExpression.getSubjectExpression() ?: return setOf()
+                val subjectType = bindingContext[BindingContext.EXPRESSION_TYPE, subject] ?: return setOf()
+                val classDescriptor = TypeUtils.getClassDescriptor(subjectType)
+                if (classDescriptor != null && DescriptorUtils.isEnumClass(classDescriptor)) {
+                    val usedEnumEntries = HashSet<ClassDescriptor>()
+                    val conditions = whenExpression.getEntries()
+                            .flatMap { it.getConditions().toList() }
+                            .filterIsInstance(javaClass<JetWhenConditionWithExpression>())
+                    for (condition in conditions) {
+                        val selectorExpr = (condition.getExpression() as? JetDotQualifiedExpression)
+                                ?.getSelectorExpression() as? JetReferenceExpression ?: continue
+                        val target = bindingContext[BindingContext.REFERENCE_TARGET, selectorExpr] as? ClassDescriptor ?: continue
+                        if (DescriptorUtils.isEnumEntry(target)) {
+                            usedEnumEntries.add(target)
+                        }
+                    }
+                    return usedEnumEntries
+                }
+            }
         }
-        return listOf()
+        return setOf()
     }
 
     private fun toFunctionReferenceLookupElement(descriptor: DeclarationDescriptor,

idea/src/org/jetbrains/jet/plugin/completion/smart/StaticMembers.kt

@@ -40,14 +40,18 @@ import org.jetbrains.jet.plugin.util.makeNotNullable
 
 // adds java static members, enum members and members from class object
 class StaticMembers(val bindingContext: BindingContext, val resolveSession: ResolveSessionForBodies) {
-    public fun addToCollection(collection: MutableCollection<LookupElement>, expectedInfos: Collection<ExpectedInfo>, context: JetExpression) {
+    public fun addToCollection(collection: MutableCollection<LookupElement>,
+                               expectedInfos: Collection<ExpectedInfo>,
+                               context: JetExpression,
+                               enumEntriesToSkip: Set<DeclarationDescriptor>) {
+
         val scope = bindingContext[BindingContext.RESOLUTION_SCOPE, context]
         if (scope == null) return
 
         val expectedInfosByClass = expectedInfos.groupBy { TypeUtils.getClassDescriptor(it.`type`) }
         for ((classDescriptor, expectedInfosForClass) in expectedInfosByClass) {
             if (classDescriptor != null && !classDescriptor.getName().isSpecial()) {
-                addToCollection(collection, classDescriptor, expectedInfosForClass, scope)
+                addToCollection(collection, classDescriptor, expectedInfosForClass, scope, enumEntriesToSkip)
             }
         }
     }
@@ -56,7 +60,8 @@ class StaticMembers(val bindingContext: BindingContext, val resolveSession: Reso
             collection: MutableCollection<LookupElement>,
             classDescriptor: ClassDescriptor,
             expectedInfos: Collection<ExpectedInfo>,
-            scope: JetScope) {
+            scope: JetScope,
+            enumEntriesToSkip: Set<DeclarationDescriptor>) {
 
         fun processMember(descriptor: DeclarationDescriptor) {
             if (descriptor is DeclarationDescriptorWithVisibility && !Visibilities.isVisible(descriptor, scope.getContainingDeclaration())) return
@@ -74,7 +79,7 @@ class StaticMembers(val bindingContext: BindingContext, val resolveSession: Reso
                         }
                 }
             }
-            else if (DescriptorUtils.isEnumEntry(descriptor)) {
+            else if (DescriptorUtils.isEnumEntry(descriptor) && !enumEntriesToSkip.contains(descriptor)) {
                 classifier = { ExpectedInfoClassification.MATCHES } /* we do not need to check type of enum entry because it's taken from proper enum */
             }
             else {

idea/src/org/jetbrains/jet/plugin/completion/smart/Utils.kt

@@ -175,5 +175,6 @@ fun createLookupElement(descriptor: DeclarationDescriptor, resolveSession: Resol
 fun JetType.isSubtypeOf(expectedType: JetType) = !isError() && JetTypeChecker.INSTANCE.isSubtypeOf(this, expectedType)
 
 fun <T : Any> T?.toList(): List<T> = if (this != null) listOf(this) else listOf()
+fun <T : Any> T?.toSet(): Set<T> = if (this != null) setOf(this) else setOf()
 
 fun String?.isNullOrEmpty() = this == null || this.isEmpty()

idea/testData/completion/handlers/smart/WhenElse.kt

@@ -0,0 +1,7 @@
+fun foo(s: String) {
+    when(s) {
+        <caret>
+    }
+}
+
+// ELEMENT: else

idea/testData/completion/handlers/smart/WhenElse.kt.after

@@ -0,0 +1,7 @@
+fun foo(s: String) {
+    when(s) {
+        else -> <caret>
+    }
+}
+
+// ELEMENT: else

idea/testData/completion/smart/WhenEntryValue1.kt

@@ -0,0 +1,16 @@
+enum class E {
+    A
+    B
+    C
+}
+
+fun foo(e: E) {
+    when(e) {
+        <caret>
+    }
+}
+
+// EXIST: E.A
+// EXIST: E.B
+// EXIST: E.C
+// EXIST: {"lookupString":"else","tailText":" ->","itemText":"else"}

idea/testData/completion/smart/WhenEntryValue2.kt

@@ -0,0 +1,16 @@
+enum class E {
+    A
+    B
+    C
+}
+
+fun foo(e: E) {
+    when(e) {
+        E.A, <caret>
+    }
+}
+
+// ABSENT: E.A
+// EXIST: E.B
+// EXIST: E.C
+// ABSENT:else

idea/testData/completion/smart/WhenEntryValue3.kt

@@ -0,0 +1,17 @@
+enum class E {
+  A
+  B
+  C
+}
+
+fun foo(e: E) {
+    when(e) {
+        E.A -> x()
+        <caret>
+    }
+}
+
+// ABSENT: E.A
+// EXIST: E.B
+// EXIST: E.C
+// EXIST: {"lookupString":"else","tailText":" ->","itemText":"else"}

idea/testData/completion/smart/WhenEntryValue4.kt

@@ -0,0 +1,17 @@
+enum class E {
+    A
+    B
+    C
+}
+
+fun foo(e: E) {
+    when(e) {
+        <caret>
+        else -> x()
+    }
+}
+
+// EXIST: E.A
+// EXIST: E.B
+// EXIST: E.C
+// ABSENT: else

idea/testData/completion/smart/WhenEntryValue5.kt

@@ -0,0 +1,17 @@
+enum class E {
+    A
+    B
+    C
+}
+
+fun foo(e: E) {
+    when(e) {
+        E.A -> x()
+        E.<caret>
+    }
+}
+
+// ABSENT: A
+// EXIST: B
+// EXIST: C
+// ABSENT: else

idea/testData/completion/smart/WhenEntryValue6.kt

@@ -0,0 +1,18 @@
+enum class E {
+  A
+  B
+  C
+}
+
+fun foo(e: E?) {
+    when(e) {
+        E.A -> x()
+        <caret>
+    }
+}
+
+// ABSENT: E.A
+// EXIST: E.B
+// EXIST: E.C
+// EXIST: null
+// EXIST: {"lookupString":"else","tailText":" ->","itemText":"else"}

idea/testData/completion/smart/WhenEntryValue7.kt

@@ -0,0 +1,15 @@
+enum class E {
+  A
+  B
+  C
+}
+
+fun foo(e: E) {
+    when(e) {
+        <caret>
+        E.B -> x()
+    }
+}
+
+// EXIST: E.A
+// ABSENT: else

idea/testData/completion/smart/WhenEntryValue8.kt

@@ -0,0 +1,11 @@
+import java.lang.annotation.ElementType
+
+fun foo(e: ElementType) {
+    when(e) {
+        ElementType.FIELD -> x()
+        <caret>
+    }
+}
+
+// ABSENT: ElementType.FIELD
+// EXIST: ElementType.TYPE

idea/testData/completion/smart/WhenWithNoSubjectEntryValue1.kt

@@ -0,0 +1,11 @@
+fun foo(b1: Boolean, b2: Boolean) {
+    when {
+        <caret>
+    }
+}
+
+// EXIST: b1
+// EXIST: b2
+// ABSENT: true
+// ABSENT: false
+// EXIST: {"lookupString":"else","tailText":" ->","itemText":"else"}

idea/testData/completion/smart/WhenWithNoSubjectEntryValue2.kt

@@ -0,0 +1,8 @@
+fun foo(s: String) {
+    when {
+        s.<caret>
+    }
+}
+
+// EXIST: equals
+// ABSENT: else

idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java

@@ -476,6 +476,56 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest("idea/testData/completion/smart/VariableInitializer.kt");
     }
     
+    @TestMetadata("WhenEntryValue1.kt")
+    public void testWhenEntryValue1() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue1.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue2.kt")
+    public void testWhenEntryValue2() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue2.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue3.kt")
+    public void testWhenEntryValue3() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue3.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue4.kt")
+    public void testWhenEntryValue4() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue4.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue5.kt")
+    public void testWhenEntryValue5() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue5.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue6.kt")
+    public void testWhenEntryValue6() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue6.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue7.kt")
+    public void testWhenEntryValue7() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue7.kt");
+    }
+    
+    @TestMetadata("WhenEntryValue8.kt")
+    public void testWhenEntryValue8() throws Exception {
+        doTest("idea/testData/completion/smart/WhenEntryValue8.kt");
+    }
+    
+    @TestMetadata("WhenWithNoSubjectEntryValue1.kt")
+    public void testWhenWithNoSubjectEntryValue1() throws Exception {
+        doTest("idea/testData/completion/smart/WhenWithNoSubjectEntryValue1.kt");
+    }
+    
+    @TestMetadata("WhenWithNoSubjectEntryValue2.kt")
+    public void testWhenWithNoSubjectEntryValue2() throws Exception {
+        doTest("idea/testData/completion/smart/WhenWithNoSubjectEntryValue2.kt");
+    }
+    
     @TestMetadata("WithPrefix.kt")
     public void testWithPrefix() throws Exception {
         doTest("idea/testData/completion/smart/WithPrefix.kt");

idea/tests/org/jetbrains/jet/completion/handlers/SmartCompletionHandlerTestGenerated.java

@@ -401,4 +401,9 @@ public class SmartCompletionHandlerTestGenerated extends AbstractSmartCompletion
         doTest("idea/testData/completion/handlers/smart/True2.kt");
     }
     
+    @TestMetadata("WhenElse.kt")
+    public void testWhenElse() throws Exception {
+        doTest("idea/testData/completion/handlers/smart/WhenElse.kt");
+    }
+    
 }