From 96b9ab4f804a214b360bd1497e9781daef96ebb3 Mon Sep 17 00:00:00 2001
From: Pavel Dolgov <pavel.dolgov@jetbrains.com>
Date: Thu, 21 Feb 2019 16:14:01 +0300
Subject: [PATCH] Java: Intention that wraps list/set/map with
 Collections.unmodifiable - handle the caret at the end of the expression
 (IDEA-93154)

---
 .../impl/WrapWithUnmodifiableAction.java      | 25 +++++++++++++------
 .../afterCaretAtTheEnd.java                   | 11 ++++++++
 .../beforeCaretAtTheEnd.java                  | 10 ++++++++
 3 files changed, 39 insertions(+), 7 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/afterCaretAtTheEnd.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/beforeCaretAtTheEnd.java

diff --git a/java/java-impl/src/com/intellij/codeInsight/intention/impl/WrapWithUnmodifiableAction.java b/java/java-impl/src/com/intellij/codeInsight/intention/impl/WrapWithUnmodifiableAction.java
index 1a2092f10770..82728b599c54 100644
--- a/java/java-impl/src/com/intellij/codeInsight/intention/impl/WrapWithUnmodifiableAction.java
+++ b/java/java-impl/src/com/intellij/codeInsight/intention/impl/WrapWithUnmodifiableAction.java
@@ -2,7 +2,6 @@
 package com.intellij.codeInsight.intention.impl;
 
 import com.intellij.codeInsight.CodeInsightBundle;
-import com.intellij.codeInsight.intention.PsiElementBaseIntentionAction;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
@@ -24,13 +23,16 @@ import static com.intellij.util.ObjectUtils.tryCast;
 /**
  * @author Pavel.Dolgov
  */
-public class WrapWithUnmodifiableAction extends PsiElementBaseIntentionAction {
+public class WrapWithUnmodifiableAction extends BaseIntentionAction {
   private static final String JAVA_UTIL_SORTED_SET = "java.util.SortedSet";
   private static final String JAVA_UTIL_SORTED_MAP = "java.util.SortedMap";
 
   @Override
-  public void invoke(@NotNull Project project, Editor editor, @NotNull PsiElement element) throws IncorrectOperationException {
-    PsiExpression expression = getParentExpression(element);
+  public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+    if (editor == null || file == null || !canModify(file)) {
+      return;
+    }
+    PsiExpression expression = getParentExpression(editor, file);
     if (expression != null) {
       PsiClass psiClass = PsiUtil.resolveClassInClassTypeOnly(expression.getType());
       if (psiClass != null) {
@@ -59,8 +61,14 @@ public class WrapWithUnmodifiableAction extends PsiElementBaseIntentionAction {
     }
   }
 
-  private static PsiExpression getParentExpression(@NotNull PsiElement element) {
+  @Nullable
+  private static PsiExpression getParentExpression(@NotNull Editor editor, @NotNull PsiFile file) {
+    PsiElement element = file.findElementAt(editor.getCaretModel().getOffset());
     PsiExpression expression = PsiTreeUtil.getNonStrictParentOfType(element, PsiExpression.class);
+    if (expression == null) {
+      element = file.findElementAt(editor.getCaretModel().getOffset() - 1);
+      expression = PsiTreeUtil.getNonStrictParentOfType(element, PsiExpression.class);
+    }
     if (expression != null) {
       PsiMethodCallExpression methodCall = tryCast(expression.getParent(), PsiMethodCallExpression.class);
       if (methodCall != null && methodCall.getMethodExpression() == expression) {
@@ -78,8 +86,11 @@ public class WrapWithUnmodifiableAction extends PsiElementBaseIntentionAction {
   }
 
   @Override
-  public boolean isAvailable(@NotNull Project project, Editor editor, @NotNull PsiElement element) {
-    PsiExpression expression = getParentExpression(element);
+  public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+    if (editor == null || file == null || !canModify(file)) {
+      return false;
+    }
+    PsiExpression expression = getParentExpression(editor, file);
     if (expression != null) {
       if (isUnmodifiable(expression)) {
         return false;
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/afterCaretAtTheEnd.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/afterCaretAtTheEnd.java
new file mode 100644
index 000000000000..0b7116dec24b
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/afterCaretAtTheEnd.java
@@ -0,0 +1,11 @@
+// "Wrap with unmodifiable set" "true"
+import java.util.Collections;
+import java.util.Set;
+import java.util.HashSet;
+
+class C {
+    Set<String> test() {
+        Set<String> result = new HashSet<>();
+        return Collections.unmodifiableSet(result);
+    }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/beforeCaretAtTheEnd.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/beforeCaretAtTheEnd.java
new file mode 100644
index 000000000000..70d4478e3610
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/wrapWithUnmodifiable/beforeCaretAtTheEnd.java
@@ -0,0 +1,10 @@
+// "Wrap with unmodifiable set" "true"
+import java.util.Set;
+import java.util.HashSet;
+
+class C {
+    Set<String> test() {
+        Set<String> result = new HashSet<>();
+        return result<caret>;
+    }
+}
\ No newline at end of file