From 2dbfc7291569fb2b9cf2fcb4bd1760af31656cf4 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Sun, 18 Feb 2018 17:29:54 +0700
Subject: [PATCH] AnnotateIntentionAction: support Unmodifiable(View)
 annotations in internal mode

Part of IDEA-182125
---
 .../codeInspection/dataFlow/Mutability.java   |  7 ++-
 .../impl/AnnotateIntentionAction.java         | 56 +++++++++++++++++--
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/Mutability.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/Mutability.java
index 3621136b9f52..db25e4e2bdb6 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/Mutability.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/Mutability.java
@@ -14,6 +14,9 @@ import java.util.Collections;
 public enum Mutability {
   UNKNOWN, MUTABLE, UNMODIFIABLE, UNMODIFIABLE_VIEW;
 
+  public static final String UNMODIFIABLE_ANNOTATION = "org.jetbrains.annotations.Unmodifiable";
+  public static final String UNMODIFIABLE_VIEW_ANNOTATION = "org.jetbrains.annotations.UnmodifiableView";
+
   public boolean isUnmodifiable() {
     return this == UNMODIFIABLE || this == UNMODIFIABLE_VIEW;
   }
@@ -37,13 +40,13 @@ public enum Mutability {
         return UNKNOWN;
       }
     }
-    if (AnnotationUtil.isAnnotated(owner, Collections.singleton("org.jetbrains.annotations.Unmodifiable"),
+    if (AnnotationUtil.isAnnotated(owner, Collections.singleton(UNMODIFIABLE_ANNOTATION),
                                    AnnotationUtil.CHECK_HIERARCHY |
                                    AnnotationUtil.CHECK_EXTERNAL |
                                    AnnotationUtil.CHECK_INFERRED)) {
       return UNMODIFIABLE;
     }
-    if (AnnotationUtil.isAnnotated(owner, Collections.singleton("org.jetbrains.annotations.UnmodifiableView"),
+    if (AnnotationUtil.isAnnotated(owner, Collections.singleton(UNMODIFIABLE_VIEW_ANNOTATION),
                                    AnnotationUtil.CHECK_HIERARCHY |
                                    AnnotationUtil.CHECK_EXTERNAL |
                                    AnnotationUtil.CHECK_INFERRED)) {
diff --git a/java/java-impl/src/com/intellij/codeInsight/intention/impl/AnnotateIntentionAction.java b/java/java-impl/src/com/intellij/codeInsight/intention/impl/AnnotateIntentionAction.java
index 1bb96d87a116..9e43716bb3b3 100644
--- a/java/java-impl/src/com/intellij/codeInsight/intention/impl/AnnotateIntentionAction.java
+++ b/java/java-impl/src/com/intellij/codeInsight/intention/impl/AnnotateIntentionAction.java
@@ -9,18 +9,18 @@ import com.intellij.codeInsight.NullableNotNullManager;
 import com.intellij.codeInsight.intention.AddAnnotationFix;
 import com.intellij.codeInsight.intention.AddAnnotationPsiFix;
 import com.intellij.codeInsight.intention.LowPriorityAction;
+import com.intellij.codeInspection.dataFlow.Mutability;
+import com.intellij.openapi.application.ex.ApplicationManagerEx;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.popup.JBPopupFactory;
 import com.intellij.openapi.ui.popup.PopupStep;
 import com.intellij.openapi.ui.popup.util.BaseListPopupStep;
 import com.intellij.openapi.util.text.StringUtil;
-import com.intellij.psi.CommonClassNames;
-import com.intellij.psi.PsiAnnotation;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiModifierListOwner;
+import com.intellij.psi.*;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.siyeh.ig.psiutils.ClassUtils;
 import one.util.streamex.MoreCollectors;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
@@ -34,7 +34,9 @@ public class AnnotateIntentionAction extends BaseIntentionAction implements LowP
   private static final AnnotationProvider[] PROVIDERS = {
     new DeprecationAnnotationProvider(),
     new NullableAnnotationProvider(),
-    new NotNullAnnotationProvider()
+    new NotNullAnnotationProvider(),
+    new UnmodifiableAnnotationProvider(),
+    new UnmodifiableViewAnnotationProvider()
   };
   private AnnotationProvider myAnnotationProvider;
   private boolean mySingleMode;
@@ -203,4 +205,48 @@ public class AnnotateIntentionAction extends BaseIntentionAction implements LowP
       return NullableNotNullManager.getInstance(project).getNullables().toArray(ArrayUtil.EMPTY_STRING_ARRAY);
     }
   }
+
+  static class UnmodifiableAnnotationProvider implements AnnotationProvider {
+
+    @NotNull
+    @Override
+    public String getName(Project project) {
+      return Mutability.UNMODIFIABLE_ANNOTATION;
+    }
+
+    @Override
+    public boolean isAvailable(PsiModifierListOwner owner) {
+      return ApplicationManagerEx.getApplicationEx().isInternal() &&
+             owner instanceof PsiMethod &&
+             !ClassUtils.isImmutable(((PsiMethod)owner).getReturnType());
+    }
+
+    @NotNull
+    @Override
+    public String[] getAnnotationsToRemove(Project project) {
+      return new String[]{Mutability.UNMODIFIABLE_VIEW_ANNOTATION};
+    }
+  }
+
+  static class UnmodifiableViewAnnotationProvider implements AnnotationProvider {
+
+    @NotNull
+    @Override
+    public String getName(Project project) {
+      return Mutability.UNMODIFIABLE_VIEW_ANNOTATION;
+    }
+
+    @Override
+    public boolean isAvailable(PsiModifierListOwner owner) {
+      return ApplicationManagerEx.getApplicationEx().isInternal() &&
+             owner instanceof PsiMethod &&
+             !ClassUtils.isImmutable(((PsiMethod)owner).getReturnType());
+    }
+
+    @NotNull
+    @Override
+    public String[] getAnnotationsToRemove(Project project) {
+      return new String[]{Mutability.UNMODIFIABLE_ANNOTATION};
+    }
+  }
 }
\ No newline at end of file