CollapseIntoLoopAction: support no-selection mode; disable if code can break the outer loop

GitOrigin-RevId: 87c99d17438c8db668e936f379f2edabde692c80

java/java-impl/src/com/intellij/codeInsight/intention/impl/CollapseIntoLoopAction.java

@@ -13,8 +13,11 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.JavaCodeStyleManager;
 import com.intellij.psi.codeStyle.VariableKind;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.intellij.util.containers.ContainerUtil;
+import com.siyeh.ig.psiutils.ControlFlowUtils;
 import com.siyeh.ig.psiutils.EquivalenceChecker;
 import com.siyeh.ig.psiutils.VariableNameGenerator;
 import one.util.streamex.MoreCollectors;
@@ -41,28 +44,14 @@ public class CollapseIntoLoopAction implements IntentionAction {
     return getText();
   }
 
-  private static List<PsiStatement> extractStatements(Editor editor, PsiFile file) {
-    if (!(file instanceof PsiJavaFile) || !PsiUtil.isLanguageLevel5OrHigher(file)) return Collections.emptyList();
-    SelectionModel model = editor.getSelectionModel();
-    int startOffset = model.getSelectionStart();
-    int endOffset = model.getSelectionEnd();
-    PsiElement[] elements = CodeInsightUtil.findStatementsInRange(file, startOffset, endOffset);
-    return StreamEx.of(elements)
-      .map(e -> tryCast(e, PsiStatement.class))
-      .collect(MoreCollectors.ifAllMatch(Objects::nonNull, Collectors.toList()))
-      .orElse(Collections.emptyList());
-  }
-
   @Override
   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
-    List<PsiStatement> statements = extractStatements(editor, file);
-    return LoopModel.from(statements) != null;
+    return LoopModel.from(editor, file) != null;
   }
 
   @Override
   public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
-    List<PsiStatement> statements = extractStatements(editor, file);
-    LoopModel model = LoopModel.from(statements);
+    LoopModel model = LoopModel.from(editor, file);
     if (model == null) return;
     model.generate();
   }
@@ -161,13 +150,42 @@ public class CollapseIntoLoopAction implements IntentionAction {
       return "for(" + initial + ";" + condition + ";" + increment + ")";
     }
 
-    static @Nullable LoopModel from(List<PsiStatement> statements) {
+    private static @NotNull List<PsiStatement> extractStatements(PsiFile file, SelectionModel model) {
+      int startOffset = model.getSelectionStart();
+      int endOffset = model.getSelectionEnd();
+      PsiElement[] elements = CodeInsightUtil.findStatementsInRange(file, startOffset, endOffset);
+      return StreamEx.of(elements)
+        .map(e -> tryCast(e, PsiStatement.class))
+        .collect(MoreCollectors.ifAllMatch(LoopModel::isAllowedStatement, Collectors.toList()))
+        .orElse(Collections.emptyList());
+    }
+
+    private static @NotNull List<PsiStatement> extractStatements(PsiFile file, int offset) {
+      PsiElement pos = file.findElementAt(offset);
+      PsiStatement statement = PsiTreeUtil.getParentOfType(pos, PsiStatement.class, false, PsiMember.class, PsiCodeBlock.class);
+      if (statement == null) return Collections.emptyList();
+      return StreamEx.iterate(statement, LoopModel::isAllowedStatement, 
+                              st -> PsiTreeUtil.getNextSiblingOfType(st, PsiStatement.class)).toList();
+    }
+
+    static @Nullable LoopModel from(Editor editor, PsiFile file) {
+      if (!(file instanceof PsiJavaFile) || !PsiUtil.isLanguageLevel5OrHigher(file)) return null;
+      SelectionModel selectionModel = editor.getSelectionModel();
+      boolean mayTrimTail;
+      List<PsiStatement> statements;
+      if (selectionModel.hasSelection()) {
+        mayTrimTail = false;
+        statements = extractStatements(file, selectionModel);
+      } else {
+        mayTrimTail = true;
+        statements = extractStatements(file, editor.getCaretModel().getOffset());
+      }
       int size = statements.size();
-      if (size <= 1 || size > 1000) return null;
+      if (size <= 1 || size > (mayTrimTail ? 100 : 1000)) return null;
       if (!(statements.get(0).getParent() instanceof PsiCodeBlock)) return null;
       for (int count = 1; count <= size / 2; count++) {
-        if (size % count != 0) continue;
-        LoopModel model = from(statements, count);
+        if (!mayTrimTail && size % count != 0) continue;
+        LoopModel model = from(statements, count, mayTrimTail);
         if (model != null) {
           return model;
         }
@@ -175,52 +193,66 @@ public class CollapseIntoLoopAction implements IntentionAction {
       return null;
     }
 
-    private static @Nullable LoopModel from(List<PsiStatement> statements, int count) {
-      EquivalenceChecker equivalence = EquivalenceChecker.getCanonicalPsiEquivalence();
+    private static @Nullable LoopModel from(List<PsiStatement> statements, int count, boolean mayTrimTail) {
       int size = statements.size();
-      PsiType type = null;
       List<PsiExpression> expressionsToReplace = new ArrayList<>();
       List<PsiExpression> expressionsToIterate = new ArrayList<>();
-      boolean secondIteration = true;
-      for (int offset = count; offset < size; offset += count) {
-        PsiExpression firstIterationExpression = null;
-        PsiExpression curIterationExpression = null;
-        for (int index = 0; index < count; index++) {
-          PsiStatement first = statements.get(index);
-          PsiStatement cur = statements.get(index + offset);
-          EquivalenceChecker.Match match = equivalence.statementsMatch(first, cur);
-          if (match.isExactMismatch()) return null;
-          if (match.isExactMatch()) continue;
-          PsiElement leftDiff = match.getLeftDiff();
-          PsiElement rightDiff = match.getRightDiff();
-          if (!(leftDiff instanceof PsiExpression) || !(rightDiff instanceof PsiExpression)) return null;
-          curIterationExpression = (PsiExpression)rightDiff;
-          firstIterationExpression = (PsiExpression)leftDiff;
-          if (secondIteration) {
-            if (!expressionsToReplace.isEmpty() &&
-                !equivalence.expressionsAreEquivalent(expressionsToReplace.get(0), (PsiExpression)leftDiff)) {
-              return null;
-            }
-            expressionsToReplace.add((PsiExpression)leftDiff);
-          }
-          else {
-            if (!expressionsToReplace.contains(leftDiff)) return null;
-          }
-        }
-        if (secondIteration) {
-          if (firstIterationExpression != null) {
-            expressionsToIterate.add(firstIterationExpression);
-            PsiType expressionType = GenericsUtil.getVariableTypeByExpressionType(firstIterationExpression.getType());
-            if (expressionType == null) return null;
-            type = expressionType;
-          }
-        }
-        secondIteration = false;
-        if (curIterationExpression != null) {
-          expressionsToIterate.add(curIterationExpression);
+      int offset;
+      for (offset = count; offset + count <= size; offset += count) {
+        if (!tryConsumeIteration(statements, count, offset, expressionsToReplace, expressionsToIterate)) {
+          if (!mayTrimTail || offset == count) return null;
+          break;
         }
       }
+      statements = statements.subList(0, offset);
+      PsiType type = expressionsToIterate.isEmpty() ? null : expressionsToIterate.get(0).getType();
       return new LoopModel(expressionsToIterate, expressionsToReplace, statements, count, type);
     }
+
+    private static boolean tryConsumeIteration(@NotNull List<PsiStatement> statements,
+                                               int count,
+                                               int offset,
+                                               @NotNull List<PsiExpression> expressionsToReplace,
+                                               @NotNull List<PsiExpression> expressionsToIterate) {
+      EquivalenceChecker equivalence = EquivalenceChecker.getCanonicalPsiEquivalence();
+      PsiExpression firstIterationExpression = null;
+      PsiExpression curIterationExpression = null;
+      boolean secondIteration = count == offset;
+      for (int index = 0; index < count; index++) {
+        PsiStatement first = statements.get(index);
+        PsiStatement cur = statements.get(index + offset);
+        EquivalenceChecker.Match match = equivalence.statementsMatch(first, cur);
+        if (match.isExactMismatch()) return false;
+        if (match.isExactMatch()) continue;
+        PsiElement leftDiff = match.getLeftDiff();
+        PsiElement rightDiff = match.getRightDiff();
+        if (!(leftDiff instanceof PsiExpression) || !(rightDiff instanceof PsiExpression)) return false;
+        curIterationExpression = (PsiExpression)rightDiff;
+        firstIterationExpression = (PsiExpression)leftDiff;
+        PsiType curType = curIterationExpression.getType();
+        PsiType firstType = firstIterationExpression.getType();
+        if (curType == null || !curType.equals(firstType)) return false;
+        if (secondIteration) {
+          if (!expressionsToReplace.isEmpty()) {
+            PsiExpression firstExpressionToReplace = expressionsToReplace.get(0);
+            if (!equivalence.expressionsAreEquivalent(firstExpressionToReplace, firstIterationExpression)) return false;
+            if (!firstType.equals(firstExpressionToReplace.getType())) return false;
+          }
+          expressionsToReplace.add(firstIterationExpression);
+        }
+        else {
+          if (!expressionsToReplace.contains(firstIterationExpression)) return false;
+        }
+      }
+      if (secondIteration) {
+        ContainerUtil.addIfNotNull(expressionsToIterate, firstIterationExpression);
+      }
+      ContainerUtil.addIfNotNull(expressionsToIterate, curIterationExpression);
+      return true;
+    }
+
+    private static boolean isAllowedStatement(PsiStatement st) {
+      return st != null && !ControlFlowUtils.statementContainsNakedBreak(st) && !ControlFlowUtils.statementContainsNakedContinue(st);
+    }
   }
 }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/afterBreakOuterLoopLabel.java

@@ -0,0 +1,13 @@
+// "Collapse into loop" "true"
+class X {
+  void test() {
+    LOOP:
+    for (int i = 0; i < 10; i++) {
+        for (int j = 2; j < 6; j++) {
+            if (i % j == 0) break LOOP;
+        }
+    }
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/afterNoSelection.java

@@ -0,0 +1,13 @@
+import java.util.Arrays;
+
+// "Collapse into loop" "true"
+class X {
+  void test() {
+      for (String s : Arrays.asList("foo", "bar")) {
+          foo(s);
+      }
+      foo(123);
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/afterThreeLinesNoSelection.java

@@ -0,0 +1,10 @@
+// "Collapse into loop" "true"
+class X {
+  void test(int[] data) {
+      for (int i = 0; i < 6; i++) {
+          System.out.print("data["+ i +"]");
+          System.out.println("=");
+          System.out.println(data[i]);
+      }
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/afterThreeLinesNoSelection3.java

@@ -0,0 +1,11 @@
+// "Collapse into loop" "true"
+class X {
+  void test(int[] data) {
+      for (int i = 0; i < 6; i++) {
+          System.out.print("data["+ i +"]");
+          System.out.println("=");
+          System.out.println(data[i]);
+      }
+      System.out.print("data["+6+"]");
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeBreakOuterLoop.java

@@ -0,0 +1,13 @@
+// "Collapse into loop" "false"
+class X {
+  void test() {
+    for (int i = 0; i < 10; i++) {
+      <selection>if (i % 2 == 0) break;
+      if (i % 3 == 0) break;
+      if (i % 4 == 0) break;
+      if (i % 5 == 0) break;</selection>
+    }
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeBreakOuterLoopLabel.java

@@ -0,0 +1,14 @@
+// "Collapse into loop" "true"
+class X {
+  void test() {
+    LOOP:
+    for (int i = 0; i < 10; i++) {
+      <selection>if (i % 2 == 0) break LOOP;
+      if (i % 3 == 0) break LOOP;
+      if (i % 4 == 0) break LOOP;
+      if (i % 5 == 0) break LOOP;</selection>
+    }
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeContinueOuterLoop.java

@@ -0,0 +1,13 @@
+// "Collapse into loop" "false"
+class X {
+  void test() {
+    for (int i = 0; i < 10; i++) {
+      <selection>if (i % 2 == 0) continue;
+      if (i % 3 == 0) continue;
+      if (i % 4 == 0) continue;
+      if (i % 5 == 0) continue;</selection>
+    }
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeDifferentTypes.java

@@ -0,0 +1,10 @@
+// "Collapse into loop" "false"
+class X {
+  void test() {
+    <selection>foo("foo");
+    foo("bar");
+    foo(123);</selection>
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeNoSelection.java

@@ -0,0 +1,10 @@
+// "Collapse into loop" "true"
+class X {
+  void test() {
+    <caret>foo("foo");
+    foo("bar");
+    foo(123);
+  }
+  
+  void foo(Object obj) {}
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeThreeLinesNoSelection.java

@@ -0,0 +1,23 @@
+// "Collapse into loop" "true"
+class X {
+  void test(int[] data) {
+    <caret>System.out.print("data["+0+"]");
+    System.out.println("=");
+    System.out.println(data[0]);
+    System.out.print("data["+1+"]");
+    System.out.println("=");
+    System.out.println(data[1]);
+    System.out.print("data["+2+"]");
+    System.out.println("=");
+    System.out.println(data[2]);
+    System.out.print("data["+3+"]");
+    System.out.println("=");
+    System.out.println(data[3]);
+    System.out.print("data["+4+"]");
+    System.out.println("=");
+    System.out.println(data[4]);
+    System.out.print("data["+5+"]");
+    System.out.println("=");
+    System.out.println(data[5]);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeThreeLinesNoSelection2.java

@@ -0,0 +1,23 @@
+// "Collapse into loop" "false"
+class X {
+  void test(int[] data) {
+    System.out.print("data["+0+"]");
+    <caret>System.out.println("=");
+    System.out.println(data[0]);
+    System.out.print("data["+1+"]");
+    System.out.println("=");
+    System.out.println(data[1]);
+    System.out.print("data["+2+"]");
+    System.out.println("=");
+    System.out.println(data[2]);
+    System.out.print("data["+3+"]");
+    System.out.println("=");
+    System.out.println(data[3]);
+    System.out.print("data["+4+"]");
+    System.out.println("=");
+    System.out.println(data[4]);
+    System.out.print("data["+5+"]");
+    System.out.println("=");
+    System.out.println(data[5]);
+  }
+}

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/collapseIntoLoop/beforeThreeLinesNoSelection3.java

@@ -0,0 +1,24 @@
+// "Collapse into loop" "true"
+class X {
+  void test(int[] data) {
+    <caret>System.out.print("data["+0+"]");
+    System.out.println("=");
+    System.out.println(data[0]);
+    System.out.print("data["+1+"]");
+    System.out.println("=");
+    System.out.println(data[1]);
+    System.out.print("data["+2+"]");
+    System.out.println("=");
+    System.out.println(data[2]);
+    System.out.print("data["+3+"]");
+    System.out.println("=");
+    System.out.println(data[3]);
+    System.out.print("data["+4+"]");
+    System.out.println("=");
+    System.out.println(data[4]);
+    System.out.print("data["+5+"]");
+    System.out.println("=");
+    System.out.println(data[5]);
+    System.out.print("data["+6+"]");
+  }
+}

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ControlFlowUtils.java

@@ -575,6 +575,11 @@ public class ControlFlowUtils {
     return !containsReturn(codeBlock) && !codeBlockMayCompleteNormally(codeBlock);
   }
 
+  /**
+   * @param statement statement to test
+   * @return true if statement contains a break without a label that could jump outside of the supplied statement
+   */
+  @Contract("null -> false")
   public static boolean statementContainsNakedBreak(PsiStatement statement) {
     if (statement == null) {
       return false;
@@ -584,6 +589,20 @@ public class ControlFlowUtils {
     return breakFinder.breakFound();
   }
 
+  /**
+   * @param statement statement to test
+   * @return true if statement contains a continue without a label that could jump outside of the supplied statement
+   */
+  @Contract("null -> false")
+  public static boolean statementContainsNakedContinue(PsiStatement statement) {
+    if (statement == null) {
+      return false;
+    }
+    final NakedContinueFinder breakFinder = new NakedContinueFinder();
+    statement.accept(breakFinder);
+    return breakFinder.continueFound();
+  }
+
   /**
    * Checks whether the given statement effectively breaks given loop. Returns true
    * if the statement is {@link PsiBreakStatement} having given loop as a target. Also may return
@@ -1026,16 +1045,8 @@ public class ControlFlowUtils {
     }
 
     @Override
-    public void visitElement(@NotNull PsiElement element) {
-      if (m_found) {
-        return;
-      }
-      super.visitElement(element);
-    }
-
-    @Override
-    public void visitReferenceExpression(
-      PsiReferenceExpression expression) {
+    public void visitExpression(PsiExpression expression) {
+      // don't drill down
     }
 
     @Override
@@ -1044,6 +1055,7 @@ public class ControlFlowUtils {
         return;
       }
       m_found = true;
+      stopWalking();
     }
 
     @Override
@@ -1072,6 +1084,47 @@ public class ControlFlowUtils {
     }
   }
 
+  private static class NakedContinueFinder extends JavaRecursiveElementWalkingVisitor {
+    private boolean m_found;
+
+    private boolean continueFound() {
+      return m_found;
+    }
+
+    @Override
+    public void visitExpression(PsiExpression expression) {
+      // don't drill down
+    }
+
+    @Override
+    public void visitContinueStatement(PsiContinueStatement statement) {
+      if (statement.getLabelIdentifier() != null) {
+        return;
+      }
+      m_found = true;
+    }
+
+    @Override
+    public void visitDoWhileStatement(PsiDoWhileStatement statement) {
+      // don't drill down
+    }
+
+    @Override
+    public void visitForStatement(PsiForStatement statement) {
+      // don't drill down
+    }
+
+    @Override
+    public void visitForeachStatement(PsiForeachStatement statement) {
+      // don't drill down
+    }
+
+    @Override
+    public void visitWhileStatement(PsiWhileStatement statement) {
+      // don't drill down
+    }
+  }
+
   private static class SystemExitFinder extends JavaRecursiveElementWalkingVisitor {
 
     private boolean m_found;

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/EquivalenceChecker.java

@@ -485,10 +485,7 @@ public class EquivalenceChecker {
     final Match conditionEq = expressionsMatch(condition1, condition2);
     final Match thenEq = statementsMatch(thenBranch1, thenBranch2);
     final Match elseEq = statementsMatch(elseBranch1, elseBranch2);
-    if (conditionEq == EXACT_MATCH && thenEq == EXACT_MATCH && elseEq == EXACT_MATCH) {
-      return EXACT_MATCH;
-    }
-    return EXACT_MISMATCH;
+    return conditionEq.combine(thenEq).combine(elseEq);
   }
 
   protected Match expressionStatementsMatch(@NotNull PsiExpressionStatement statement1, @NotNull PsiExpressionStatement statement2) {