CommentTracker#replaceAndRestoreComments: try to keep comments after the statement (not move them to the front)

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

@@ -274,7 +274,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     if (nextUnreachable) {
       setElseBranch(ifStatement, thenBranch, flow, ct);
 
-      PsiElement first = ifStatement.getNextSibling();
+      PsiElement first = PsiTreeUtil.skipWhitespacesForward(ifStatement);
       if (first != null) {
         PsiElement last = first;
         PsiElement next = last.getNextSibling();

java/java-tests/testData/codeInsight/createSwitch/string_after.java

@@ -3,9 +3,8 @@ class X {
   void m() {
 
       String s = getStr();
-      //keep me
       switch (s) {
-      }
+      }//keep me
 
   }
 

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/afterString.java

@@ -5,10 +5,9 @@ class X {
       //comment4
       //comment6
       //comment7
-      //comment8
       if ("case1".equals(value)) {//comment2
           //comment3
       } else if ("case2".equals(value)) {//comment5
-      }
+      }//comment8
   }
 }

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

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

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

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

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/continue/afterForEachContinue.java

@@ -5,10 +5,9 @@ import java.util.Collection;
 
 public class Test {
   void test(int[] arr) {
-      // comment
       Arrays.stream(arr).forEach(x -> {
           int y = x * 2;
-          if (x > y) return;
+          if (x > y) return; // comment
           System.out.println(x);
       });
   }

java/java-tests/testData/codeInsight/invertIfCondition/afterConditionWithComment.java

@@ -4,9 +4,8 @@ class A {
     public void foo() {
         String value ="not-null";
 
-        // Another comment
         if (value == null) {
-            return;
+            return; // Another comment
         }
         System.out.println(value);
         // Comment gets deleted.

java/java-tests/testData/codeInsight/invertIfCondition/afterContinue3.java

@@ -4,7 +4,6 @@ class B {
         for (int i = 0; i < 10; i++) {
             if (i % 2 == 0) {
                 if (i != 0) {
-
                     System.out.println("!= 0");
                     continue;
                 }
@@ -12,6 +11,7 @@ class B {
                     System.out.println("== 0");
                     continue;
                 }
+
             }
 
             System.out.println("i = " + i);

java/java-tests/testData/codeInsight/invertIfCondition/afterContinue4.java

@@ -4,12 +4,12 @@ class B {
         for (int i = 0; i < 10; i++) {
             if (i % 2 == 0) {
                 if (i != 0) {
-
                     System.out.println("!= 0");
                 }
                 else {
                     System.out.println("== 0");
                 }
+
             }
         }
     }

java/java-tests/testData/codeInsight/invertIfCondition/afterContinueWithCommentNoBlock.java

@@ -2,10 +2,9 @@
 class A {
     void f(){
         while (true) {
-            //comment
             if (false) {
                 System.out.println();
-            }
+            }//comment
         }
     }
 }

java/java-tests/testData/codeInsight/surroundAutoCloseable/Expression_after.java

@@ -2,8 +2,7 @@ import java.io.*;
 
 class C {
     void m(File file) throws IOException {
-        //comment after expr
         try (FileInputStream fileInputStream = new FileInputStream(file)) {
-        }
+        }//comment after expr
     }
 }

java/java-tests/testData/codeInsight/surroundAutoCloseable/MixedUsages_after.java

@@ -5,7 +5,6 @@ import java.util.Arrays;
 
 class C {
     void m(File file) throws IOException {
-        //comment
         int len;
         boolean empty;
         try (FileInputStream fileInputStream = new FileInputStream(file)) {
@@ -20,7 +19,7 @@ class C {
                 empty = false;
             }
             while (read != -1);
-        }
+        }//comment
         System.out.println(len);
         System.out.println(empty);
     }

java/java-tests/testData/inspection/collectionAddAllCanBeReplacedWithConstructor/afterComment.java

@@ -7,8 +7,7 @@ public class Test {
   void test(Object c) {
       //noinspection unchecked
       /*inside*/
-      // after
-      List<String> list = new ArrayList<>((/*cast!*/Collection<String>) c);
-    System.out.println(list);
+      List<String> list = new ArrayList<>((/*cast!*/Collection<String>) c); // after
+      System.out.println(list);
   }
 }

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

@@ -1,12 +1,12 @@
 // Copyright 2000-2018 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.siyeh.ig.psiutils;
 
-import com.intellij.lang.ASTFactory;
-import com.intellij.lang.ASTNode;
+import com.intellij.lang.*;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.source.tree.ChildRole;
 import com.intellij.psi.impl.source.tree.CompositeElement;
 import com.intellij.psi.impl.source.tree.LeafPsiElement;
+import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiTreeUtil;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
@@ -208,6 +208,21 @@ public final class CommentTracker {
    * @return the element which was actually inserted in the tree (either {@code replacement} or its copy)
    */
   public @NotNull PsiElement replaceAndRestoreComments(@NotNull PsiElement element, @NotNull PsiElement replacement) {
+    List<PsiElement> suffix = new ArrayList<>();
+    if (element instanceof PsiStatement) {
+      PsiElement lastChild = element.getLastChild();
+      boolean hasComment = false;
+      while (lastChild instanceof PsiComment || lastChild instanceof PsiWhiteSpace) {
+        hasComment |= lastChild instanceof PsiComment;
+        if (!(lastChild instanceof PsiComment) || !(shouldIgnore((PsiComment)lastChild))) {
+          suffix.add(markUnchanged(lastChild).copy());
+        }
+        lastChild = lastChild.getPrevSibling();
+      }
+      if (!hasComment) {
+        suffix.clear();
+      }
+    }
     PsiElement result = replace(element, replacement);
     PsiElement anchor = PsiTreeUtil
       .getNonStrictParentOfType(result, PsiStatement.class, PsiLambdaExpression.class, PsiVariable.class, PsiNameValuePair.class);
@@ -221,6 +236,24 @@ public final class CommentTracker {
       anchor = anchor.getParent();
     }
     if (anchor == null) anchor = result;
+    if (!suffix.isEmpty()) {
+      Commenter commenter = LanguageCommenters.INSTANCE.forLanguage(result.getLanguage());
+      if (commenter instanceof CodeDocumentationAwareCommenter) {
+        IElementType lineCommentTokenType = ((CodeDocumentationAwareCommenter)commenter).getLineCommentTokenType();
+        if (lineCommentTokenType != null) {
+          PsiElement lastChild = result.getLastChild();
+          if (lastChild instanceof PsiComment && lineCommentTokenType.equals(((PsiComment)lastChild).getTokenType())) {
+            PsiElement nextSibling = result.getNextSibling();
+            if (nextSibling instanceof PsiWhiteSpace) {
+              result.add(nextSibling);
+            } else {
+              result.add(ASTFactory.whitespace("\n").getPsi());
+            }
+          }
+        }
+      }
+      StreamEx.ofReversed(suffix).forEach(result::add);
+    }
     insertCommentsBefore(anchor);
     return result;
   }

plugins/InspectionGadgets/test/com/siyeh/igfixes/controlflow/conditional_expression/Comment.after.java

@@ -4,8 +4,7 @@ class Comment {
   public String get() {
       /*before then*/
       /*after then*/
-      //comment
       if (239 > 42) return "239";
-      else return "42";
+      else return "42";//comment
   }
 }

plugins/InspectionGadgets/test/com/siyeh/igfixes/exceptions/try_multiple_resource/AtFinalKeyword.after.java

@@ -2,11 +2,10 @@ import java.io.*;
 class AtFinalKeyword {
     void foo(File file1, File file2) throws IOException {
         //after first resource var
-        //after end
         try (final FileInputStream in = /*comment in first*/new FileInputStream(file1)) {
             try (FileOutputStream out /*comment in out*/ = new FileOutputStream(file2)) {//comment
                 System.out.println(in + ", " /*inside statement*/ + out); //commnt at sout
             }
-        }
+        }//after end
     }
 }

plugins/InspectionGadgets/testsrc/com/siyeh/ig/fixes/ExtractParameterAsLocalVariableFixTest.java

@@ -60,8 +60,7 @@ public class ExtractParameterAsLocalVariableFixTest extends IGQuickFixesTestCase
       "}",
       "class X {\n" +
       "    void m(String s) {\n" +
-      "        //end of line comment\n" +
-      "        String hello = \"hello\";\n" +
+      "        String hello = \"hello\";//end of line comment\n" +
       "        System.out.println(hello);\n" +
       "    }\n" +
       "}"

plugins/IntentionPowerPak/test/com/siyeh/ipp/trivialif/convert_to_nested_if/NoBraces_after.java

@@ -3,12 +3,11 @@ package com.siyeh.ipp.trivialif.convert_to_nested_if;
 public class X {
   
     boolean m(boolean a, boolean b, boolean c) {
-        //c1
         if (a) {
             if (b) r<caret>eturn true;
             if (c) return true;
             return false;
-        }
+        }//c1
         return false;//c2
     }
 }

plugins/IntentionPowerPak/test/com/siyeh/ipp/trivialif/convert_to_nested_if/StaircaseWithParenthesis_after.java

@@ -7,8 +7,7 @@ public class X {
       if (a > c) if (a < b) if (!bar1(a//comment in bar1
       )) if (!bar2(//comment in bar2
               a)) return true;/*inside nested*///comment after first condition
-//after end
-      return false;
+      return false;//after end
   }
 
   private static boolean bar1(double a)

plugins/IntentionPowerPak/test/com/siyeh/ipp/trivialif/expand_boolean/Incomplete1_after.java

@@ -1,8 +1,7 @@
 class Incomplete1 {
 
   boolean f(boolean a, boolean b, boolean c) {
-      //keep me
       if (a ? b :) return true;
-      else return false;
+      else return false; //keep me
   }
 }

plugins/IntentionPowerPak/test/com/siyeh/ipp/whileloop/replace_do_while_with_while_loop/InfiniteLoop_after.java

@@ -3,10 +3,9 @@ package com.siyeh.ipp.whileloop.replace_do_while_with_while_loop;
 class InfiniteLoop {
   void m() {
       /*before code block*/
-      //after end
       while ((true)) { //comment
           int i = 10;
           System.out.println(i);
-      }
+      } //after end
   }
 }