Java: Simplify handling the comments by using CommentTracker (IDEA-182669)

java/java-impl/src/com/intellij/codeInspection/intermediaryVariable/ReturnSeparatedFromComputationInspection.java

@@ -5,11 +5,11 @@ import com.intellij.codeInspection.*;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
-import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.controlFlow.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiTypesUtil;
 import com.intellij.refactoring.util.RefactoringUtil;
+import com.siyeh.ig.psiutils.CommentTracker;
 import com.siyeh.ig.psiutils.ExpressionUtils;
 import gnu.trove.THashMap;
 import gnu.trove.THashSet;
@@ -218,49 +218,11 @@ public class ReturnSeparatedFromComputationInspection extends AbstractBaseJavaLo
   }
 
   private static void replaceStatementKeepComments(PsiStatement replacedStatement, PsiReturnStatement returnStatement) {
-    List<PsiComment> keptComments = getComments(replacedStatement);
-    if (!keptComments.isEmpty()) {
-      returnStatement = (PsiReturnStatement)returnStatement.copy();
-      PsiElement lastReturnChild = returnStatement.getLastChild();
-      Project project = returnStatement.getProject();
-      for (PsiComment comment : keptComments) {
-        lastReturnChild = returnStatement.addAfter(comment, lastReturnChild);
-      }
-      CodeStyleManager.getInstance(project).reformat(returnStatement, true);
-    }
-    replacedStatement.replace(returnStatement);
+    new CommentTracker().replaceAndRestoreComments(replacedStatement, returnStatement);
   }
 
   private static void removeElementKeepComments(PsiElement removedElement) {
-    List<PsiComment> keptComments = getComments(removedElement);
-    if (!keptComments.isEmpty()) {
-      PsiComment firstComment = keptComments.get(0);
-      PsiElement lastComment = removedElement.replace(firstComment);
-      PsiElement parent = lastComment.getParent();
-      Project project = parent.getProject();
-      CodeStyleManager styleManager = CodeStyleManager.getInstance(project);
-      styleManager.reformat(lastComment, true);
-      if (keptComments.size() > 1) {
-        for (PsiComment comment : keptComments.subList(1, keptComments.size())) {
-          lastComment = parent.addAfter(comment, lastComment);
-          styleManager.reformat(lastComment, true);
-        }
-      }
-    }
-    else {
-      removedElement.delete();
-    }
-  }
-
-  private static List<PsiComment> getComments(PsiElement commentedElement) {
-    final List<PsiComment> comments = new ArrayList<>();
-    commentedElement.accept(new JavaRecursiveElementVisitor() {
-      @Override
-      public void visitComment(PsiComment comment) {
-        comments.add(comment);
-      }
-    });
-    return comments;
+    new CommentTracker().deleteAndRestoreComments(removedElement);
   }
 
   private static class Mover {

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/returnSeparatedFromComputation/afterForBreakComment.java

@@ -5,10 +5,12 @@ class T {
         for (int i=0; i<a.length; i++) {
             if (n < a[i]) n = a[i];
             if (n > 100) {
-                return /* return 1 */ n /* return 2 */;// at the end 1
+                // at the end 1
+                return /* return 1 */ n /* return 2 */;
             }
             if (n < 0) {
-                return /* return 1 */ 0 /* return 2 */;// at the end 2
+                // at the end 2
+                return /* return 1 */ 0 /* return 2 */;
                 /* inline */
             }
         }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/returnSeparatedFromComputation/afterIfComment.java

@@ -3,16 +3,17 @@ class T {
     String f(String a) {
         String s = a;
         if (s == null) {
-            return ""; // return comment
             // end of line
+            return ""; // return comment
         }
         else if (s.startsWith("@")) {
+            /* inline 1 */
+            /* inline 2 */
             return s.substring(1); // return comment
-            /* inline 1 *//* inline 2 */
         }
         else if (s.startsWith("#")) {
-            return "#"; // return comment
             /* inline */
+            return "#"; // return comment
         }
         return s; // return comment
     }