[java-inspections] IDEA-339341 improve process last comments

GitOrigin-RevId: 6313e67d7bd5180a43d20faffb20c207f46452ae

java/java-impl-inspections/src/com/intellij/codeInspection/EnhancedSwitchBackwardMigrationInspection.java

@@ -18,6 +18,7 @@ import it.unimi.dsi.fastutil.ints.IntList;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -210,11 +211,15 @@ public final class EnhancedSwitchBackwardMigrationInspection extends AbstractBas
       IntList caseCounts = new IntArrayList();
       StringJoiner joiner = new StringJoiner("\n");
       boolean addDefaultBranch = mySwitchBlock instanceof PsiSwitchExpression;
-      for (PsiSwitchLabeledRuleStatement rule : rules) {
+      for (int i = 0; i < rules.size(); i++) {
+        PsiSwitchLabeledRuleStatement rule = rules.get(i);
         CommentTracker ct = new CommentTracker();
         branchTrackers.add(ct);
         String generate = collectCommentsBefore(rule, mainCommentTracker);
         generate += generateBranch(rule, ct, switchCopy);
+        if (i == rules.size() - 1) {
+          generate += collectCommentsBefore(body.getRBrace(), mainCommentTracker);
+        }
         PsiCaseLabelElementList labelElementList = rule.getCaseLabelElementList();
         int caseCount = labelElementList == null ? 1 : labelElementList.getElementCount();
         caseCounts.add(caseCount);
@@ -224,6 +229,7 @@ public final class EnhancedSwitchBackwardMigrationInspection extends AbstractBas
         joiner.add(generate);
         mainCommentTracker.markUnchanged(rule);
         addDefaultBranch &= !SwitchUtils.isDefaultLabel(rule);
+
       }
       if (addDefaultBranch) {
         joiner.add("default:throw new java.lang.IllegalArgumentException();");
@@ -250,11 +256,18 @@ public final class EnhancedSwitchBackwardMigrationInspection extends AbstractBas
     }
 
     @NotNull
-    private static String collectCommentsBefore(@NotNull PsiSwitchLabeledRuleStatement rule, @NotNull CommentTracker ct) {
+    private static String collectCommentsBefore(@Nullable PsiElement rule, @NotNull CommentTracker ct) {
+      boolean commentFound = false;
+      if (rule == null) {
+        return "";
+      }
       List<String> lists = new ArrayList<>();
       PsiElement previous = rule.getPrevSibling();
       while (true) {
         if (previous instanceof PsiComment || previous instanceof PsiWhiteSpace) {
+          if (previous instanceof PsiComment) {
+            commentFound = true;
+          }
           lists.add(ct.text(previous));
           previous = previous.getPrevSibling();
         }
@@ -262,6 +275,9 @@ public final class EnhancedSwitchBackwardMigrationInspection extends AbstractBas
           break;
         }
       }
+      if (!commentFound) {
+        return "";
+      }
       Collections.reverse(lists);
       return String.join("", lists);
     }

java/java-tests/testData/inspection/switchExpressionBackwardMigration/afterLastComment.java

@@ -0,0 +1,31 @@
+// "Replace with old style 'switch' statement" "true"
+
+
+import java.util.function.BiFunction;
+
+public class AAA {
+  public static void main(String[] args) {
+    BiFunction<Integer, Integer, Integer> operation = getOperation(1);
+    System.out.println(operation.apply(5, 3));
+  }
+
+  private static BiFunction<Integer, Integer, Integer> getOperation(int operationCode) {
+      switch (operationCode) {
+          case 1:
+              return (a, b) -> a + b;
+          // Addition
+          case 2:
+              return (a, b) -> a - b;
+          // Subtraction
+          case 3:
+              return (a, b) -> a * b;
+          // Multiplication
+          case 4:
+              return (a, b) -> a / b;
+          // Division
+          default:
+              return (a, b) -> 0; // Default case
+      }
+  }
+
+}

java/java-tests/testData/inspection/switchExpressionBackwardMigration/afterSwitchVarDeclarationComments.java

@@ -8,8 +8,7 @@ class SwitchExpressionMigration {
       /*2*/
       /*3*/
       int x;
-      switch (x +/*cond*/ x) {
-          /*5*/
+      switch (x +/*cond*/ x) {/*5*/
           case 1:
               if (true) {
                   x = 0;

java/java-tests/testData/inspection/switchExpressionBackwardMigration/beforeLastComment.java

@@ -0,0 +1,22 @@
+// "Replace with old style 'switch' statement" "true"
+
+
+import java.util.function.BiFunction;
+
+public class AAA {
+  public static void main(String[] args) {
+    BiFunction<Integer, Integer, Integer> operation = getOperation(1);
+    System.out.println(operation.apply(5, 3));
+  }
+
+  private static BiFunction<Integer, Integer, Integer> getOperation(int operationCode) {
+    return sw<caret>itch (operationCode) {
+      case 1 -> (a, b) -> a + b; // Addition
+      case 2 -> (a, b) -> a - b; // Subtraction
+      case 3 -> (a, b) -> a * b; // Multiplication
+      case 4 -> (a, b) -> a / b; // Division
+      default -> (a, b) -> 0; // Default case
+    };
+  }
+
+}