From 5c96c794edba67bd0086eb184c766812ca4192bb Mon Sep 17 00:00:00 2001
From: Bas Leijdekkers <basleijdekkers@gmail.com>
Date: Mon, 5 Aug 2024 10:48:31 +0200
Subject: [PATCH] Java: fix one more case where "invert if" loses comment
 (IDEA-252033)

GitOrigin-RevId: 8bba0c15ca98a31a29ffc8a0d46effca9d96d27f
---
 .../intention/impl/InvertIfConditionAction.java  | 16 ++++++++++++++--
 .../afterCommentAfterThenBlockStatement.java     | 11 +++++++++++
 .../beforeCommentAfterThenBlockStatement.java    | 10 ++++++++++
 3 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/invertIfCondition/afterCommentAfterThenBlockStatement.java
 create mode 100644 java/java-tests/testData/codeInsight/invertIfCondition/beforeCommentAfterThenBlockStatement.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 cfc3a6f42a7f..08c77e7877a5 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
@@ -237,8 +237,8 @@ public final class InvertIfConditionAction extends PsiUpdateModCommandAction<Psi
         element instanceof PsiForeachStatement && flow.getStartOffset(element) + 1 == endOffset) {
       PsiStatement statement = factory.createStatementFromText("continue;", ifStatement);
       statement = (PsiStatement)codeStyle.reformat(statement);
-      ifStatement = addAfterWithinCodeBlock(ifStatement, ct.markUnchanged(thenBranch), ct);
-      Objects.requireNonNull(ifStatement.getThenBranch()).replace(statement);
+      ifStatement = addAfterWithinCodeBlock(ifStatement, thenBranch, ct);
+      ct.replaceAndRestoreComments(Objects.requireNonNull(ifStatement.getThenBranch()), statement);
       return ifStatement;
     }
 
@@ -372,6 +372,18 @@ public final class InvertIfConditionAction extends PsiUpdateModCommandAction<Psi
       ifStatement = (PsiIfStatement)wrapWithCodeBlock(ifStatement);
     }
     addAfter(ifStatement, branch, ct);
+    branch = ifStatement.getThenBranch();
+    if (branch instanceof PsiBlockStatement blockStatement) {
+      PsiCodeBlock codeBlock = blockStatement.getCodeBlock();
+      PsiElement start = codeBlock.getFirstBodyElement();
+      PsiElement end = codeBlock.getLastBodyElement();
+      if (start != null && end != null) {
+        ct.markRangeUnchanged(start, end);
+      }
+    }
+    else {
+      ct.markUnchanged(branch);
+    }
     return ifStatement;
   }
 
diff --git a/java/java-tests/testData/codeInsight/invertIfCondition/afterCommentAfterThenBlockStatement.java b/java/java-tests/testData/codeInsight/invertIfCondition/afterCommentAfterThenBlockStatement.java
new file mode 100644
index 000000000000..c53ffd0c59e5
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/invertIfCondition/afterCommentAfterThenBlockStatement.java
@@ -0,0 +1,11 @@
+// "Invert 'if' condition" "true"
+class D {
+  private void f(char[] buffer) {
+    for (char c : buffer) {
+        if (false) {
+            continue; // THIS COMMENT DISAPPEARS
+        }
+        break;
+    }
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/invertIfCondition/beforeCommentAfterThenBlockStatement.java b/java/java-tests/testData/codeInsight/invertIfCondition/beforeCommentAfterThenBlockStatement.java
new file mode 100644
index 000000000000..6d5b778306b0
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/invertIfCondition/beforeCommentAfterThenBlockStatement.java
@@ -0,0 +1,10 @@
+// "Invert 'if' condition" "true"
+class D {
+  private void f(char[] buffer) {
+    for (char c : buffer) {
+      if (true)<caret> {
+        break;
+      } // THIS COMMENT DISAPPEARS
+    }
+  }
+}
\ No newline at end of file