Smart completion after == and !=

#KT-4912 Fixed

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

@@ -40,6 +40,8 @@ import java.util.HashSet
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall
 import org.jetbrains.jet.lang.descriptors.ModuleDescriptor
 import org.jetbrains.jet.lang.types.JetType
+import org.jetbrains.jet.lang.psi.JetBinaryExpression
+import org.jetbrains.jet.lexer.JetTokens
 
 enum class Tail {
     COMMA
@@ -50,14 +52,10 @@ data class ExpectedInfo(val `type`: JetType, val tail: Tail?)
 
 class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: ModuleDescriptor) {
     public fun calculate(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
-        val expectedInfos1 = calculateForArgument(expressionWithType)
-        if (expectedInfos1 != null) return expectedInfos1
-
-        val expectedInfos2 = calculateForFunctionLiteralArgument(expressionWithType)
-        if (expectedInfos2 != null) return expectedInfos2
-
-        val expectedType = bindingContext[BindingContext.EXPECTED_EXPRESSION_TYPE, expressionWithType] ?: return null
-        return listOf(ExpectedInfo(expectedType, null))
+        return calculateForArgument(expressionWithType)
+            ?: calculateForFunctionLiteralArgument(expressionWithType)
+            ?: calculateForEq(expressionWithType)
+            ?: getFromBindingContext(expressionWithType)
     }
 
     private fun calculateForArgument(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
@@ -128,4 +126,24 @@ class ExpectedInfos(val bindingContext: BindingContext, val moduleDescriptor: Mo
         }
         return expectedInfos
     }
+
+    private fun calculateForEq(expressionWithType: JetExpression): Collection<ExpectedInfo>? {
+        val binaryExpression = expressionWithType.getParent() as? JetBinaryExpression
+        if (binaryExpression != null) {
+            val operationToken = binaryExpression.getOperationToken()
+            if (operationToken == JetTokens.EQEQ || operationToken == JetTokens.EXCLEQ) {
+                val otherOperand = if (expressionWithType == binaryExpression.getRight()) binaryExpression.getLeft() else binaryExpression.getRight()
+                if (otherOperand != null) {
+                    val expressionType = bindingContext[BindingContext.EXPRESSION_TYPE, otherOperand] ?: return null
+                    return listOf(ExpectedInfo(TypeUtils.makeNullable(expressionType), null))
+                }
+            }
+        }
+        return 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/SmartCompletion.kt

@@ -104,10 +104,13 @@ class SmartCompletion(val expression: JetSimpleNameExpression,
             }
 
             is JetBinaryExpression -> {
-                if (parent.getRight() == expression && parent.getOperationToken() == JetTokens.EQ) {
-                    val left = parent.getLeft()
-                    if (left is JetReferenceExpression) {
-                        return resolveSession.resolveToElement(left)[BindingContext.REFERENCE_TARGET, left].toList()
+                if (parent.getRight() == expression) {
+                    val operationToken = parent.getOperationToken()
+                    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()
+                        }
                     }
                 }
             }

idea/testData/completion/smart/EqOperator.kt

@@ -0,0 +1,14 @@
+enum class E {
+    A
+    B
+}
+
+fun f(e1: E, e2: E?, x: Any) {
+    if (e1 == <caret>
+}
+
+// EXIST: E.A
+// EXIST: E.B
+// EXIST: e2
+// ABSENT: e1
+// ABSENT: x

idea/testData/completion/smart/NotEqOperator.kt

@@ -0,0 +1,14 @@
+enum class E {
+    A
+    B
+}
+
+fun f(e1: E, e2: E?, x: Any) {
+    if (e1 != <caret>
+}
+
+// EXIST: E.A
+// EXIST: E.B
+// EXIST: e2
+// ABSENT: e1
+// ABSENT: x

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

@@ -131,6 +131,11 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest("idea/testData/completion/smart/EnumMembers.kt");
     }
     
+    @TestMetadata("EqOperator.kt")
+    public void testEqOperator() throws Exception {
+        doTest("idea/testData/completion/smart/EqOperator.kt");
+    }
+    
     @TestMetadata("FunctionReference1.kt")
     public void testFunctionReference1() throws Exception {
         doTest("idea/testData/completion/smart/FunctionReference1.kt");
@@ -251,6 +256,11 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest("idea/testData/completion/smart/NoSillyAssignment.kt");
     }
     
+    @TestMetadata("NotEqOperator.kt")
+    public void testNotEqOperator() throws Exception {
+        doTest("idea/testData/completion/smart/NotEqOperator.kt");
+    }
+    
     @TestMetadata("NotSillyAssignment.kt")
     public void testNotSillyAssignment() throws Exception {
         doTest("idea/testData/completion/smart/NotSillyAssignment.kt");