[java-inspections] IDEA-322460 Convert to enhanced switch: put comment into braces

GitOrigin-RevId: fe4e2187916c721211c6ed9d7c1f1456fcf2d817

java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java

@@ -282,7 +282,7 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
 
   //Right part of switch rule (case labels -> result)
   private interface SwitchRuleResult {
-    String generate(CommentTracker ct);
+    String generate(CommentTracker ct, SwitchBranch branch);
   }
 
   private static class ReplaceWithSwitchExpressionFix extends PsiUpdateModCommandQuickFix {
@@ -910,7 +910,7 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
     }
 
     @Override
-    public String generate(CommentTracker ct) {
+    public String generate(CommentTracker ct, SwitchBranch branch) {
       if (myResultStatements.length == 1) {
         PsiStatement first = myResultStatements[0];
         if (first instanceof PsiExpressionStatement || first instanceof PsiBlockStatement || first instanceof PsiThrowStatement) return ct.textWithComments(myResultStatements[0]) + "\n";
@@ -935,7 +935,11 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
         }
         sb.append(text).append("\n");
       }
-      sb.append("\n}");
+      addCommentsUntilNextLabel(ct, branch, sb);
+      if (sb.charAt(sb.length() - 1) != '\n') {
+        sb.append("\n");
+      }
+      sb.append("}");
       return sb.toString();
     }
   }
@@ -946,11 +950,43 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
     private SwitchRuleExpressionResult(@NotNull PsiExpression expression) { myExpression = expression; }
 
     @Override
-    public String generate(CommentTracker ct) {
+    public String generate(CommentTracker ct, SwitchBranch branch) {
       return ct.textWithComments(myExpression) + ";";
     }
   }
 
+  private static void addCommentsUntilNextLabel(CommentTracker ct, SwitchBranch branch, StringBuilder builder) {
+    PsiElement label = ContainerUtil.find(branch.myUsedElements, e -> e instanceof PsiSwitchLabelStatement);
+    if (!(label instanceof PsiSwitchLabelStatement labelStatement)) {
+      return;
+    }
+    PsiSwitchLabelStatement nextLabelStatement = PsiTreeUtil.getNextSiblingOfType(labelStatement, PsiSwitchLabelStatement.class);
+    PsiElement untilComment = null;
+    if (nextLabelStatement != null) {
+      untilComment = PsiTreeUtil.getPrevSiblingOfType(nextLabelStatement, PsiStatement.class);
+    }
+    if (untilComment == null) {
+      PsiElement next = labelStatement.getNextSibling();
+      if (next != null) {
+        while (next.getNextSibling() != null) {
+          next = next.getNextSibling();
+        }
+      }
+      untilComment = next;
+    }
+    if (untilComment != null) {
+      String commentsBefore = ct.commentsBefore(untilComment).stripTrailing();
+      String previousText = builder.toString().stripTrailing();
+      if (previousText.length() > 1 && previousText.charAt(builder.length() - 1) == '\n') {
+        commentsBefore = StringUtil.trimStart(commentsBefore, "\n");
+      }
+      if (!commentsBefore.isEmpty()) {
+        commentsBefore += '\n';
+      }
+      builder.append(commentsBefore);
+    }
+  }
+
   private static final class SwitchBranch {
     final boolean myIsDefault;
     final List<? extends PsiCaseLabelElement> myCaseExpressions;
@@ -1005,7 +1041,7 @@ public class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInsp
       }
       grabCommentsBeforeColon(label, ct, sb);
       sb.append("->");
-      sb.append(myRuleResult.generate(ct));
+      sb.append(myRuleResult.generate(ct, this));
       return sb.toString();
     }
 

java/java-tests/testData/inspection/switchExpressionMigration/afterCommentsInside.java

@@ -0,0 +1,14 @@
+// "Replace with enhanced 'switch' statement" "true-preview"
+import java.util.*;
+
+class CommentsInside {
+
+  void test(int x) {
+      switch (x) {
+          case 0 -> {
+              // nothing to do
+          }
+          case 1 -> System.out.println(x);
+      }
+  }
+}

java/java-tests/testData/inspection/switchExpressionMigration/beforeCommentsInside.java

@@ -0,0 +1,15 @@
+// "Replace with enhanced 'switch' statement" "true-preview"
+import java.util.*;
+
+class CommentsInside {
+
+  void test(int x) {
+    sw<caret>itch (x) {
+      case 0:
+        // nothing to do
+        break;
+      case 1:
+        System.out.println(x);
+    }
+  }
+}