CollapseIntoLoopAction: handle comments; handle case if x-th iteration equals to the first one

GitOrigin-RevId: 104dc5f3fec65604aef4288e24df1113227f8ac2

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

@@ -110,7 +110,11 @@ public class CollapseIntoLoopAction implements IntentionAction {
       block.addRangeBefore(myStatements.get(0), myStatements.get(myStatementCount - 1), brace);
       PsiElement origBlock = context.getParent();
       JavaCodeStyleManager.getInstance(block.getProject()).shortenClassReferences(origBlock.addBefore(loop, myStatements.get(0)));
-      origBlock.deleteChildRange(myStatements.get(0), myStatements.get(myStatements.size() - 1));
+      CommentTracker ct = new CommentTracker();
+      myLoopElements.forEach(ct::markUnchanged);
+      ct.delete(myStatements.subList(myStatementCount, myStatements.size()).toArray(PsiStatement.EMPTY_ARRAY));
+      ct.insertCommentsBefore(myStatements.get(0));
+      origBlock.deleteChildRange(myStatements.get(0), myStatements.get(myStatementCount - 1));
     }
 
     private String tryCollapseIntoCountingLoop(String varName) {
@@ -221,12 +225,14 @@ public class CollapseIntoLoopAction implements IntentionAction {
       PsiExpression firstIterationExpression = null;
       PsiExpression curIterationExpression = null;
       boolean secondIteration = count == offset;
+      int mismatchedStatements = 0;
       for (int index = 0; index < count; index++) {
         PsiStatement first = statements.get(index);
         PsiStatement cur = statements.get(index + offset);
         EquivalenceChecker.Match match = new TrackingEquivalenceChecker().statementsMatch(first, cur);
         if (match.isExactMismatch()) return false;
         if (match.isExactMatch()) continue;
+        mismatchedStatements++;
         PsiElement leftDiff = match.getLeftDiff();
         PsiElement rightDiff = match.getRightDiff();
         if (!(leftDiff instanceof PsiExpression) || !(rightDiff instanceof PsiExpression)) return false;
@@ -258,6 +264,10 @@ public class CollapseIntoLoopAction implements IntentionAction {
       }
       if (secondIteration) {
         ContainerUtil.addIfNotNull(expressionsToIterate, firstIterationExpression);
+      } else {
+        if (mismatchedStatements != expressionsToReplace.size()) {
+          return false;
+        }
       }
       ContainerUtil.addIfNotNull(expressionsToIterate, curIterationExpression);
       return true;

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

@@ -0,0 +1,13 @@
+import java.util.Arrays;
+
+// "Collapse into loop" "true"
+class X {
+    void test() {
+        for (String s : Arrays.asList("Hello", "Hello1", "Hello2")) {
+            System.out.println(s);
+        }
+        System.out.println("Hello");
+        System.out.println("Hello");
+        System.out.println("Hello");
+    }
+}

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

@@ -0,0 +1,14 @@
+import java.util.Arrays;
+
+// "Collapse into loop" "true"
+class X {
+    void test() {
+        // 1
+        for (String s : Arrays.asList("Hello", "Hello" +/*2.1*/"world", "Hello1")) {
+            System.out.println(s);
+        }/*2.0*/
+        // 2
+        // 3
+        // 4
+    }
+}

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

@@ -0,0 +1,11 @@
+// "Collapse into loop" "true"
+class X {
+    void test() {
+        <caret>System.out.println("Hello");
+        System.out.println("Hello1");
+        System.out.println("Hello2");
+        System.out.println("Hello");
+        System.out.println("Hello");
+        System.out.println("Hello");
+    }
+}

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

@@ -0,0 +1,12 @@
+// "Collapse into loop" "true"
+class X {
+    void test() {
+        // 1
+        <caret>System.out.println("Hello");
+        // 2
+        System.out./*2.0*/println("Hello"+/*2.1*/"world");
+        // 3
+        System.out.println("Hello1");
+        // 4
+    }
+}