[ml-completion] introduce "incorrect_element" feature; add test cases

GitOrigin-RevId: 4e101b12a631fde7078161a00a1292e60a6bf00e

java/java-impl/intellij.java.impl.iml

@@ -84,6 +84,7 @@
     <orderEntry type="module" module-name="intellij.platform.core.ui" />
     <orderEntry type="module" module-name="intellij.platform.codeStyle.impl" />
     <orderEntry type="module" module-name="intellij.platform.ide.util.io" />
+    <orderEntry type="module" module-name="intellij.completionMlRanking" scope="TEST" />
   </component>
   <component name="copyright">
     <Base>

java/java-impl/src/com/intellij/codeInsight/completion/JavaIncorrectElements.kt

@@ -0,0 +1,199 @@
+// Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.codeInsight.completion
+
+import com.intellij.codeInsight.AnnotationTargetUtil
+import com.intellij.codeInsight.ExceptionUtil
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.UserDataHolder
+import com.intellij.patterns.ElementPattern
+import com.intellij.patterns.PsiJavaPatterns
+import com.intellij.psi.*
+import com.intellij.psi.impl.source.tree.JavaElementType
+import com.intellij.psi.util.InheritanceUtil
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.containers.addIfNotNull
+
+
+interface LookupPositionMatcher {
+  fun match(position: PsiElement): Boolean
+  fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean
+}
+
+object JavaIncorrectElements {
+  private val matcherKey = Key<(LookupElement) -> Boolean>("JavaIncorrectElements.matcher")
+  private val positions = listOf(
+    ExceptionPositionMatcher, TryWithResourcesPositionMatcher, AnnotationPositionMatcher, TypeParameterPositionMatcher,
+    ImplementsDeclarationPositionMatcher, ExtendsDeclarationPositionMatcher
+  )
+
+  fun matchPosition(position: PsiElement): LookupPositionMatcher? {
+    return positions.firstOrNull { it.match(position) }
+  }
+
+  fun putMatcher(elementMatcher: (LookupElement) -> Boolean, context: UserDataHolder) {
+    context.putUserData(matcherKey, elementMatcher)
+  }
+
+  fun tryGetMatcher(context: UserDataHolder): ((LookupElement) -> Boolean)? {
+    return context.getUserData(matcherKey)
+  }
+}
+
+private object AnnotationPositionMatcher: LookupPositionMatcher {
+  private fun tryGetAnnotation(position: PsiElement): PsiAnnotation? {
+    val parent = position.parent as? PsiJavaCodeReferenceElement ?: return null
+    return parent.parent as? PsiAnnotation
+  }
+
+  override fun match(position: PsiElement): Boolean {
+    return tryGetAnnotation(position) != null
+  }
+
+  override fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean {
+    val annotation = tryGetAnnotation(position) ?: throw AssertionError("Annotation is null")
+    val targets = AnnotationTargetUtil.getTargetsForLocation(annotation.owner)
+    return l@ { element ->
+      val psiClass = element.`object` as? PsiClass ?: return@l false
+      return@l !psiClass.isAnnotationType || AnnotationTargetUtil.findAnnotationTarget(psiClass, *targets) == null
+    }
+  }
+}
+
+private object TypeParameterPositionMatcher: LookupPositionMatcher {
+  private fun tryGetTypeElement(position: PsiElement): PsiTypeElement? {
+    val parent = position.parent as? PsiJavaCodeReferenceElement ?: return null
+    return parent.parent as? PsiTypeElement
+  }
+  
+  override fun match(position: PsiElement): Boolean {
+    return tryGetTypeElement(position) != null
+  }
+
+  override fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean {
+    val typeElement = tryGetTypeElement(position)!!
+    val bounds = PreferByKindWeigher.getTypeBounds(typeElement)
+    return l@ { element ->
+      val obj = element.`object`
+      if (obj is PsiKeyword) return@l true
+      val psiClass = obj as? PsiClass ?: return@l false
+      return@l bounds.all { !InheritanceUtil.isInheritorOrSelf(psiClass, it, true) }
+    }
+  }
+}
+
+private object TryWithResourcesPositionMatcher: LookupPositionMatcher {
+  override fun match(position: PsiElement): Boolean {
+    return PreferByKindWeigher.IN_RESOURCE.accepts(position)
+  }
+
+  override fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean {
+    return this::match
+  }
+
+  private fun match(lookupElement: LookupElement): Boolean {
+    val obj = lookupElement.`object`
+    if (obj is PsiKeyword && obj.text in JavaKeywordCompletion.PRIMITIVE_TYPES) {
+      return true
+    }
+    val psiClass = obj as? PsiClass ?: return false
+    return !InheritanceUtil.isInheritor(psiClass, CommonClassNames.JAVA_LANG_AUTO_CLOSEABLE)
+  }
+}
+
+private object ExceptionPositionMatcher: LookupPositionMatcher {
+  private fun isCatchClausePosition(position: PsiElement): Boolean {
+    return PreferByKindWeigher.IN_CATCH_TYPE.accepts(position) || PreferByKindWeigher.IN_MULTI_CATCH_TYPE.accepts(position)
+  }
+
+  override fun match(position: PsiElement): Boolean {
+    return isCatchClausePosition(position) ||
+           PreferByKindWeigher.INSIDE_METHOD_THROWS_CLAUSE.accepts(position) ||
+           JavaDocCompletionContributor.THROWS_TAG_EXCEPTION.accepts(position) ||
+           JavaSmartCompletionContributor.AFTER_THROW_NEW.accepts(position)
+  }
+
+  override fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean {
+    val thrownCheckedExceptions = mutableListOf<PsiClass>()
+    val isCatchClause = isCatchClausePosition(position)
+    if (isCatchClause) {
+      val container = PsiTreeUtil.getParentOfType(position, PsiTryStatement::class.java, PsiMethod::class.java)
+      if (container != null) {
+        val block = if (container is PsiTryStatement) container.tryBlock else container
+        if (block != null) {
+          for (type in ExceptionUtil.getThrownCheckedExceptions(block)) {
+            thrownCheckedExceptions.addIfNotNull(type.resolve())
+          }
+        }
+      }
+    }
+
+    return l@ { element ->
+      val psiClass = element.`object` as? PsiClass ?: return@l false
+
+      if (!InheritanceUtil.isInheritor(psiClass, CommonClassNames.JAVA_LANG_THROWABLE)) {
+        return@l true
+      }
+
+      val psiManager = psiClass.manager
+
+      // if exception is checked
+      if (isCatchClause) {
+        val qualifiedName = psiClass.qualifiedName
+        return@l !ExceptionUtil.isUncheckedException(psiClass)
+                 && qualifiedName != CommonClassNames.JAVA_LANG_THROWABLE
+                 && qualifiedName != CommonClassNames.JAVA_LANG_EXCEPTION
+                 && !thrownCheckedExceptions.any { psiManager.areElementsEquivalent(it, psiClass) }
+      }
+
+      return@l false
+    }
+  }
+}
+
+private object ImplementsDeclarationPositionMatcher: LookupPositionMatcher {
+  private val INSIDE_IMPLEMENTS_LIST: ElementPattern<PsiElement> = PsiJavaPatterns.psiElement().afterLeaf(
+    PsiKeyword.IMPLEMENTS, ",").inside(PsiJavaPatterns.psiElement(JavaElementType.IMPLEMENTS_LIST))
+
+  override fun match(position: PsiElement): Boolean {
+    return INSIDE_IMPLEMENTS_LIST.accepts(position)
+  }
+
+  override fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean {
+    return this::matchElement
+  }
+
+  private fun matchElement(element: LookupElement): Boolean {
+    val psiClass = element.`object` as? PsiClass ?: return false
+    return !psiClass.isInterface
+  }
+}
+
+private object ExtendsDeclarationPositionMatcher: LookupPositionMatcher {
+  private val EXTENDS_LIST: ElementPattern<PsiElement> = PsiJavaPatterns.psiElement().afterLeaf(
+    PsiKeyword.EXTENDS, ",").inside(PsiJavaPatterns.psiElement(JavaElementType.EXTENDS_LIST))
+
+  override fun match(position: PsiElement): Boolean {
+    return EXTENDS_LIST.accepts(position)
+  }
+
+  override fun createIncorrectElementMatcher(position: PsiElement): (LookupElement) -> Boolean {
+    val psiClass = position.parent.parent.parent as? PsiClass ?: return { false }
+
+    if (psiClass.isInterface) {
+      return this::matchClass
+    }
+
+    return this::matchElement
+  }
+
+  private fun matchClass(element: LookupElement): Boolean {
+    val psiClass = element.`object` as? PsiClass ?: return false
+    return !psiClass.isInterface
+  }
+
+  private fun matchElement(element: LookupElement): Boolean {
+    val psiClass = element.`object` as? PsiClass ?: return false
+    return psiClass.isInterface || psiClass.hasModifierProperty(PsiModifier.FINAL) // ...
+  }
+}

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

@@ -36,19 +36,19 @@ import static com.intellij.patterns.StandardPatterns.or;
 public class PreferByKindWeigher extends LookupElementWeigher {
   public static final Key<Boolean> INTRODUCED_VARIABLE = Key.create("INTRODUCED_VARIABLE");
 
-  private static final ElementPattern<PsiElement> IN_CATCH_TYPE =
+  static final ElementPattern<PsiElement> IN_CATCH_TYPE =
     psiElement().withParent(psiElement(PsiJavaCodeReferenceElement.class).
       withParent(psiElement(PsiTypeElement.class).
         withParent(or(psiElement(PsiCatchSection.class),
                       psiElement(PsiVariable.class).withParent(PsiCatchSection.class)))));
 
-  private static final ElementPattern<PsiElement> IN_MULTI_CATCH_TYPE =
+  static final ElementPattern<PsiElement> IN_MULTI_CATCH_TYPE =
     or(psiElement().afterLeaf(psiElement().withText("|").
          withParent(PsiTypeElement.class).withSuperParent(2, PsiCatchSection.class)),
        psiElement().afterLeaf(psiElement().withText("|").
          withParent(PsiTypeElement.class).withSuperParent(2, PsiParameter.class).withSuperParent(3, PsiCatchSection.class)));
 
-  private static final ElementPattern<PsiElement> INSIDE_METHOD_THROWS_CLAUSE =
+  static final ElementPattern<PsiElement> INSIDE_METHOD_THROWS_CLAUSE =
     psiElement().afterLeaf(PsiKeyword.THROWS, ",").inside(psiElement(JavaElementType.THROWS_LIST));
 
   static final ElementPattern<PsiElement> IN_RESOURCE =
@@ -56,7 +56,7 @@ public class PreferByKindWeigher extends LookupElementWeigher {
       psiElement(PsiJavaCodeReferenceElement.class).withParent(PsiTypeElement.class).
         withSuperParent(2, or(psiElement(PsiResourceVariable.class), psiElement(PsiResourceList.class))),
       psiElement(PsiReferenceExpression.class).withParent(PsiResourceExpression.class)));
-  private static final Function<PsiClass, MyResult>
+  static final Function<PsiClass, MyResult>
     PREFER_THROWABLE = psiClass -> preferClassIf(InheritanceUtil.isInheritor(psiClass, CommonClassNames.JAVA_LANG_THROWABLE));
 
   private final CompletionType myCompletionType;
@@ -118,7 +118,7 @@ public class PreferByKindWeigher extends LookupElementWeigher {
     return aClass -> MyResult.classNameOrGlobalStatic;
   }
 
-  private static List<PsiClass> getTypeBounds(PsiTypeElement typeElement) {
+  static List<PsiClass> getTypeBounds(PsiTypeElement typeElement) {
     PsiElement typeParent = typeElement.getParent();
     if (typeParent instanceof PsiReferenceParameterList) {
       int index = Arrays.asList(((PsiReferenceParameterList)typeParent).getTypeParameterElements()).indexOf(typeElement);

java/java-impl/src/com/intellij/codeInsight/completion/ml/JavaContextFeaturesProvider.kt

@@ -2,6 +2,7 @@
 package com.intellij.codeInsight.completion.ml
 
 import com.intellij.codeInsight.completion.JavaCompletionUtil
+import com.intellij.codeInsight.completion.JavaIncorrectElements
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.annotations.NotNull
 
@@ -22,7 +23,8 @@ class JavaContextFeaturesProvider : ContextFeatureProvider {
     if (JavaCompletionFeatures.isAfterMethodCall(environment)) {
       features["is_after_method_call"] = MLFeatureValue.binary(true)
     }
-    PsiTreeUtil.prevVisibleLeaf(environment.parameters.position)?.let { prevLeaf ->
+    val position = environment.parameters.position
+    PsiTreeUtil.prevVisibleLeaf(position)?.let { prevLeaf ->
       JavaCompletionFeatures.asKeyword(prevLeaf.text)?.let { keyword ->
         features["prev_neighbour_keyword"] = MLFeatureValue.categorical(keyword)
 
@@ -34,6 +36,13 @@ class JavaContextFeaturesProvider : ContextFeatureProvider {
         }
       }
     }
+    val positionMatcher = JavaIncorrectElements.matchPosition(position)
+    if (positionMatcher != null) {
+      val incorrectElementMatcher = positionMatcher.createIncorrectElementMatcher(position)
+      JavaIncorrectElements.putMatcher(incorrectElementMatcher, environment)
+
+      features["position_matcher"] = MLFeatureValue.className(positionMatcher::class.java)
+    }
     return features
   }
 }

java/java-impl/src/com/intellij/codeInsight/completion/ml/JavaElementFeaturesProvider.kt

@@ -2,6 +2,7 @@
 package com.intellij.codeInsight.completion.ml
 
 import com.intellij.codeInsight.completion.CompletionLocation
+import com.intellij.codeInsight.completion.JavaIncorrectElements
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.lang.jvm.JvmModifier
 import com.intellij.psi.*
@@ -59,6 +60,10 @@ class JavaElementFeaturesProvider : ElementFeatureProvider {
         }
       }
     }
+    val matcher = JavaIncorrectElements.tryGetMatcher(contextFeatures)
+    if (matcher != null && matcher(element)) {
+      features["incorrect_element"] = MLFeatureValue.binary(true)
+    }
 
     return features
   }

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/ml/JavaIncorrectElementsTest.kt

@@ -0,0 +1,274 @@
+// Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.java.codeInsight.completion.ml
+
+import com.intellij.codeInsight.completion.CompletionLocation
+import com.intellij.codeInsight.completion.LightCompletionTestCase
+import com.intellij.codeInsight.completion.ml.ContextFeatures
+import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
+import com.intellij.codeInsight.completion.ml.JavaElementFeaturesProvider
+import com.intellij.codeInsight.completion.ml.MLFeatureValue
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.lang.java.JavaLanguage
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiKeyword
+
+class JavaIncorrectElementsTest : LightCompletionTestCase() {
+  fun `test class extends position`() {
+    doTest("""
+        final class FinalClass {}
+        class Test extends <caret> {}
+      """.trimIndent(), mapOf(
+      "java.lang.Runnable" to true,
+      "java.lang.Thread" to false,
+      "java.lang.Override" to true,
+      "FinalClass" to true,
+    ))
+  }
+
+  fun `test interface extends position`() {
+    doTest("""interface Test extends <caret> {}""".trimIndent(), mapOf(
+      "java.lang.Runnable" to false,
+      "java.lang.Thread" to true,
+      "java.lang.Override" to false,
+    ))
+  }
+
+  fun `test class implements position`() {
+    doTest("""class Test implements <caret> {}""".trimIndent(), mapOf(
+      "java.lang.Runnable" to false,
+      "java.lang.Thread" to true,
+      "java.lang.Override" to false,
+    ))
+  }
+
+  fun `test catch clause position`() {
+    doTest("""
+      import java.io.IOException;
+
+      class Test {
+        void test() {
+          try {
+            
+          } catch(<caret>)
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+        "java.lang.Exception" to false,
+        "java.lang.Error" to false,
+        "java.lang.Throwable" to false,
+        "java.lang.RuntimeException" to false,
+        "java.lang.IllegalArgumentException" to false,
+        "java.lang.ClassNotFoundException" to true,
+        "java.io.IOException" to true
+      ))
+  }
+
+  fun `test catch clause with checked exception position`() {
+    doTest("""
+      import java.io.IOException;
+
+      class Test {
+        void test() {
+          try {
+            throw new IOException();
+          } catch(<caret>)
+        }
+      }""".trimIndent(), mapOf(
+        "java.io.IOException" to false
+      ))
+  }
+
+  fun `test catch clause with unchecked exception position`() {
+    doTest("""
+      class Test {
+        void test() {
+          try {
+            throw new Error();
+          } catch(<caret>)
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+        "java.lang.Error" to false,
+      ))
+  }
+
+  fun `test multi catch type position`() {
+    doTest("""
+      class Test {
+        void test() {
+          try {
+      
+          } catch (IllegalArgumentException | <caret>)
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+        "java.lang.Exception" to false,
+        "java.lang.Error" to false,
+        "java.lang.Throwable" to false,
+        "java.lang.RuntimeException" to false,
+        "java.lang.IllegalArgumentException" to false,
+        "java.lang.ClassNotFoundException" to true,
+      ))
+  }
+
+  fun `test throws method list position`() {
+    doTest("""
+      class Test {
+        void test() throws <caret> {
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Object" to true,
+        "java.lang.Thread" to true,
+        "java.lang.Exception" to false,
+        "java.lang.Error" to false,
+        "java.lang.Throwable" to false,
+        "java.lang.RuntimeException" to false,
+        "java.lang.IllegalArgumentException" to false,
+        "java.lang.ClassNotFoundException" to false,
+      ))
+  }
+
+  fun `test multi throws method list position`() {
+    doTest("""
+      import java.io.IOException;
+      
+      class Test {
+        void test() throws IOException, <caret> {
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+        "java.lang.Exception" to false,
+        "java.lang.Error" to false,
+        "java.lang.Throwable" to false,
+        "java.lang.RuntimeException" to false,
+        "java.lang.IllegalArgumentException" to false,
+        "java.lang.ClassNotFoundException" to false,
+      ))
+  }
+
+  fun `test javadoc @throws tag position`() {
+    doTest("""
+      /**
+       * @throws <caret>
+       */
+      class Test {
+        void test() {
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+        "java.lang.Exception" to false,
+        "java.lang.Error" to false,
+        "java.lang.Throwable" to false,
+        "java.lang.RuntimeException" to false,
+        "java.lang.IllegalArgumentException" to false,
+        "java.lang.ClassNotFoundException" to false,
+    ))
+  }
+
+  fun `test try-in-resources position`() {
+    doTest("""
+      import java.io.FileInputStream;
+      class Test {
+        void test() {
+          try(<caret>) {
+            
+          } catch(Exception e) {}
+        }
+      }""".trimIndent(), mapOf(
+        "java.io.FileInputStream" to false,
+        "java.lang.AutoCloseable" to false,
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+        "boolean" to true,
+      ))
+  }
+
+  fun `test multi try-with-resources position`() {
+    doTest("""
+      import java.util.zip.ZipFile;
+      import java.io.IOException;
+
+      class Test {
+        void test() {
+          try (ZipFile file = new ZipFile(""); <caret>) {
+          
+          } catch (IOException e) {
+            e.printStackTrace();
+          }
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.AutoCloseable" to false,
+        "java.util.zip.ZipFile" to false,
+        "java.lang.Thread" to true,
+        "java.lang.Object" to true,
+      ))
+  }
+
+  fun `test annotation position`() {
+    doTest("""
+      class Test {
+        @<caret>
+        void test() {
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Deprecated" to false,
+        "java.lang.FunctionalInterface" to true,
+      ))
+  }
+
+  fun `test type parameter position`() {
+    doTest("""
+      import java.io.Serializable;
+      class Test<T extends Serializable> {
+        static void test() {
+          new Test<<caret>>();
+        }
+      }""".trimIndent(), mapOf(
+        "java.lang.Runtime" to true,
+        "java.lang.String" to false,
+        "boolean" to true,
+      ))
+  }
+
+  private fun doTest(text: String, assertNames: Map<String, Boolean>) {
+    val assertedNames = assertNames.toMutableMap()
+
+    val overrideProvider = object: ElementFeatureProvider {
+      private val original = JavaElementFeaturesProvider()
+
+      override fun getName(): String = original.name
+      override fun calculateFeatures(element: LookupElement,
+                                     location: CompletionLocation,
+                                     contextFeatures: ContextFeatures): Map<String, MLFeatureValue> {
+        val elementFeatures = original.calculateFeatures(element, location, contextFeatures)
+        val actual = (elementFeatures["incorrect_element"]?.value ?: false) as Boolean
+        val obj = element.`object`
+        val keyName = when {
+          obj is PsiKeyword -> obj.text!!
+          obj is PsiClass && obj.qualifiedName != null -> obj.qualifiedName!!
+          else -> return elementFeatures
+        }
+
+        val expected = assertedNames[keyName] ?: return elementFeatures
+        assertEquals(keyName, expected, actual)
+        assertedNames.remove(keyName)
+        return elementFeatures
+      }
+    }
+
+    try {
+      ElementFeatureProvider.EP_NAME.addExplicitExtension(JavaLanguage.INSTANCE, overrideProvider)
+      configureFromFileText("test.java", text)
+      complete()
+      assertTrue("Lookup doesn't contain next elements: ${assertedNames.keys.joinToString()}", assertedNames.isEmpty())
+    }
+    finally {
+      ElementFeatureProvider.EP_NAME.removeExplicitExtension(JavaLanguage.INSTANCE, overrideProvider)
+    }
+  }
+}