Java8CollectionRemoveIf: support for-each loop (IDEA-214448)

GitOrigin-RevId: 4ac219e418cc0bef8e1930d14d53256e3c87be6e

java/java-impl/src/com/intellij/codeInspection/java18api/Java8CollectionRemoveIfInspection.java

@@ -3,6 +3,8 @@ package com.intellij.codeInspection.java18api;
 
 import com.intellij.codeInsight.daemon.QuickFixBundle;
 import com.intellij.codeInspection.*;
+import com.intellij.codeInspection.util.ForEachCollectionTraversal;
+import com.intellij.codeInspection.util.IterableTraversal;
 import com.intellij.codeInspection.util.IteratorDeclaration;
 import com.intellij.codeInspection.util.LambdaGenerationUtil;
 import com.intellij.openapi.project.Project;
@@ -15,7 +17,7 @@ import com.intellij.psi.codeStyle.SuggestedNameInfo;
 import com.intellij.psi.codeStyle.VariableKind;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiUtil;
-import com.intellij.util.ObjectUtils;
+import com.intellij.util.ArrayUtil;
 import com.siyeh.ig.psiutils.CommentTracker;
 import com.siyeh.ig.psiutils.ControlFlowUtils;
 import com.siyeh.ig.psiutils.VariableAccessUtils;
@@ -23,6 +25,8 @@ import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import static com.intellij.util.ObjectUtils.tryCast;
+
 public class Java8CollectionRemoveIfInspection extends AbstractBaseJavaLocalInspectionTool {
   @NotNull
   @Override
@@ -33,9 +37,7 @@ public class Java8CollectionRemoveIfInspection extends AbstractBaseJavaLocalInsp
     return new JavaElementVisitor() {
       void handleIteratorLoop(PsiLoopStatement statement, PsiJavaToken endToken, IteratorDeclaration declaration) {
         if (endToken == null || declaration == null || !declaration.isCollection()) return;
-        PsiStatement body = statement.getBody();
-        if(!(body instanceof PsiBlockStatement)) return;
-        PsiStatement[] statements = ((PsiBlockStatement)body).getCodeBlock().getStatements();
+        PsiStatement[] statements = ControlFlowUtils.unwrapBlock(statement.getBody());
         if (statements.length == 2 && statements[1] instanceof PsiIfStatement) {
           PsiVariable element = declaration.getNextElementVariable(statements[0]);
           if (element == null) return;
@@ -79,16 +81,15 @@ public class Java8CollectionRemoveIfInspection extends AbstractBaseJavaLocalInsp
       }
 
       @Nullable
-      private PsiExpression checkAndExtractCondition(IteratorDeclaration declaration,
-                                                     PsiIfStatement ifStatement) {
+      private PsiExpression checkAndExtractCondition(IterableTraversal traversal, PsiIfStatement ifStatement) {
         PsiExpression condition = ifStatement.getCondition();
         if (condition == null || ifStatement.getElseBranch() != null) return null;
         PsiStatement thenStatement = ControlFlowUtils.stripBraces(ifStatement.getThenBranch());
         if (!(thenStatement instanceof PsiExpressionStatement)) return null;
-        if (!declaration.isIteratorMethodCall(((PsiExpressionStatement)thenStatement).getExpression(), "remove")) return null;
+        if (!traversal.isRemoveCall(((PsiExpressionStatement)thenStatement).getExpression())) return null;
         if (!LambdaGenerationUtil.canBeUncheckedLambda(condition)) return null;
-        PsiReferenceExpression iterable = ObjectUtils.tryCast(PsiUtil.skipParenthesizedExprDown(declaration.getIterable()), PsiReferenceExpression.class);
-        PsiVariable iterableVariable = iterable != null ? ObjectUtils.tryCast(iterable.resolve(), PsiVariable.class) : null;
+        PsiReferenceExpression iterable = tryCast(PsiUtil.skipParenthesizedExprDown(traversal.getIterable()), PsiReferenceExpression.class);
+        PsiVariable iterableVariable = iterable != null ? tryCast(iterable.resolve(), PsiVariable.class) : null;
         if (iterableVariable != null && VariableAccessUtils.variableIsUsed(iterableVariable, condition)) return null;
         return condition;
       }
@@ -106,6 +107,20 @@ public class Java8CollectionRemoveIfInspection extends AbstractBaseJavaLocalInsp
         IteratorDeclaration declaration = IteratorDeclaration.fromLoop(statement);
         handleIteratorLoop(statement, statement.getRParenth(), declaration);
       }
+
+      @Override
+      public void visitForeachStatement(PsiForeachStatement statement) {
+        super.visitForeachStatement(statement);
+        ForEachCollectionTraversal traversal = ForEachCollectionTraversal.fromLoop(statement);
+        if (traversal == null) return;
+        PsiIfStatement ifStatement = tryCast(ControlFlowUtils.stripBraces(statement.getBody()), PsiIfStatement.class);
+        if (ifStatement == null) return;
+        PsiExpression condition = checkAndExtractCondition(traversal, ifStatement);
+        if (condition == null) return;
+        PsiJavaToken endToken = statement.getRParenth();
+        if (endToken == null) return;
+        registerProblem(statement, endToken);
+      }
     };
   }
 
@@ -122,51 +137,56 @@ public class Java8CollectionRemoveIfInspection extends AbstractBaseJavaLocalInsp
       PsiElement element = descriptor.getStartElement();
       if(!(element instanceof PsiLoopStatement)) return;
       PsiLoopStatement loop = (PsiLoopStatement)element;
-      IteratorDeclaration declaration;
-      declaration = IteratorDeclaration.fromLoop(loop);
-      if(declaration == null) return;
-      PsiStatement body = loop.getBody();
-      if(!(body instanceof PsiBlockStatement)) return;
-      PsiStatement[] statements = ((PsiBlockStatement)body).getCodeBlock().getStatements();
-      PsiElementFactory factory = JavaPsiFacade.getElementFactory(project);
-      String replacement = null;
+      PsiStatement[] statements = ControlFlowUtils.unwrapBlock(loop.getBody());
+      PsiIfStatement ifStatement = tryCast(ArrayUtil.getLastElement(statements), PsiIfStatement.class);
+      if (ifStatement == null) return;
+      PsiExpression condition = ifStatement.getCondition();
+      if (condition == null) return;
+      String replacement;
       CommentTracker ct = new CommentTracker();
-      if (statements.length == 2 && statements[1] instanceof PsiIfStatement) {
-        PsiVariable variable = declaration.getNextElementVariable(statements[0]);
-        if (variable == null) return;
-        PsiExpression condition = ((PsiIfStatement)statements[1]).getCondition();
-        if (condition == null) return;
-        replacement = generateRemoveIf(declaration, ct, condition, variable.getName());
+      if (loop instanceof PsiForeachStatement) {
+        ForEachCollectionTraversal traversal = ForEachCollectionTraversal.fromLoop((PsiForeachStatement)loop);
+        if (traversal == null || statements.length != 1) return;
+        replacement = generateRemoveIf(traversal, ct, condition, traversal.getParameter().getName());
       }
-      else if (statements.length == 1 && statements[0] instanceof PsiIfStatement){
-        PsiExpression condition = ((PsiIfStatement)statements[0]).getCondition();
-        if (condition == null) return;
-        PsiElement ref = declaration.findOnlyIteratorRef(condition);
-        if(ref != null) {
-          PsiElement call = ref.getParent().getParent();
-          if(!declaration.isIteratorMethodCall(call, "next")) return;
-          PsiType type = ((PsiExpression)call).getType();
-          JavaCodeStyleManager javaCodeStyleManager = JavaCodeStyleManager.getInstance(project);
-          SuggestedNameInfo info = javaCodeStyleManager.suggestVariableName(VariableKind.PARAMETER, null, null, type);
-          if(info.names.length == 0) {
-            info = javaCodeStyleManager.suggestVariableName(VariableKind.PARAMETER, "value", null, type);
-          }
-          String paramName = javaCodeStyleManager.suggestUniqueVariableName(info, condition, true).names[0];
-          ct.replace(call, factory.createIdentifier(paramName));
-          replacement = generateRemoveIf(declaration, ct, condition, paramName);
+      else {
+        IteratorDeclaration declaration = IteratorDeclaration.fromLoop(loop);
+        if (declaration == null) return;
+        switch (statements.length) {
+          case 1:
+            PsiElement ref = declaration.findOnlyIteratorRef(condition);
+            if (ref == null) return;
+            PsiElement call = ref.getParent().getParent();
+            if (!declaration.isIteratorMethodCall(call, "next")) return;
+            PsiType type = ((PsiExpression)call).getType();
+            JavaCodeStyleManager javaCodeStyleManager = JavaCodeStyleManager.getInstance(project);
+            SuggestedNameInfo info = javaCodeStyleManager.suggestVariableName(VariableKind.PARAMETER, null, null, type);
+            if (info.names.length == 0) {
+              info = javaCodeStyleManager.suggestVariableName(VariableKind.PARAMETER, "value", null, type);
+            }
+            String paramName = javaCodeStyleManager.suggestUniqueVariableName(info, condition, true).names[0];
+            ct.replace(call, JavaPsiFacade.getElementFactory(project).createIdentifier(paramName));
+            replacement = generateRemoveIf(declaration, ct, condition, paramName);
+            break;
+          case 2:
+            PsiVariable variable = declaration.getNextElementVariable(statements[0]);
+            if (variable == null) return;
+            replacement = generateRemoveIf(declaration, ct, condition, variable.getName());
+            break;
+          default:
+            return;
         }
+        ct.delete(declaration.getIterator());
       }
-      if (replacement == null) return;
-      ct.delete(declaration.getIterator());
       PsiElement result = ct.replaceAndRestoreComments(loop, replacement);
       LambdaCanBeMethodReferenceInspection.replaceAllLambdasWithMethodReferences(result);
       CodeStyleManager.getInstance(project).reformat(result);
     }
 
     @NotNull
-    private static String generateRemoveIf(IteratorDeclaration declaration, CommentTracker ct,
+    private static String generateRemoveIf(IterableTraversal traversal, CommentTracker ct,
                                            PsiExpression condition, String paramName) {
-      return (declaration.getIterable() == null ? "" : ct.text(declaration.getIterable()) + ".") +
+      return (traversal.getIterable() == null ? "" : ct.text(traversal.getIterable()) + ".") +
              "removeIf(" + paramName + "->" + ct.text(condition) + ");";
     }
   }

java/java-impl/src/com/intellij/codeInspection/util/ForEachCollectionTraversal.java

@@ -0,0 +1,80 @@
+// Copyright 2000-2019 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.codeInspection.util;
+
+import com.intellij.psi.*;
+import com.intellij.psi.util.InheritanceUtil;
+import com.intellij.psi.util.PsiUtil;
+import com.siyeh.ig.callMatcher.CallMatcher;
+import com.siyeh.ig.psiutils.ConstructionUtils;
+import com.siyeh.ig.psiutils.EquivalenceChecker;
+import com.siyeh.ig.psiutils.ExpressionUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import static com.intellij.psi.CommonClassNames.*;
+
+public class ForEachCollectionTraversal extends IterableTraversal {
+  private static final CallMatcher COLLECTION_TO_ARRAY = CallMatcher.anyOf(
+    CallMatcher.instanceCall(JAVA_UTIL_COLLECTION, "toArray").parameterCount(0),
+    CallMatcher.instanceCall(JAVA_UTIL_COLLECTION, "toArray").parameterTypes("T[]"));
+
+  private static final CallMatcher COLLECTION_REMOVE =
+    CallMatcher.instanceCall(JAVA_UTIL_COLLECTION, "remove").parameterTypes(JAVA_LANG_OBJECT);
+
+  private final PsiParameter myParameter;
+
+  ForEachCollectionTraversal(@Nullable PsiExpression iterable, PsiParameter parameter) {
+    super(iterable, true);
+    myParameter = parameter;
+  }
+
+  public PsiParameter getParameter() {
+    return myParameter;
+  }
+
+  @Override
+  public boolean isRemoveCall(PsiElement candidate) {
+    while (candidate instanceof PsiParenthesizedExpression) {
+      candidate = ((PsiParenthesizedExpression)candidate).getExpression();
+    }
+    if (!(candidate instanceof PsiMethodCallExpression)) return false;
+    PsiMethodCallExpression call = (PsiMethodCallExpression)candidate;
+    if (!COLLECTION_REMOVE.test(call)) return false;
+    PsiExpression qualifier = call.getMethodExpression().getQualifierExpression();
+    if (!EquivalenceChecker.getCanonicalPsiEquivalence().expressionsAreEquivalent(qualifier, myIterable)) return false;
+    PsiExpression arg = call.getArgumentList().getExpressions()[0];
+    return ExpressionUtils.isReferenceTo(arg, myParameter);
+  }
+
+  @Nullable
+  public static ForEachCollectionTraversal fromLoop(@NotNull PsiForeachStatement loop) {
+    PsiExpression collection = extractCollectionExpression(loop.getIteratedValue());
+    if (collection == null) return null;
+    PsiType collectionElement = PsiUtil.substituteTypeParameter(collection.getType(), JAVA_UTIL_COLLECTION, 0, false);
+    if (collectionElement == null) return null;
+    PsiParameter parameter = loop.getIterationParameter();
+    if (!parameter.getType().equals(collectionElement)) return null;
+    return new ForEachCollectionTraversal(collection, parameter);
+  }
+
+  private static PsiExpression extractCollectionExpression(PsiExpression expression) {
+    expression = PsiUtil.skipParenthesizedExprDown(expression);
+    if (expression instanceof PsiMethodCallExpression && COLLECTION_TO_ARRAY.test((PsiMethodCallExpression)expression)) {
+      return PsiUtil.skipParenthesizedExprDown(((PsiMethodCallExpression)expression).getMethodExpression().getQualifierExpression());
+    }
+    if (expression instanceof PsiNewExpression &&
+        ConstructionUtils.isReferenceTo(((PsiNewExpression)expression).getClassReference(), JAVA_UTIL_ARRAY_LIST)) {
+      PsiExpressionList argumentList = ((PsiNewExpression)expression).getArgumentList();
+      if (argumentList != null) {
+        PsiExpression[] args = argumentList.getExpressions();
+        if (args.length == 1 && InheritanceUtil.isInheritor(args[0].getType(), JAVA_UTIL_COLLECTION)) {
+          return args[0];
+        }
+      }
+    }
+    if (expression != null && InheritanceUtil.isInheritor(expression.getType(), JAVA_UTIL_COLLECTION)) {
+      return expression;
+    }
+    return null;
+  }
+}

java/java-impl/src/com/intellij/codeInspection/util/IterableTraversal.java

@@ -0,0 +1,43 @@
+// Copyright 2000-2019 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.codeInspection.util;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiExpression;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Represents a loop which traverses an iterable (e.g. via iterator or for-each loop statement).
+ *
+ * @see IteratorDeclaration
+ * @see ForEachCollectionTraversal
+ */
+public abstract class IterableTraversal {
+  protected final @Nullable PsiExpression myIterable;
+  protected final boolean myCollection;
+
+  IterableTraversal(@Nullable PsiExpression iterable, boolean collection) {
+    myIterable = iterable;
+    myCollection = collection;
+  }
+
+  /**
+   * @return an expression which represent an iterable
+   */
+  @Nullable
+  public final PsiExpression getIterable() {
+    return myIterable;
+  }
+
+  /**
+   * @return true if iterable is known to be a collection
+   */
+  public final boolean isCollection() {
+    return myCollection;
+  }
+
+  /**
+   * @param candidate element to check
+   * @return true if given element is a method call which removes current element from iterable
+   */
+  public abstract boolean isRemoveCall(PsiElement candidate);
+}

java/java-impl/src/com/intellij/codeInspection/util/IteratorDeclaration.java

@@ -36,15 +36,12 @@ import java.util.Objects;
  *
  * @author Tagir Valeev
  */
-public class IteratorDeclaration {
+public class IteratorDeclaration extends IterableTraversal {
   private final @NotNull PsiLocalVariable myIterator;
-  private final @Nullable PsiExpression myIterable;
-  private final boolean myCollection;
 
   private IteratorDeclaration(@NotNull PsiLocalVariable iterator, @Nullable PsiExpression iterable, boolean collection) {
+    super(iterable, collection);
     myIterator = iterator;
-    myIterable = iterable;
-    myCollection = collection;
   }
 
   @NotNull
@@ -52,15 +49,6 @@ public class IteratorDeclaration {
     return myIterator;
   }
 
-  @Nullable
-  public PsiExpression getIterable() {
-    return myIterable;
-  }
-
-  public boolean isCollection() {
-    return myCollection;
-  }
-
   public boolean isHasNextCall(PsiExpression condition) {
     return isIteratorMethodCall(condition, "hasNext");
   }
@@ -87,6 +75,11 @@ public class IteratorDeclaration {
     return method.equals(expression.getReferenceName()) && ExpressionUtils.isReferenceTo(expression.getQualifierExpression(), myIterator);
   }
 
+  @Override
+  public boolean isRemoveCall(PsiElement candidate) {
+    return isIteratorMethodCall(candidate, "remove");
+  }
+
   public PsiVariable getNextElementVariable(PsiStatement statement) {
     PsiLocalVariable var = getDeclaredVariable(statement);
     if (var == null || !isIteratorMethodCall(var.getInitializer(), "next")) return null;

java/java-tests/testData/inspection/java8CollectionRemoveIf/afterForEachRemoveArrayList.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.*;
+
+public class Main {
+    public void removeEmpty(List<String> list) throws Exception {
+        // Copy to avoid CME
+        for(<caret>String item : new ArrayList<>(list)) {
+            if(item.isEmpty()) list.remove(item);
+        }
+    }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/afterForEachRemoveSimple.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.*;
+
+public class Main {
+    public void removeEmpty(List<String> list) throws Exception {
+        for(<caret>String item : list) {
+            // Presumably CopyOnWriteArrayList
+            if(item.isEmpty()) list.remove(item);
+        }
+    }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/afterForEachRemoveToArray.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.*;
+
+public class Main {
+    public void removeEmpty(List<String> list) throws Exception {
+        // Copy to avoid CME
+        for(<caret>String item : list.toArray(new String[list.size()])) {
+            if(item.isEmpty()) list.remove(item);
+        }
+    }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/afterIteratorRemoveInlineNoBrace.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.Iterator;
+import java.util.List;
+
+public class Main {
+  public void testIterator(List<List<String>> data, boolean b) {
+    for(Ite<caret>rator<List<String>> iter = (data.iterator()); iter.hasNext();)
+      if(iter.next().isEmpty() && /* also check the flag */ b)
+        iter.remove();
+  }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/beforeForEachRemoveArrayList.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.*;
+
+public class Main {
+    public void removeEmpty(List<String> list) throws Exception {
+        // Copy to avoid CME
+        for(<caret>String item : new ArrayList<>(list)) {
+            if(item.isEmpty()) list.remove(item);
+        }
+    }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/beforeForEachRemoveSimple.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.*;
+
+public class Main {
+    public void removeEmpty(List<String> list) throws Exception {
+        for(<caret>String item : list) {
+            // Presumably CopyOnWriteArrayList
+            if(item.isEmpty()) list.remove(item);
+        }
+    }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/beforeForEachRemoveToArray.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.*;
+
+public class Main {
+    public void removeEmpty(List<String> list) throws Exception {
+        // Copy to avoid CME
+        for(<caret>String item : list.toArray(new String[list.size()])) {
+            if(item.isEmpty()) list.remove(item);
+        }
+    }
+}

java/java-tests/testData/inspection/java8CollectionRemoveIf/beforeIteratorRemoveInlineNoBrace.java

@@ -0,0 +1,11 @@
+// "Replace the loop with Collection.removeIf" "true"
+import java.util.Iterator;
+import java.util.List;
+
+public class Main {
+  public void testIterator(List<List<String>> data, boolean b) {
+    for(Ite<caret>rator<List<String>> iter = (data.iterator()); iter.hasNext();)
+      if(iter.next().isEmpty() && /* also check the flag */ b)
+        iter.remove();
+  }
+}