Moving logic into ExpectedInfos

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

@@ -51,6 +51,7 @@ import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils
 import org.jetbrains.kotlin.types.typeUtil.containsError
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
+import org.jetbrains.kotlin.types.typeUtil.makeNullable
 import org.jetbrains.kotlin.utils.addToStdlib.check
 import org.jetbrains.kotlin.utils.addToStdlib.singletonOrEmptyList
 import java.util.LinkedHashSet
@@ -138,20 +139,22 @@ class ExpectedInfos(
         val useOuterCallsExpectedTypeCount: Int = 0
 ) {
     public fun calculate(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
-        return calculateForArgument(expressionWithType)
-            ?: calculateForFunctionLiteralArgument(expressionWithType)
-            ?: calculateForEqAndAssignment(expressionWithType)
-            ?: calculateForIf(expressionWithType)
-            ?: calculateForElvis(expressionWithType)
-            ?: calculateForBlockExpression(expressionWithType)
-            ?: calculateForWhenEntryValue(expressionWithType)
-            ?: calculateForExclOperand(expressionWithType)
-            ?: calculateForInitializer(expressionWithType)
-            ?: calculateForExpressionBody(expressionWithType)
-            ?: calculateForReturn(expressionWithType)
-            ?: calculateForLoopRange(expressionWithType)
-            ?: calculateForInOperatorArgument(expressionWithType)
-            ?: getFromBindingContext(expressionWithType)
+        val expectedInfos = calculateForArgument(expressionWithType)
+                            ?: calculateForFunctionLiteralArgument(expressionWithType)
+                            ?: calculateForEqAndAssignment(expressionWithType)
+                            ?: calculateForIf(expressionWithType)
+                            ?: calculateForElvis(expressionWithType)
+                            ?: calculateForBlockExpression(expressionWithType)
+                            ?: calculateForWhenEntryValue(expressionWithType)
+                            ?: calculateForExclOperand(expressionWithType)
+                            ?: calculateForInitializer(expressionWithType)
+                            ?: calculateForExpressionBody(expressionWithType)
+                            ?: calculateForReturn(expressionWithType)
+                            ?: calculateForLoopRange(expressionWithType)
+                            ?: calculateForInOperatorArgument(expressionWithType)
+                            ?: getFromBindingContext(expressionWithType)
+                            ?: return null
+        return expectedInfos.filterNot { it.fuzzyType?.type?.isError ?: false }
     }
 
     private fun calculateForArgument(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
@@ -328,12 +331,17 @@ class ExpectedInfos(
         val binaryExpression = expressionWithType.getParent() as? JetBinaryExpression
         if (binaryExpression != null) {
             val operationToken = binaryExpression.getOperationToken()
-            if (operationToken == JetTokens.EQ || operationToken == JetTokens.EQEQ || operationToken == JetTokens.EXCLEQ
-                || operationToken == JetTokens.EQEQEQ || operationToken == JetTokens.EXCLEQEQEQ) {
+            if (operationToken == JetTokens.EQ || operationToken in COMPARISON_TOKENS) {
                 val otherOperand = if (expressionWithType == binaryExpression.getRight()) binaryExpression.getLeft() else binaryExpression.getRight()
                 if (otherOperand != null) {
-                    val expressionType = bindingContext.getType(otherOperand) ?: return null
-                    return listOf(ExpectedInfo(expressionType, expectedNameFromExpression(otherOperand), null))
+                    var expectedType = bindingContext.getType(otherOperand) ?: return null
+
+                    // if we complete argument of == or !=, make types in expected info's nullable to allow nullable items too
+                    if (operationToken in COMPARISON_TOKENS) {
+                        expectedType = expectedType.makeNullable()
+                    }
+
+                    return listOf(ExpectedInfo(expectedType, expectedNameFromExpression(otherOperand), null))
                 }
             }
         }
@@ -524,3 +532,6 @@ class ExpectedInfos(
     private fun String.unpluralize()
             = StringUtil.unpluralize(this)
 }
+
+val COMPARISON_TOKENS = setOf(JetTokens.EQEQ, JetTokens.EXCLEQ, JetTokens.EQEQEQ, JetTokens.EXCLEQEQEQ)
+

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

@@ -21,13 +21,11 @@ import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.codeInsight.lookup.LookupElementDecorator
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.idea.completion.COMPARISON_TOKENS
 import org.jetbrains.kotlin.idea.completion.ExpectedInfo
 import org.jetbrains.kotlin.idea.completion.fuzzyType
 import org.jetbrains.kotlin.idea.completion.handlers.WithTailInsertHandler
-import org.jetbrains.kotlin.psi.JetExpression
-import org.jetbrains.kotlin.psi.JetWhenConditionWithExpression
-import org.jetbrains.kotlin.psi.JetWhenEntry
-import org.jetbrains.kotlin.psi.JetWhenExpression
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.types.TypeSubstitutor
 
 object KeywordValues {
@@ -60,14 +58,21 @@ object KeywordValues {
             collection.addLookupElements(null, expectedInfos, booleanInfoClassifier) { LookupElementBuilder.create("false").bold().assignSmartCompletionPriority(SmartCompletionItemPriority.FALSE) }
         }
 
-        val classifier = { info: ExpectedInfo ->
-            if (info.fuzzyType != null && info.fuzzyType!!.type.isMarkedNullable())
-                ExpectedInfoClassification.match(TypeSubstitutor.EMPTY)
-            else
-                ExpectedInfoClassification.noMatch
-        }
-        collection.addLookupElements(null, expectedInfos, classifier) {
-            LookupElementBuilder.create("null").bold().assignSmartCompletionPriority(SmartCompletionItemPriority.NULL)
+        if (!shouldSkipNull(expressionWithType)) {
+            val classifier = { info: ExpectedInfo ->
+                if (info.fuzzyType != null && info.fuzzyType!!.type.isMarkedNullable())
+                    ExpectedInfoClassification.match(TypeSubstitutor.EMPTY)
+                else
+                    ExpectedInfoClassification.noMatch
+            }
+            collection.addLookupElements(null, expectedInfos, classifier) {
+                LookupElementBuilder.create("null").bold().assignSmartCompletionPriority(SmartCompletionItemPriority.NULL)
+            }
         }
     }
+
+    private fun shouldSkipNull(expressionWithType: JetExpression): Boolean {
+        val binaryExpression = expressionWithType.parent as? JetBinaryExpression ?: return false
+        return binaryExpression.operationToken in COMPARISON_TOKENS
+    }
 }

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

@@ -30,10 +30,8 @@ import org.jetbrains.kotlin.descriptors.*
 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.completion.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
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
@@ -113,14 +111,7 @@ class SmartCompletion(
 
         val expressionWithType = expression.toExpressionWithType()
 
-        var originalExpectedInfos = calcExpectedInfos(expressionWithType) ?: return null
-        originalExpectedInfos = originalExpectedInfos.filterNot { it.fuzzyType?.type?.isError ?: false }
-
-        // if we complete argument of == or !=, make types in expected info's nullable to allow nullable items too
-        val expectedInfos = if ((expressionWithType.getParent() as? JetBinaryExpression)?.getOperationToken() in COMPARISON_TOKENS)
-            originalExpectedInfos.map { if (it.fuzzyType != null) ExpectedInfo(it.fuzzyType!!.makeNullable(), it.expectedName, it.tail) else it }
-        else
-            originalExpectedInfos
+        val expectedInfos = calcExpectedInfos(expressionWithType) ?: return null
 
         val smartCastCalculator = SmartCastCalculator(bindingContext, moduleDescriptor, expression)
 
@@ -166,7 +157,7 @@ class SmartCompletion(
 
             LambdaItems.addToCollection(additionalItems, functionExpectedInfos)
 
-            KeywordValues.addToCollection(additionalItems, originalExpectedInfos/* use originalExpectedInfos to not include null after == */, expression)
+            KeywordValues.addToCollection(additionalItems, expectedInfos, expression)
 
             MultipleArgumentsItemProvider(bindingContext, smartCastCalculator).addToCollection(additionalItems, expectedInfos, expression)
         }
@@ -349,7 +340,5 @@ class SmartCompletion(
     companion object {
         public val OLD_ARGUMENTS_REPLACEMENT_OFFSET: OffsetKey = OffsetKey.create("nonFunctionReplacementOffset")
         public val MULTIPLE_ARGUMENTS_REPLACEMENT_OFFSET: OffsetKey = OffsetKey.create("multipleArgumentsReplacementOffset")
-
-        private val COMPARISON_TOKENS = setOf(JetTokens.EQEQ, JetTokens.EXCLEQ, JetTokens.EQEQEQ, JetTokens.EXCLEQEQEQ)
     }
 }