java completion: prefer expected enum constants after == (IDEA-233490)

GitOrigin-RevId: 6db7e389ccd575d1d18218f3c5a7352a05ac5c0c

java/java-impl/src/com/intellij/codeInsight/completion/JavaCompletionStatistician.java

@@ -23,7 +23,6 @@ import com.intellij.psi.*;
 import com.intellij.psi.statistics.JavaStatisticsManager;
 import com.intellij.psi.statistics.StatisticsInfo;
 import com.intellij.psi.util.PsiTreeUtil;
-import com.intellij.psi.util.PsiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -51,7 +50,9 @@ public class JavaCompletionStatistician extends CompletionStatistician{
     }
 
     PsiElement position = location.getCompletionParameters().getPosition();
-    if (SUPER_CALL.accepts(position) || JavaCompletionContributor.IN_SWITCH_LABEL.accepts(position)) {
+    if (SUPER_CALL.accepts(position) ||
+        JavaCompletionContributor.IN_SWITCH_LABEL.accepts(position) ||
+        PreferByKindWeigher.isComparisonRhs(position)) {
       return StatisticsInfo.EMPTY;
     }
 
@@ -67,10 +68,7 @@ public class JavaCompletionStatistician extends CompletionStatistician{
   }
 
   private static boolean isInEnumAnnotationParameter(PsiElement position, ExpectedTypeInfo firstInfo) {
-    if (PsiTreeUtil.getParentOfType(position, PsiNameValuePair.class) == null) return false;
-    
-    PsiClass expectedClass = PsiUtil.resolveClassInType(firstInfo.getType());
-    return expectedClass != null && expectedClass.isEnum();
+    return PsiTreeUtil.getParentOfType(position, PsiNameValuePair.class) != null && PreferByKindWeigher.isEnumClass(firstInfo);
   }
 
   @Nullable

java/java-impl/src/com/intellij/codeInsight/completion/PreferByKindWeigher.java

@@ -147,6 +147,11 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     return condition ? MyResult.suitableClass : MyResult.classNameOrGlobalStatic;
   }
 
+  static boolean isEnumClass(@NotNull ExpectedTypeInfo info) {
+    PsiClass expectedClass = PsiUtil.resolveClassInType(info.getType());
+    return expectedClass != null && expectedClass.isEnum();
+  }
+
   enum MyResult {
     annoMethod,
     probableKeyword,
@@ -171,7 +176,7 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     nonInitialized,
     classNameOrGlobalStatic,
     introducedVariable,
-    unlikelyClass,
+    unlikelyItem,
     improbableKeyword,
   }
 
@@ -194,7 +199,9 @@ public class PreferByKindWeigher extends LookupElementWeigher {
         object instanceof PsiThisExpression ||
         object instanceof PsiField && !((PsiField)object).hasModifierProperty(PsiModifier.STATIC)) {
       if (PsiTreeUtil.getParentOfType(myPosition, PsiDocComment.class) == null) {
-        return isExpectedTypeItem(item) ? MyResult.expectedTypeVariable : MyResult.variable;
+        return isComparisonWithItself((PsiElement)object) ? MyResult.unlikelyItem :
+               isExpectedTypeItem(item) ? MyResult.expectedTypeVariable :
+               MyResult.variable;
       }
     }
 
@@ -218,7 +225,7 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     if (object instanceof PsiClass &&
         CommonClassNames.JAVA_LANG_STRING.equals(((PsiClass)object).getQualifiedName()) &&
         JavaSmartCompletionContributor.AFTER_NEW.accepts(myPosition)) {
-      return MyResult.unlikelyClass;
+      return MyResult.unlikelyItem;
     }
     Boolean expectedTypeMember = item.getUserData(MembersGetter.EXPECTED_TYPE_MEMBER);
     if (expectedTypeMember != null) {
@@ -281,6 +288,16 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     return MyResult.normal;
   }
 
+  private boolean isComparisonWithItself(PsiElement itemObject) {
+    if (isComparisonRhs(myPosition) && myPosition.getParent().getParent() instanceof PsiPolyadicExpression) {
+      PsiExpression[] operands = ((PsiPolyadicExpression)myPosition.getParent().getParent()).getOperands();
+      if (operands[0] instanceof PsiReferenceExpression && ((PsiReferenceExpression)operands[0]).resolve() == itemObject) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private boolean isExpectedTypeItem(@NotNull LookupElement item) {
     TypedLookupItem typed = item.as(TypedLookupItem.CLASS_CONDITION_KEY);
     PsiType itemType = typed == null ? null : typed.getType();
@@ -325,8 +342,9 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     if (PsiKeyword.INTERFACE.equals(keyword) && psiElement().afterLeaf("@").accepts(myPosition)) {
       return ThreeState.NO;
     }
-    if (PsiKeyword.NULL.equals(keyword) && psiElement().afterLeaf(psiElement().withElementType(elementType().oneOf(JavaTokenType.EQEQ, JavaTokenType.NE))).accepts(myPosition)) {
-      return ThreeState.YES;
+    if (PsiKeyword.NULL.equals(keyword) && isComparisonRhs(myPosition)) {
+      boolean expectsNotNull = Arrays.stream(myExpectedTypes).anyMatch(PreferByKindWeigher::isEnumClass);
+      return expectsNotNull ? ThreeState.NO : ThreeState.YES;
     }
     if (JavaKeywordCompletion.PRIMITIVE_TYPES.contains(keyword) || PsiKeyword.VOID.equals(keyword)) {
       boolean inCallArg = psiElement().withParents(PsiReferenceExpression.class, PsiExpressionList.class).accepts(myPosition);
@@ -338,6 +356,10 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     return ThreeState.UNSURE;
   }
 
+  static boolean isComparisonRhs(PsiElement position) {
+    return psiElement().afterLeaf(psiElement().withElementType(elementType().oneOf(JavaTokenType.EQEQ, JavaTokenType.NE))).accepts(position);
+  }
+
   private boolean isBeforeVariableOnSameLine(@Nullable PsiStatement parentStatement) {
     return parentStatement != null &&
            parentStatement.getTextRange().getStartOffset() == myPosition.getTextRange().getStartOffset() &&

java/java-tests/testData/codeInsight/completion/normalSorting/PreferExpectedEnumConstantsInComparison.java

@@ -0,0 +1,7 @@
+class Foo {
+  boolean foo(MyEnum e) {
+    return e == <caret>
+  }
+}
+
+enum MyEnum { const1, const2 }

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/NormalCompletionOrderingTest.groovy

@@ -270,6 +270,12 @@ class NormalCompletionOrderingTest extends CompletionSortingTestCase {
     assertPreferredItems 0, 'constx1', 'constx2', 'const1', 'const2'
   }
 
+  void testPreferExpectedEnumConstantsInComparison() {
+    checkPreferredItems 0, 'MyEnum.const1', 'MyEnum', 'MyEnum.const2'
+    incUseCount(lookup, myFixture.lookupElementStrings.indexOf('String')) // select some unrelated class
+    assertPreferredItems 0, 'MyEnum.const1', 'MyEnum', 'MyEnum.const2'
+  }
+
   void testPreferElse() {
     checkPreferredItems(0, "else", "element")
   }