From adf14b7ded3a316e90ef59cbfa2cbc2b5fa071f0 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Mon, 12 Mar 2018 15:53:28 +0700
Subject: [PATCH] InvertIfConditionAction: better comment handling

---
 .../impl/InvertIfConditionAction.java         | 29 +++++++++++--------
 .../afterContinueWithCommentNoBlock.java      | 11 +++++++
 .../beforeContinueWithCommentNoBlock.java     |  9 ++++++
 3 files changed, 37 insertions(+), 12 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/invertIfCondition/afterContinueWithCommentNoBlock.java
 create mode 100644 java/java-tests/testData/codeInsight/invertIfCondition/beforeContinueWithCommentNoBlock.java

diff --git a/java/java-impl/src/com/intellij/codeInsight/intention/impl/InvertIfConditionAction.java b/java/java-impl/src/com/intellij/codeInsight/intention/impl/InvertIfConditionAction.java
index 4bfed8eb99ac..c7e830af8b22 100644
--- a/java/java-impl/src/com/intellij/codeInsight/intention/impl/InvertIfConditionAction.java
+++ b/java/java-impl/src/com/intellij/codeInsight/intention/impl/InvertIfConditionAction.java
@@ -163,13 +163,14 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     PsiElementFactory factory = JavaPsiFacade.getInstance(ifStatement.getProject()).getElementFactory();
     Project project = ifStatement.getProject();
 
+    CommentTracker ct = new CommentTracker();
     PsiStatement thenBranch = Objects.requireNonNull(ifStatement.getThenBranch());
     PsiStatement elseBranch = ifStatement.getElseBranch();
 
     if (elseBranch != null) {
       elseBranch = (PsiStatement) elseBranch.copy();
-      setElseBranch(ifStatement, thenBranch, flow);
-      ifStatement.getThenBranch().replace(elseBranch);
+      setElseBranch(ifStatement, thenBranch, flow, ct);
+      ct.replaceAndRestoreComments(ifStatement.getThenBranch(), elseBranch);
       return ifStatement;
     }
 
@@ -178,7 +179,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
       ifStatement.setElseBranch(thenBranch);
       PsiStatement statement = factory.createStatementFromText("{}", ifStatement);
       statement = (PsiStatement) codeStyle.reformat(statement);
-      statement = (PsiStatement) ifStatement.getThenBranch().replace(statement);
+      statement = (PsiStatement) ct.replaceAndRestoreComments(ifStatement.getThenBranch(), statement);
       codeStyle.reformat(statement);
       return ifStatement;
     }
@@ -197,12 +198,13 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
         if (firstElement != null && lastElement != null) {
           ifStatement.getParent().addRangeAfter(firstElement, lastElement, ifStatement);
         }
+        ct.markUnchanged(thenBranch);
       } else {
         if (!(thenBranch instanceof PsiReturnStatement)) {
           ifStatement = addAfterWithinCodeBlock(ifStatement, thenBranch);
         }
       }
-      Objects.requireNonNull(ifStatement.getThenBranch()).replace(statement);
+      ct.replaceAndRestoreComments(Objects.requireNonNull(ifStatement.getThenBranch()), statement);
       return ifStatement;
     }
     PsiElement element = flow.getElement(endOffset);
@@ -221,7 +223,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
     if (element instanceof PsiReturnStatement) {
       PsiReturnStatement returnStatement = (PsiReturnStatement) element;
       ifStatement = addAfterWithinCodeBlock(ifStatement, thenBranch);
-      Objects.requireNonNull(ifStatement.getThenBranch()).replace(returnStatement.copy());
+      ct.replaceAndRestoreComments(Objects.requireNonNull(ifStatement.getThenBranch()), returnStatement.copy());
 
       ControlFlow flow2 = buildControlFlow(findCodeBlock(ifStatement));
       if (!ControlFlowUtil.isInstructionReachable(flow2, flow2.getStartOffset(returnStatement), 0)) returnStatement.delete();
@@ -244,7 +246,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
       }
     }
     if (nextUnreachable) {
-      setElseBranch(ifStatement, thenBranch, flow);
+      setElseBranch(ifStatement, thenBranch, flow, ct);
 
       PsiElement first = ifStatement.getNextSibling();
       if (first != null) {
@@ -261,21 +263,24 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
         PsiBlockStatement codeBlock = (PsiBlockStatement) factory.createStatementFromText("{}", ifStatement);
         codeBlock.getCodeBlock().addRange(first, last);
         first.getParent().deleteChildRange(first, last);
-        ifStatement.getThenBranch().replace(codeBlock);
+        ct.replaceAndRestoreComments(ifStatement.getThenBranch(), codeBlock);
       }
       codeStyle.reformat(ifStatement);
       return ifStatement;
     }
 
-    setElseBranch(ifStatement, thenBranch, flow);
+    setElseBranch(ifStatement, thenBranch, flow, ct);
     PsiStatement statement = factory.createStatementFromText("{}", ifStatement);
     statement = (PsiStatement) codeStyle.reformat(statement);
-    statement = (PsiStatement) Objects.requireNonNull(ifStatement.getThenBranch()).replace(statement);
+    statement = (PsiStatement)ct.replaceAndRestoreComments(Objects.requireNonNull(ifStatement.getThenBranch()), statement);
     codeStyle.reformat(statement);
     return ifStatement;
   }
 
-  private static void setElseBranch(PsiIfStatement ifStatement, PsiStatement thenBranch, ControlFlow flow)
+  private static void setElseBranch(PsiIfStatement ifStatement,
+                                    PsiStatement thenBranch,
+                                    ControlFlow flow,
+                                    CommentTracker ct)
     throws IncorrectOperationException {
     if (flow.getEndOffset(ifStatement) == flow.getEndOffset(thenBranch)) {
       final PsiLoopStatement loopStmt = PsiTreeUtil.getParentOfType(ifStatement, PsiLoopStatement.class);
@@ -285,7 +290,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
           final PsiStatement[] statements = ((PsiBlockStatement)body).getCodeBlock().getStatements();
           if (statements.length > 0 && !PsiTreeUtil.isAncestor(statements[statements.length - 1], ifStatement, false) &&
               ArrayUtilRt.find(statements, ifStatement) < 0) {
-            ifStatement.setElseBranch(thenBranch);
+            ifStatement.setElseBranch(ct.markUnchanged(thenBranch));
             return;
           }
         }
@@ -304,7 +309,7 @@ public class InvertIfConditionAction extends PsiElementBaseIntentionAction {
         }
       }
     }
-    ifStatement.setElseBranch(thenBranch);
+    ifStatement.setElseBranch(ct.markUnchanged(thenBranch));
   }
 
   private static PsiStatement wrapWithCodeBlock(@NotNull PsiStatement statement) {
diff --git a/java/java-tests/testData/codeInsight/invertIfCondition/afterContinueWithCommentNoBlock.java b/java/java-tests/testData/codeInsight/invertIfCondition/afterContinueWithCommentNoBlock.java
new file mode 100644
index 000000000000..c758a140372a
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/invertIfCondition/afterContinueWithCommentNoBlock.java
@@ -0,0 +1,11 @@
+// "Invert 'if' condition" "true"
+class A {
+    void f(){
+        while (true) {
+            if (false) //comment
+            {
+                System.out.println();
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/invertIfCondition/beforeContinueWithCommentNoBlock.java b/java/java-tests/testData/codeInsight/invertIfCondition/beforeContinueWithCommentNoBlock.java
new file mode 100644
index 000000000000..cedc3858469f
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/invertIfCondition/beforeContinueWithCommentNoBlock.java
@@ -0,0 +1,9 @@
+// "Invert 'if' condition" "true"
+class A {
+    void f(){
+        while (true) {
+            i<caret>f (true) continue;//comment
+            System.out.println();
+        }
+    }
+}
\ No newline at end of file