From afc371e02957e2fecec92dc521547a1684da5002 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Fri, 19 Mar 2021 13:40:51 +0700
Subject: [PATCH] [java-inspections] Quick-fix for missing not-null annotation
 on type element

GitOrigin-RevId: 6c7d91b8f30015eb1b5f51be422951ba8ef6b36b
---
 .../messages/JavaAnalysisBundle.properties    |  1 +
 .../intention/AddTypeAnnotationFix.java       | 48 +++++++++++++++++
 .../nullable/NullableStuffInspectionBase.java |  7 ++-
 .../QuickFixOnTypeArgument.java               |  7 +++
 .../QuickFixOnTypeArgumentNullable.java       |  7 +++
 .../QuickFixOnTypeArgumentNullable_after.java |  7 +++
 .../QuickFixOnTypeArgument_after.java         |  7 +++
 .../NullableStuffInspectionTest.java          | 52 +++++++++++++------
 8 files changed, 119 insertions(+), 17 deletions(-)
 create mode 100644 java/java-analysis-impl/src/com/intellij/codeInsight/intention/AddTypeAnnotationFix.java
 create mode 100644 java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument.java
 create mode 100644 java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable.java
 create mode 100644 java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable_after.java
 create mode 100644 java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument_after.java

diff --git a/java/java-analysis-api/resources/messages/JavaAnalysisBundle.properties b/java/java-analysis-api/resources/messages/JavaAnalysisBundle.properties
index d39c2abf025c..9826b9de7bfe 100644
--- a/java/java-analysis-api/resources/messages/JavaAnalysisBundle.properties
+++ b/java/java-analysis-api/resources/messages/JavaAnalysisBundle.properties
@@ -155,6 +155,7 @@ service.uses=Uses service <a href="#javaClass/{0}">{0}</a><br><div style='margin
 suppress.for.parameter=Suppress for parameter
 suppress.for.statement.with.comment=Suppress for statement with comment
 intention.add.annotation.family=Add annotation
+intention.add.type.annotation.family=Add type annotation
 intention.text.remove.annotation=Remove
 
 scope.package=Package {0}
diff --git a/java/java-analysis-impl/src/com/intellij/codeInsight/intention/AddTypeAnnotationFix.java b/java/java-analysis-impl/src/com/intellij/codeInsight/intention/AddTypeAnnotationFix.java
new file mode 100644
index 000000000000..df5543ed3e89
--- /dev/null
+++ b/java/java-analysis-impl/src/com/intellij/codeInsight/intention/AddTypeAnnotationFix.java
@@ -0,0 +1,48 @@
+// Copyright 2000-2021 JetBrains s.r.o. 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.intention;
+
+import com.intellij.codeInspection.LocalQuickFix;
+import com.intellij.codeInspection.ProblemDescriptor;
+import com.intellij.java.analysis.JavaAnalysisBundle;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.PsiAnnotation;
+import com.intellij.psi.PsiTypeElement;
+import com.intellij.util.ObjectUtils;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collection;
+
+public class AddTypeAnnotationFix implements LocalQuickFix {
+  private final @NotNull String myAnnotationToAdd;
+  @SafeFieldForPreview
+  private final @NotNull Collection<String> myAnnotationsToRemove;
+
+  public AddTypeAnnotationFix(@NotNull String annotationToAdd, @NotNull Collection<String> annotationsToRemove) {
+    myAnnotationToAdd = annotationToAdd;
+    myAnnotationsToRemove = annotationsToRemove;
+  }
+
+  @Override
+  public @NotNull String getName() {
+    return JavaAnalysisBundle.message("inspection.i18n.quickfix.annotate.as", StringUtil.getShortName(myAnnotationToAdd));
+  }
+
+  @Override
+  public @NotNull String getFamilyName() {
+    return JavaAnalysisBundle.message("intention.add.type.annotation.family");
+  }
+
+  @Override
+  public void applyFix(@NotNull Project project,
+                       @NotNull ProblemDescriptor descriptor) {
+    PsiTypeElement typeElement = ObjectUtils.tryCast(descriptor.getStartElement(), PsiTypeElement.class);
+    if (typeElement == null || !typeElement.acceptsAnnotations()) return;
+    for (PsiAnnotation annotation : typeElement.getAnnotations()) {
+      if (myAnnotationsToRemove.contains(annotation.getQualifiedName())) {
+        annotation.delete();
+      }
+    }
+    typeElement.addAnnotation(myAnnotationToAdd);
+  }
+}
diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/nullable/NullableStuffInspectionBase.java b/java/java-analysis-impl/src/com/intellij/codeInspection/nullable/NullableStuffInspectionBase.java
index 20a0476b0ccb..ed25449e0c7e 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/nullable/NullableStuffInspectionBase.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/nullable/NullableStuffInspectionBase.java
@@ -7,6 +7,7 @@ import com.intellij.codeInsight.NullabilityAnnotationInfo;
 import com.intellij.codeInsight.NullableNotNullManager;
 import com.intellij.codeInsight.daemon.impl.analysis.JavaGenericsUtil;
 import com.intellij.codeInsight.intention.AddAnnotationPsiFix;
+import com.intellij.codeInsight.intention.AddTypeAnnotationFix;
 import com.intellij.codeInspection.*;
 import com.intellij.codeInspection.dataFlow.DfaPsiUtil;
 import com.intellij.codeInspection.dataFlow.instructions.MethodCallInstruction;
@@ -257,9 +258,11 @@ public class NullableStuffInspectionBase extends AbstractBaseJavaLocalInspection
           PsiTypeElement[] typeArguments = getReferenceTypeArguments(reference);
           if (typeParameters.length > 0 && typeParameters.length == typeArguments.length && !(typeArguments[0].getType() instanceof PsiDiamondType)) {
             for (int i = 0; i < typeParameters.length; i++) {
+              PsiTypeElement typeArgument = typeArguments[i];
               if (DfaPsiUtil.getTypeNullability(JavaPsiFacade.getElementFactory(element.getProject()).createType(typeParameters[i])) ==
-                  Nullability.NOT_NULL && DfaPsiUtil.getTypeNullability(typeArguments[i].getType()) != Nullability.NOT_NULL) {
-                reportProblem(holder, typeArguments[i], "non.null.type.argument.is.expected");
+                  Nullability.NOT_NULL && DfaPsiUtil.getTypeNullability(typeArgument.getType()) != Nullability.NOT_NULL) {
+                AddTypeAnnotationFix fix = new AddTypeAnnotationFix(manager.getDefaultNotNull(), manager.getNullables());
+                reportProblem(holder, typeArgument, fix, "non.null.type.argument.is.expected");
               }
             }
           }
diff --git a/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument.java b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument.java
new file mode 100644
index 000000000000..be3a57398580
--- /dev/null
+++ b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument.java
@@ -0,0 +1,7 @@
+import typeUse.*;
+
+class Bar<T extends @NotNull Object> {
+  public static void main(String[] args) {
+    Bar<<warning descr="Non-null type argument is expected">S<caret>tring</warning>> bar;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable.java b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable.java
new file mode 100644
index 000000000000..0a266e25a26c
--- /dev/null
+++ b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable.java
@@ -0,0 +1,7 @@
+import typeUse.*;
+
+class Bar<T extends @NotNull Object> {
+  public static void main(String[] args) {
+    Bar<<warning descr="Non-null type argument is expected">@Nullable S<caret>tring</warning>> bar;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable_after.java b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable_after.java
new file mode 100644
index 000000000000..914cafb5343e
--- /dev/null
+++ b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgumentNullable_after.java
@@ -0,0 +1,7 @@
+import typeUse.*;
+
+class Bar<T extends @NotNull Object> {
+  public static void main(String[] args) {
+    Bar<@NotNull S<caret>tring> bar;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument_after.java b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument_after.java
new file mode 100644
index 000000000000..914cafb5343e
--- /dev/null
+++ b/java/java-tests/testData/inspection/nullableProblems/QuickFixOnTypeArgument_after.java
@@ -0,0 +1,7 @@
+import typeUse.*;
+
+class Bar<T extends @NotNull Object> {
+  public static void main(String[] args) {
+    Bar<@NotNull S<caret>tring> bar;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testSrc/com/intellij/java/codeInspection/NullableStuffInspectionTest.java b/java/java-tests/testSrc/com/intellij/java/codeInspection/NullableStuffInspectionTest.java
index 214c91cee718..14f2e29ca815 100644
--- a/java/java-tests/testSrc/com/intellij/java/codeInspection/NullableStuffInspectionTest.java
+++ b/java/java-tests/testSrc/com/intellij/java/codeInspection/NullableStuffInspectionTest.java
@@ -45,6 +45,12 @@ public class NullableStuffInspectionTest extends LightJavaCodeInsightFixtureTest
     myFixture.testHighlighting(true, false, true, getTestName(false) + ".java");
   }
 
+  private void doTestWithFix(String intentionAction) {
+    doTest();
+    myFixture.launchAction(myFixture.findSingleIntention(intentionAction));
+    myFixture.checkResultByFile(getTestName(false) + "_after.java");
+  }
+
   @Override
   public void setUp() throws Exception {
     super.setUp();
@@ -289,30 +295,22 @@ public class NullableStuffInspectionTest extends LightJavaCodeInsightFixtureTest
   }
 
   public void testAnnotateQuickFixOnMethodReference() {
-    doTest();
-    myFixture.launchAction(myFixture.findSingleIntention("Annotate"));
-    myFixture.checkResultByFile(getTestName(false) + "_after.java");
+    doTestWithFix("Annotate");
   }
 
   public void testAnnotateOverridingParametersOnNotNullMethod() {
     myInspection.REPORT_ANNOTATION_NOT_PROPAGATED_TO_OVERRIDERS = true;
-    doTest();
-    myFixture.launchAction(myFixture.findSingleIntention("Annotate overridden method parameters"));
-    myFixture.checkResultByFile(getTestName(false) + "_after.java");
+    doTestWithFix("Annotate overridden method parameters");
   }
 
   public void testRemoveMethodAnnotationRemovesOverriders() {
     myInspection.REPORT_ANNOTATION_NOT_PROPAGATED_TO_OVERRIDERS = true;
-    doTest();
-    myFixture.launchAction(myFixture.findSingleIntention("Remove annotation"));
-    myFixture.checkResultByFile(getTestName(false) + "_after.java");
+    doTestWithFix("Remove annotation");
   }
 
   public void testRemoveParameterAnnotationRemovesOverriders() {
     myInspection.REPORT_ANNOTATION_NOT_PROPAGATED_TO_OVERRIDERS = true;
-    doTest();
-    myFixture.launchAction(myFixture.findSingleIntention("Remove annotation"));
-    myFixture.checkResultByFile(getTestName(false) + "_after.java");
+    doTestWithFix("Remove annotation");
   }
 
   public void testNullPassedToNullableParameter() {
@@ -343,8 +341,32 @@ public class NullableStuffInspectionTest extends LightJavaCodeInsightFixtureTest
   public void testOverridersHaveNonDefaultAnnotation() {
     myFixture.addClass("package org.eclipse.jdt.annotation;\n\nimport java.lang.annotation.*;\n\n@Target(ElementType.PARAMETER) public @interface NonNull { }");
     myInspection.REPORT_ANNOTATION_NOT_PROPAGATED_TO_OVERRIDERS = true;
-    doTest();
-    myFixture.launchAction(myFixture.findSingleIntention("Annotate overridden method parameters"));
-    myFixture.checkResultByFile(getTestName(false) + "_after.java");
+    doTestWithFix("Annotate overridden method parameters");
+  }
+
+  public void testQuickFixOnTypeArgument() {
+    DataFlowInspection8Test.setupTypeUseAnnotations("typeUse", myFixture);
+    NullableNotNullManager manager = NullableNotNullManager.getInstance(getProject());
+    String oldDefault = manager.getDefaultNotNull();
+    try {
+      manager.setDefaultNotNull("typeUse.NotNull");
+      doTestWithFix("Annotate as @NotNull");
+    }
+    finally {
+      manager.setDefaultNotNull(oldDefault);
+    }
+  }
+
+  public void testQuickFixOnTypeArgumentNullable() {
+    DataFlowInspection8Test.setupTypeUseAnnotations("typeUse", myFixture);
+    NullableNotNullManager manager = NullableNotNullManager.getInstance(getProject());
+    String oldDefault = manager.getDefaultNotNull();
+    try {
+      manager.setDefaultNotNull("typeUse.NotNull");
+      doTestWithFix("Annotate as @NotNull");
+    }
+    finally {
+      manager.setDefaultNotNull(oldDefault);
+    }
   }
 }
\ No newline at end of file