Java8MapApiInspection: suggest to replace for each loop with replaceAll method call (IDEA-208005)

2026-09-27 10:03:11 +07:00 · 2019-03-05 12:02:49 +07:00
parent 2bf3db47a1
commit 45d2ba086f
18 changed files with 474 additions and 12 deletions
@@ -14,6 +14,9 @@ import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.JavaCodeStyleManager;
 import com.intellij.psi.search.searches.ReferencesSearch;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ArrayUtil;
+import com.intellij.util.ObjectUtils;
+import com.siyeh.ig.PsiReplacementUtil;
 import com.siyeh.ig.psiutils.*;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Nls;
@@ -21,6 +24,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

 import javax.swing.*;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
@@ -93,6 +97,30 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool {
        }
      }

+      @Override
+      public void visitForeachStatement(PsiForeachStatement statement) {
+        MapLoopCondition condition = MapLoopCondition.create(statement);
+        if (condition == null) return;
+        PsiMethodCallExpression putCall = condition.extractPut(statement);
+        if (putCall == null) return;
+        PsiExpression[] args = putCall.getArgumentList().getExpressions();
+        if (args.length != 2) return;
+        PsiExpression key = args[0];
+        if (!condition.isKeyAccess(key)) return;
+        PsiExpression value = args[1];
+        if (condition.isEntrySet() && isUsedAsReference(value, condition)) return;
+
+        ReplaceWithSingleMapOperation fix = ReplaceWithSingleMapOperation.create("replaceAll", putCall, value);
+        holder.registerProblem(statement.getFirstChild(), QuickFixBundle.message("java.8.map.api.inspection.description", fix.myMethodName),
+                               fix);
+      }
+
+      private boolean isUsedAsReference(@NotNull PsiElement value, @NotNull MapLoopCondition condition) {
+        VariableUsedAsReferenceVisitor visitor = new VariableUsedAsReferenceVisitor(condition.getIterParam(), "getValue", "getKey");
+        value.accept(visitor);
+        return visitor.isUsed;
+      }
+
      private void processMerge(MapCheckCondition condition,
                                PsiStatement existsBranch,
                                PsiStatement noneBranch) {
@@ -221,6 +249,44 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool {
    return nameCandidate.toLowerCase(Locale.ENGLISH);
  }

+  private static class VariableUsedAsReferenceVisitor extends JavaRecursiveElementWalkingVisitor {
+    private final PsiVariable myVariable;
+    private final String[] myExcludeMethodNames;
+
+    private boolean isUsed = false;
+
+    private VariableUsedAsReferenceVisitor(PsiVariable variable, String... excludeMethodNames) {
+      myVariable = variable;
+      myExcludeMethodNames = excludeMethodNames;
+    }
+
+    @Override
+    public void visitReferenceExpression(PsiReferenceExpression expression) {
+      if (isUsed) {
+        return;
+      }
+      super.visitReferenceExpression(expression);
+      if (expression.resolve() == myVariable) {
+        isUsed = true;
+      }
+    }
+
+    @Override
+    public void visitMethodCallExpression(PsiMethodCallExpression expression) {
+      if (isUsed) {
+        return;
+      }
+      if (ArrayUtil.contains(expression.getMethodExpression().getReferenceName(), myExcludeMethodNames)) {
+        PsiReferenceExpression qualifier =
+          ObjectUtils.tryCast(expression.getMethodExpression().getQualifierExpression(), PsiReferenceExpression.class);
+        if (qualifier != null && qualifier.resolve() == myVariable) {
+          return;
+        }
+      }
+      super.visitMethodCallExpression(expression);
+    }
+  }
+
  private static class ReplaceWithSingleMapOperation implements LocalQuickFix {
    private final String myMethodName;
    private final SmartPsiElementPointer<PsiMethodCallExpression> myCallPointer;
@@ -237,16 +303,19 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool {

    @Override
    public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
-      PsiElement conditional = PsiTreeUtil.getParentOfType(descriptor.getStartElement(), PsiIfStatement.class, PsiConditionalExpression.class);
-      if(conditional == null) return;
-      MapCheckCondition condition = fromConditional(conditional, true);
+      PsiElement outerElement = PsiTreeUtil.getParentOfType(descriptor.getStartElement(), PsiIfStatement.class,
+                                                            PsiConditionalExpression.class, PsiForeachStatement.class);
+      if (outerElement == null) return;
+      MapCondition condition = outerElement instanceof PsiForeachStatement ?
+                               MapLoopCondition.create((PsiForeachStatement)outerElement) :
+                               fromConditional(outerElement, true);
      if(condition == null) return;
      PsiMethodCallExpression call = myCallPointer.getElement();
      if (call == null) return;
      PsiExpressionList argsList = call.getArgumentList();
      PsiExpression[] args = argsList.getExpressions();
      if(args.length == 0) return;
-      if(myMethodName.equals("merge") && args.length != 2) return;
+      if ((myMethodName.equals("merge") || myMethodName.equals("replaceAll")) && args.length != 2) return;
      PsiExpression value = myValuePointer.getElement();
      if (value == null) return;
      PsiElement result = myResultPointer.getElement();
@@ -285,21 +354,34 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool {
        }
        replacement = factory.createExpressionFromText(varName + " -> " + ct.text(value), value);
      } else if (myMethodName.equals("merge")) {
+        MapCheckCondition checkCondition = ObjectUtils.tryCast(condition, MapCheckCondition.class);
+        if (checkCondition == null) return;
        PsiExpression absentValue = args[1];
        String aVar = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName("a", value, true);
        String bVar = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName("b", value, true);
        for(PsiElement e : PsiTreeUtil.collectElements(value, e -> PsiEquivalenceUtil.areElementsEquivalent(e, absentValue))) {
          ct.replace(e, factory.createIdentifier(bVar));
        }
-        for(PsiElement e : PsiTreeUtil.collectElements(value, e -> condition.extractGetCall(e) != null || condition.isValueReference(e))) {
+        for (PsiElement e : PsiTreeUtil
+          .collectElements(value, e -> checkCondition.extractGetCall(e) != null || checkCondition.isValueReference(e))) {
          ct.replace(e, factory.createIdentifier(aVar));
        }
        replacement = factory.createExpressionFromText("("+aVar+","+bVar+") -> "+ct.text(value), value);
-      } else {
+      }
+      else if (myMethodName.equals("replaceAll")) {
+        MapLoopCondition loopCondition = ObjectUtils.tryCast(condition, MapLoopCondition.class);
+        if (loopCondition == null) return;
+        JavaCodeStyleManager styleManager = JavaCodeStyleManager.getInstance(project);
+        String kVar = suggestKeyName(styleManager, loopCondition, value);
+        String vVar = styleManager.suggestUniqueVariableName("v", value, true);
+        replacement = createLambdaForLoopReplacement(factory, kVar, vVar, loopCondition, value, ct);
+        ct.delete(args);
+      }
+      else {
        replacement = ct.markUnchanged(value);
      }
      PsiElement newArg;
-      if(args.length == 2 && !myMethodName.equals("merge")) {
+      if (args.length == 2 && !myMethodName.equals("merge") && !myMethodName.equals("replaceAll")) {
        newArg = ct.replace(args[1], replacement);
      } else {
        newArg = argsList.add(replacement);
@@ -307,18 +389,65 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool {
      if(newArg instanceof PsiLambdaExpression) {
        LambdaCanBeMethodReferenceInspection.replaceLambdaWithMethodReference((PsiLambdaExpression)newArg);
      }
-      if(PsiTreeUtil.isAncestor(conditional, result, true)) {
-        result = ct.replaceAndRestoreComments(conditional, result);
+      if (PsiTreeUtil.isAncestor(outerElement, result, true)) {
+        result = ct.replaceAndRestoreComments(outerElement, result);
      } else {
-        ct.deleteAndRestoreComments(conditional);
+        ct.deleteAndRestoreComments(outerElement);
      }
-      PsiVariable variable = condition.extractDeclaration();
+      PsiVariable variable = condition instanceof MapCheckCondition ? ((MapCheckCondition)condition).extractDeclaration() : null;
      if (variable != null && !PsiTreeUtil.isAncestor(result, variable, true) && ReferencesSearch.search(variable).findFirst() == null) {
        new CommentTracker().deleteAndRestoreComments(variable);
      }
      CodeStyleManager.getInstance(project).reformat(result);
    }

+    @NotNull
+    private static String suggestKeyName(@NotNull JavaCodeStyleManager codeStyleManager,
+                                         @NotNull MapLoopCondition loopCondition,
+                                         @Nullable PsiElement value) {
+      String nameCandidate = "k";
+      if (!loopCondition.isEntrySet()) {
+        String origName = loopCondition.getIterParam().getName();
+        if (origName != null) {
+          nameCandidate = getNameCandidate(origName);
+          if (origName.equals(nameCandidate)) return nameCandidate;
+        }
+      }
+      return codeStyleManager.suggestUniqueVariableName(nameCandidate, value, true);
+    }
+
+    @NotNull
+    private static PsiExpression createLambdaForLoopReplacement(@NotNull PsiElementFactory factory,
+                                                                @NotNull String kVar,
+                                                                @NotNull String vVar,
+                                                                @NotNull MapLoopCondition loopCondition,
+                                                                @NotNull PsiExpression value,
+                                                                @NotNull CommentTracker tracker) {
+      if (!loopCondition.isEntrySet()) {
+        PsiParameter param = loopCondition.getIterParam();
+        StreamEx.ofTree((PsiElement)value, e -> StreamEx.of(e.getChildren()))
+          .select(PsiReferenceExpression.class)
+          .filter(ref -> param.equals(ref.resolve()))
+          .forEach(ref -> ExpressionUtils.bindReferenceTo(ref, kVar));
+      }
+      else {
+        if (value instanceof PsiMethodCallExpression) {
+          if (loopCondition.isKeyAccess(value)) return factory.createExpressionFromText("(" + kVar + "," + vVar + ") ->" + kVar, value);
+          if (loopCondition.isValueAccess(value)) return factory.createExpressionFromText("(" + kVar + "," + vVar + ") ->" + vVar, value);
+        }
+        Collection<PsiMethodCallExpression> calls = PsiTreeUtil.collectElementsOfType(value, PsiMethodCallExpression.class);
+        for (PsiMethodCallExpression call : calls) {
+          if (loopCondition.isKeyAccess(call)) {
+            PsiReplacementUtil.replaceExpression(call, kVar, new CommentTracker());
+          }
+          else if (loopCondition.isValueAccess(call)) {
+            PsiReplacementUtil.replaceExpression(call, vVar, new CommentTracker());
+          }
+        }
+      }
+      return factory.createExpressionFromText("(" + kVar + "," + vVar + ") ->" + tracker.text(value), value);
+    }
+
    @Nls
    @NotNull
    @Override
@@ -336,6 +465,11 @@ public class Java8MapApiInspection extends AbstractBaseJavaLocalInspectionTool {
    @NotNull
    static ReplaceWithSingleMapOperation fromIf(String methodName, MapCheckCondition condition, PsiExpression value) {
      PsiMethodCallExpression call = condition.getCheckCall();
+      return create(methodName, call, value);
+    }
+
+    @NotNull
+    static ReplaceWithSingleMapOperation create(String methodName, PsiMethodCallExpression call, PsiExpression value) {
      PsiStatement result = PsiTreeUtil.getParentOfType(call, PsiStatement.class);
      LOG.assertTrue(result != null);
      return new ReplaceWithSingleMapOperation(methodName, call, value, result);
@@ -0,0 +1,13 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    String defaultValue = "42";
+    map.put("foo", "bar");
+      map.replaceAll((k, v) -> v + k + defaultValue + "baz");
+  }
+}
@@ -0,0 +1,13 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    String k = "another var";
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+      map.replaceAll((k1, v) -> k);
+  }
+}
@@ -0,0 +1,12 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.Map;
+
+public class Main {
+  public static void main(String[] args) {
+    Map<String, String> vals = new HashMap<>();
+    vals.put("foo", "bar");
+    String defaultValue = "42";
+      vals.replaceAll((k, v) -> defaultValue);
+  }
+}
@@ -0,0 +1,12 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+      map.replaceAll((k, v) -> v);
+  }
+}
@@ -0,0 +1,12 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+      map.replaceAll((k, v) -> k);
+  }
+}
@@ -0,0 +1,13 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    String defaultValue = "42";
+    map.put("foo", "bar");
+      map.replaceAll((k, v) -> k + defaultValue + "baz");
+  }
+}
@@ -0,0 +1,13 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public static void main(String[] args) {
+    Map<String, String> vals = new HashMap<>();
+    vals.put("foo", "bar");
+    String defaultValue = "42";
+      vals.replaceAll((k, v) -> defaultValue);
+  }
+}
@@ -0,0 +1,14 @@
+// "Replace with 'replaceAll' method call" "false"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+    for<caret> (Map.Entry<String, String> entry : map.entrySet()) {
+      map.put(entry.getKey(), String.valueOf(entry));
+    }
+  }
+}
@@ -0,0 +1,14 @@
+// "Replace with 'replaceAll' method call" "false"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+    for<caret> (Map.Entry<String, String> entry : map.entrySet()) {
+      map.put(entry.getKey(), entry.toString());
+    }
+  }
+}
@@ -0,0 +1,15 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    String defaultValue = "42";
+    map.put("foo", "bar");
+    for<caret> (Map.Entry<String, String> entry : map.entrySet()) {
+      map.put(entry.getKey(), entry.getValue() + entry.getKey() + defaultValue + "baz");
+    }
+  }
+}
@@ -0,0 +1,15 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    String k = "another var";
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+    for<caret> (Map.Entry<String, String> entry : map.entrySet()) {
+      map.put(entry.getKey(), k);
+    }
+  }
+}
@@ -0,0 +1,14 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.Map;
+
+public class Main {
+  public static void main(String[] args) {
+    Map<String, String> vals = new HashMap<>();
+    vals.put("foo", "bar");
+    String defaultValue = "42";
+    for<caret> (Entry entry : vals.entrySet()) {
+      vals.put(entry.getKey(), defaultValue);
+    }
+  }
+}
@@ -0,0 +1,14 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+    for<caret> (Map.Entry<String, String> entry : map.entrySet()) {
+      map.put(entry.getKey(), entry.getValue());
+    }
+  }
+}
@@ -0,0 +1,14 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    map.put("foo", "bar");
+    for<caret> (String key : map.keySet()) {
+      map.put(key, key);
+    }
+  }
+}
@@ -0,0 +1,15 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+class Main {
+  public void test() {
+    Map<String, String> map = new HashMap<>();
+    String defaultValue = "42";
+    map.put("foo", "bar");
+    for<caret> (String key : map.keySet()) {
+      map.put(key, key + defaultValue + "baz");
+    }
+  }
+}
@@ -0,0 +1,15 @@
+// "Replace with 'replaceAll' method call" "true"
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class Main {
+  public static void main(String[] args) {
+    Map<String, String> vals = new HashMap<>();
+    vals.put("foo", "bar");
+    String defaultValue = "42";
+    for<caret> (String key : vals.keySet()) {
+      vals.put(key, defaultValue);
+    }
+  }
+}
@@ -21,11 +21,14 @@ import com.intellij.psi.*;
 import com.intellij.psi.search.searches.ReferencesSearch;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
+import com.intellij.util.ObjectUtils;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

+import java.util.Map;
+
 public class Java8MigrationUtils {
  @Nullable
  private static Java8MigrationUtils.MapCheckCondition tryExtract(PsiExpression fullCondition,
@@ -159,10 +162,13 @@ public class Java8MigrationUtils {
  }


+  public interface MapCondition {
+  }
+
  /**
   * Class represents check when working with map: there is 2 ways - when there is a value that matches the key, and when value doesn't exists
   */
-  public static class MapCheckCondition {
+  public static class MapCheckCondition implements MapCondition {
    private final @Nullable PsiReferenceExpression myValueReference;
    private final PsiExpression myMapExpression;
    private final PsiExpression myKeyExpression;
@@ -281,4 +287,108 @@ public class Java8MigrationUtils {
      return null;
    }
  }
+
+  /**
+   * This class is used to store enhanced for loop info when iterating {@link Map} keys or entries.
+   * It supports two iteration patterns:
+   * 1. for (String value : map.keySet())
+   * 2. for (Entry entry : map.entrySet())
+   */
+  public static class MapLoopCondition implements MapCondition {
+    private final PsiParameter myIterParam;
+    private final boolean myIsEntrySet;
+    private final PsiReferenceExpression myMapExpression;
+
+    private MapLoopCondition(@NotNull PsiParameter iterParam, boolean isEntrySet, @NotNull PsiReferenceExpression mapExpression) {
+      myIterParam = iterParam;
+      myIsEntrySet = isEntrySet;
+      myMapExpression = mapExpression;
+    }
+
+    /**
+     * Try to extract {@link Map#put(Object, Object)} call from for each statement body.
+     * This call is expected to be applied to the same map instance as for each iteration, otherwise null will be returned.
+     *
+     * @see ControlFlowUtils#stripBraces(PsiStatement)
+     */
+    @Nullable
+    public PsiMethodCallExpression extractPut(@NotNull PsiForeachStatement statement) {
+      PsiExpressionStatement putStatement =
+        ObjectUtils.tryCast(ControlFlowUtils.stripBraces(statement.getBody()), PsiExpressionStatement.class);
+      if (putStatement == null) return null;
+      PsiMethodCallExpression putCall = extractMapMethodCall(putStatement.getExpression(), "put");
+      if (putCall == null || !isMap(putCall.getMethodExpression().getQualifierExpression())) return null;
+      return putCall;
+    }
+
+    /**
+     * Check if given expression either entry.getKey() call (for entry set based loop) or
+     * loop parameter reference (for key based loop).
+     */
+    public boolean isKeyAccess(@NotNull PsiExpression expression) {
+      if (!myIsEntrySet) {
+        PsiReferenceExpression keyRef = ObjectUtils.tryCast(expression, PsiReferenceExpression.class);
+        return keyRef != null && keyRef.resolve() == myIterParam;
+      }
+      return isParamCall(expression, "getKey");
+    }
+
+    /**
+     * Check if given expression is entry.getValue() call (for entry set based loop).
+     */
+    public boolean isValueAccess(@NotNull PsiExpression expression) {
+      return myIsEntrySet && isParamCall(expression, "getValue");
+    }
+
+    /**
+     * Check if given call is invoked on loop iteration parameter.
+     */
+    public boolean isParamCall(@NotNull PsiMethodCallExpression call) {
+      PsiReferenceExpression qualifier =
+        ObjectUtils.tryCast(call.getMethodExpression().getQualifierExpression(), PsiReferenceExpression.class);
+      return qualifier != null && qualifier.resolve() == myIterParam;
+    }
+
+    public PsiParameter getIterParam() {
+      return myIterParam;
+    }
+
+    public boolean isEntrySet() {
+      return myIsEntrySet;
+    }
+
+    private boolean isParamCall(@NotNull PsiExpression expression, @NotNull String expectedName) {
+      PsiMethodCallExpression call = ObjectUtils.tryCast(expression, PsiMethodCallExpression.class);
+      if (call == null) return false;
+      String name = call.getMethodExpression().getReferenceName();
+      return expectedName.equals(name) && isParamCall(call);
+    }
+
+    private boolean isMap(@Nullable PsiElement element) {
+      return element != null && PsiEquivalenceUtil.areElementsEquivalent(myMapExpression, element);
+    }
+
+    /**
+     * Create {@link MapLoopCondition} from enhanced for statement.
+     * Loop condition instance created only if iteration is done using {@link Map#keySet()} or {@link Map#entrySet()} methods.
+     */
+    @Nullable
+    public static MapLoopCondition create(@NotNull PsiForeachStatement statement) {
+      PsiExpression iteratedValue = statement.getIteratedValue();
+      PsiParameter iterParam = statement.getIterationParameter();
+      PsiMethodCallExpression iterCall = extractMapMethodCall(iteratedValue, "keySet");
+      if (iterCall != null) return create(iterParam, false, iterCall.getMethodExpression().getQualifierExpression());
+
+      iterCall = extractMapMethodCall(iteratedValue, "entrySet");
+      if (iterCall != null) return create(iterParam, true, iterCall.getMethodExpression().getQualifierExpression());
+
+      return null;
+    }
+
+    private static MapLoopCondition create(@NotNull PsiParameter iterParam, boolean isEntrySet, @Nullable PsiExpression qualifier) {
+      PsiReferenceExpression ref = ObjectUtils.tryCast(qualifier, PsiReferenceExpression.class);
+      if (ref == null) return null;
+      return new MapLoopCondition(iterParam, isEntrySet, ref);
+    }
+  }
 }