[Java. Code Formatting] IDEA-168035 Treat jspecify annotations as type annotations in formatter

- Use language level from module and fix tests after this
- Check name of the import module explicitly

GitOrigin-RevId: a0cd8502f4d3d00c9de8b38cea66f61eb3e940d1

java/java-frontback-impl/src/com/intellij/psi/formatter/java/JavaFormatterAnnotationUtil.kt

@@ -4,8 +4,7 @@ package com.intellij.psi.formatter.java
 import com.intellij.codeInsight.AnnotationUtil
 import com.intellij.codeInsight.DumbAwareAnnotationUtil
 import com.intellij.lang.ASTNode
-import com.intellij.openapi.roots.LanguageLevelProjectExtension
-import com.intellij.pom.java.LanguageLevel
+import com.intellij.pom.java.JavaFeature
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiKeyword
 import com.intellij.psi.PsiModifierListOwner
@@ -14,6 +13,7 @@ import com.intellij.psi.formatter.FormatterUtil
 import com.intellij.psi.formatter.java.JavaFormatterAnnotationUtil.isFieldWithAnnotations
 import com.intellij.psi.impl.source.tree.JavaElementType
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.PsiUtil
 
 internal object JavaFormatterAnnotationUtil {
   private val KNOWN_TYPE_ANNOTATIONS: Set<String> = setOf(
@@ -34,8 +34,7 @@ internal object JavaFormatterAnnotationUtil {
   fun isTypeAnnotation(annotation: ASTNode): Boolean {
     val node = annotation.psi as? PsiAnnotation ?: return false
 
-    val languageLevel = LanguageLevelProjectExtension.getInstance(node.project).languageLevel
-    if (languageLevel.isLessThan(LanguageLevel.JDK_1_8)) return false
+    if (!PsiUtil.isAvailable(JavaFeature.TYPE_ANNOTATIONS, node)) return false
 
     val next = PsiTreeUtil.skipSiblingsForward(node, PsiWhiteSpace::class.java, PsiAnnotation::class.java)
     if (next is PsiKeyword) return false

java/java-psi-api/src/com/intellij/codeInsight/DumbAwareAnnotationUtil.kt

@@ -3,14 +3,14 @@ package com.intellij.codeInsight
 
 import com.intellij.codeInsight.DumbAwareAnnotationUtil.KNOWN_ANNOTATIONS
 import com.intellij.codeInsight.DumbAwareAnnotationUtil.hasAnnotation
-import com.intellij.openapi.roots.LanguageLevelProjectExtension
 import com.intellij.openapi.util.NlsSafe
 import com.intellij.openapi.util.text.StringUtil
-import com.intellij.pom.java.LanguageLevel
+import com.intellij.pom.java.JavaFeature
 import com.intellij.psi.*
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.PsiModificationTracker
+import com.intellij.psi.util.PsiUtil
 
 /**
  * Utility which helps to detect annotation in `Dumb mode`.
@@ -18,6 +18,9 @@ import com.intellij.psi.util.PsiModificationTracker
 object DumbAwareAnnotationUtil {
   private const val JAVA_LANG_PACKAGE = "java.lang"
 
+  /**
+   * Represents a list of fully qualified names for annotations that are treated as a type annotations.
+   */
   private val KNOWN_ANNOTATIONS = setOf(
     AnnotationUtil.NOT_NULL,
     AnnotationUtil.NULLABLE,
@@ -26,6 +29,14 @@ object DumbAwareAnnotationUtil {
     AnnotationUtil.J_SPECIFY_NULLABLE
   )
 
+  /**
+   * Represents a mapping from a fully qualified name of a module to a set of fully qualified names of annotations
+   * that are treated as a type annotations and located in this module
+   */
+  private val KNOWN_MODULE_TO_ANNOTATIONS_MAP = mapOf(
+    "org.jspecify" to setOf(AnnotationUtil.J_SPECIFY_NON_NULL, AnnotationUtil.J_SPECIFY_NULLABLE)
+  )
+
   /**
    * Checks if the given `PsiModifierListOwner` has an annotation with the specified fully qualified name. If the annotation has not
    * fqn (which is more likely), it will find it if and only if it is present in [KNOWN_ANNOTATIONS].
@@ -82,12 +93,12 @@ object DumbAwareAnnotationUtil {
   }
 
   private fun isAnnotationInModuleImportList(annotationFqn: String, moduleList: PsiImportList): Boolean {
-    val languageLevel = LanguageLevelProjectExtension.getInstance(moduleList.project).languageLevel
-    if (languageLevel.isLessThan(LanguageLevel.JDK_25)) return false
+    if (!PsiUtil.isAvailable(JavaFeature.MODULE_IMPORT_DECLARATIONS, moduleList)) return false
     return moduleList.importModuleStatements.any { statement: PsiImportModuleStatement ->
-      val referenceElement = statement.moduleReference ?: return@any false
-      val referenceElementText = getCanonicalTextOfTheReference(referenceElement)
-      annotationFqn.startsWith(referenceElementText)
+      val referenceName = statement.referenceName ?: return@any false
+      val formattedReferenceName = getFormattedReferenceFqn(referenceName)
+      if (formattedReferenceName !in KNOWN_MODULE_TO_ANNOTATIONS_MAP) return@any false
+      annotationFqn in KNOWN_MODULE_TO_ANNOTATIONS_MAP.getValue(formattedReferenceName)
     }
   }
 
@@ -97,10 +108,7 @@ object DumbAwareAnnotationUtil {
     return AnnotationImportInfo(packageName, className)
   }
 
-  private fun getCanonicalTextOfTheReference(reference: PsiElement): String {
-    if (reference !is PsiJavaCodeReferenceElement && reference !is PsiJavaModuleReferenceElement) error("Only code and module references are supported ")
-    return getFormattedReferenceFqn(reference.text)
-  }
+  private fun getCanonicalTextOfTheReference(reference: PsiJavaCodeReferenceElement): String = getFormattedReferenceFqn(reference.text)
 
   private data class AnnotationImportInfo(val packageName: String, val className: String)
 }

java/java-tests/testData/psi/formatter/java/jSpecifyTypeAnnotation/UnknownModuleImport.java

@@ -0,0 +1,23 @@
+package org.example;
+
+import module org.jspecify.annotations;
+
+public class Formatter {
+    @NonNull @Nullable String breakLineBetweenAnnotations() {
+        return "";
+    }
+
+    @Nullable
+    @NonNull String breakLineBetweenTypeAndAnnotations() {
+        return null;
+    }
+
+    @Nullable <T> String breakLineBetweenTypeParameterAndAnnotation() {
+        return null;
+    }
+
+    @Nullable
+    @NonNull <T> String breakLineMixed() {
+      return null;
+    }
+}

java/java-tests/testData/psi/formatter/java/jSpecifyTypeAnnotation/UnknownModuleImport_after.java

@@ -0,0 +1,28 @@
+package org.example;
+
+import module org.jspecify.annotations;
+
+public class Formatter {
+    @NonNull
+    @Nullable
+    String breakLineBetweenAnnotations() {
+        return "";
+    }
+
+    @Nullable
+    @NonNull
+    String breakLineBetweenTypeAndAnnotations() {
+        return null;
+    }
+
+    @Nullable
+    <T> String breakLineBetweenTypeParameterAndAnnotation() {
+        return null;
+    }
+
+    @Nullable
+    @NonNull
+    <T> String breakLineMixed() {
+        return null;
+    }
+}

java/java-tests/testSrc/com/intellij/codeInsight/daemon/impl/quickfix/AnnotateMethodInGeneratedFilesTest.java

@@ -8,7 +8,9 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.roots.GeneratedSourcesFilter;
 import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.pom.java.LanguageLevel;
 import com.intellij.psi.PsiClass;
+import com.intellij.testFramework.IdeaTestUtil;
 import com.intellij.testFramework.fixtures.LightJavaCodeInsightFixtureTestCase;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
@@ -38,6 +40,7 @@ public class AnnotateMethodInGeneratedFilesTest extends LightJavaCodeInsightFixt
   }
 
   public void testAnnotateOverriddenMethod() {
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.getModule(), LanguageLevel.JDK_1_6);
     doTest("Add missing nullability annotation");
   }
 

java/java-tests/testSrc/com/intellij/java/codeInsight/ExternalAnnotationsTest.java

@@ -14,12 +14,14 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.roots.JavaModuleExternalPaths;
 import com.intellij.openapi.roots.ModuleRootModificationUtil;
 import com.intellij.openapi.vfs.VfsUtilCore;
+import com.intellij.pom.java.LanguageLevel;
 import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiMethod;
 import com.intellij.psi.codeStyle.JavaCodeStyleSettings;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.testFramework.IdeaTestUtil;
 import com.intellij.testFramework.UsefulTestCase;
 import com.intellij.testFramework.builders.JavaModuleFixtureBuilder;
 import com.intellij.testFramework.fixtures.*;
@@ -70,6 +72,7 @@ public class ExternalAnnotationsTest extends UsefulTestCase {
 
   public void testAddedAnnotationInCodeWhenAlreadyPresent() {
     myFixture.configureByFile("src/withAnnotation/Foo.java");
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.getModule(), LanguageLevel.JDK_1_8);
     PsiMethod method = PsiTreeUtil.getParentOfType(myFixture.getElementAtCaret(), PsiMethod.class, false);
     assertNotNull(method);
     ActionContext context = myFixture.getActionContext();

java/java-tests/testSrc/com/intellij/java/psi/formatter/java/JSpecifyTypeAnnotationFormatterTest.kt

@@ -24,7 +24,7 @@ class JSpecifyTypeAnnotationFormatterTest : LightJavaCodeInsightFixtureTestCase(
     super.setUp()
     commonSettings.KEEP_LINE_BREAKS = false
     commonSettings.METHOD_ANNOTATION_WRAP = WRAP_ALWAYS
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_1_8)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_1_8)
     ModuleRootModificationUtil.updateModel(module) { model ->
       MavenDependencyUtil.addFromMaven(model, "org.jspecify:jspecify:1.0.0")
     }
@@ -58,33 +58,38 @@ class JSpecifyTypeAnnotationFormatterTest : LightJavaCodeInsightFixtureTestCase(
   }
 
   fun testLowLanguageLevel() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_1_7)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_1_7)
     doTest()
   }
 
   fun testModuleImport() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_25)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_25)
     doTest()
   }
 
   fun testModuleImportWithSpaces() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_25)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_25)
     doTest()
   }
 
   fun testModuleImportMixedWithPackageImport() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_25)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_25)
     doTest()
   }
 
   fun testModuleImportMixedWithFqn() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_25)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_25)
     doTest()
   }
 
 
   fun testLowLanguageLevelForModuleImport() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_24)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_24)
+    doTest()
+  }
+
+  fun testUnknownModuleImport() {
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_25)
     doTest()
   }
 

java/java-tests/testSrc/com/intellij/java/psi/formatter/java/TypeAnnotationFormatterTest.kt

@@ -24,7 +24,7 @@ class TypeAnnotationFormatterTest : LightJavaCodeInsightFixtureTestCase() {
     super.setUp()
     commonSettings.KEEP_LINE_BREAKS = false
     commonSettings.METHOD_ANNOTATION_WRAP = WRAP_ALWAYS
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_1_8)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_1_8)
     ModuleRootModificationUtil.updateModel(module, DefaultLightProjectDescriptor::addJetBrainsAnnotations)
   }
 
@@ -56,7 +56,7 @@ class TypeAnnotationFormatterTest : LightJavaCodeInsightFixtureTestCase() {
   }
 
   fun testLowLanguageLevel() {
-    IdeaTestUtil.setProjectLanguageLevel(myFixture.project, LanguageLevel.JDK_1_7)
+    IdeaTestUtil.setModuleLanguageLevel(myFixture.module, LanguageLevel.JDK_1_7)
     doTest()
   }