From 7addf9e702d1ed1fa5937f29fb08154cbaedc60d Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Wed, 26 Dec 2018 17:56:52 +0700
Subject: [PATCH] IDEA-204717 Can't convert enhanced 'switch' to 'if' if a
 branch's body contains an end-of-line comment

---
 .../impl/quickfix/ConvertSwitchToIfIntention.java    |  9 +++++++--
 .../afterCommentAfterReturnJava12Preview.java        | 10 ++++++++++
 .../beforeCommentAfterReturnJava12Preview.java       | 12 ++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/afterCommentAfterReturnJava12Preview.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/beforeCommentAfterReturnJava12Preview.java

diff --git a/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/ConvertSwitchToIfIntention.java b/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/ConvertSwitchToIfIntention.java
index 1ee8c1ed3567..2956b5911aa9 100644
--- a/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/ConvertSwitchToIfIntention.java
+++ b/java/java-impl/src/com/intellij/codeInsight/daemon/impl/quickfix/ConvertSwitchToIfIntention.java
@@ -338,8 +338,13 @@ public class ConvertSwitchToIfIntention implements IntentionAction {
       if (bodyStatement instanceof PsiBlockStatement) {
         final PsiBlockStatement blockStatement = (PsiBlockStatement)bodyStatement;
         final PsiCodeBlock codeBlock = blockStatement.getCodeBlock();
-        for (PsiStatement statement : codeBlock.getStatements()) {
-          out.append(commentTracker.text(statement));
+        PsiElement start = PsiTreeUtil.skipWhitespacesForward(codeBlock.getFirstBodyElement());
+        PsiElement end = PsiTreeUtil.skipWhitespacesBackward(codeBlock.getLastBodyElement());
+        if (start != null && end != null) {
+          for (PsiElement child = start; child != null; child = child.getNextSibling()) {
+            out.append(commentTracker.text(child));
+            if (child == end) break;
+          }
         }
       }
       else {
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/afterCommentAfterReturnJava12Preview.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/afterCommentAfterReturnJava12Preview.java
new file mode 100644
index 000000000000..f0c3355541bb
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/afterCommentAfterReturnJava12Preview.java
@@ -0,0 +1,10 @@
+// "Replace 'switch' with 'if'" "true"
+class Test {
+  void f(int n, int k) {
+      if (n == 0) {
+          if (k == 0) return; //this comment causes the exception IDEA-204717
+          if (k == 1) {
+          }
+      }
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/beforeCommentAfterReturnJava12Preview.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/beforeCommentAfterReturnJava12Preview.java
new file mode 100644
index 000000000000..568df7077387
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/convertSwitchToIf/beforeCommentAfterReturnJava12Preview.java
@@ -0,0 +1,12 @@
+// "Replace 'switch' with 'if'" "true"
+class Test {
+  void f(int n, int k) {
+    sw<caret>itch (n) {
+      case 0 -> {
+        if (k == 0) return; //this comment causes the exception IDEA-204717
+        if (k == 1) {
+        }
+      }
+    }
+  }
+}
\ No newline at end of file