StatementSwitchToExpressionSwitch: for return switch transformation, …

…retain comments appearing in the switch block before the first case

PiperOrigin-RevId: 695861404

core/src/main/java/com/google/errorprone/bugpatterns/StatementSwitchToExpressionSwitch.java

@@ -78,6 +78,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Predicate;
 import java.util.regex.Pattern;
 import javax.inject.Inject;
 import javax.lang.model.element.ElementKind;
@@ -646,6 +647,7 @@ private static SuggestedFix convertToReturnSwitch(
         state.getTokensForNode(switchTree).stream()
             .flatMap(errorProneToken -> errorProneToken.comments().stream())
             .collect(toImmutableList());
+
     StringBuilder replacementCodeBuilder = new StringBuilder();
     replacementCodeBuilder
         .append("return switch ")
@@ -662,7 +664,12 @@ private static SuggestedFix convertToReturnSwitch(
           transformReturnOrThrowBlock(caseTree, state, getStatements(caseTree));
 
       if (firstCaseInGroup) {
-        groupedCaseCommentsAccumulator = new StringBuilder();
+        groupedCaseCommentsAccumulator =
+            caseIndex == 0
+                ? new StringBuilder(
+                    extractCommentsBeforeFirstCase(switchTree, allSwitchComments).orElse(""))
+                : new StringBuilder();
+
         replacementCodeBuilder.append("\n  ");
         if (!isDefaultCase) {
           replacementCodeBuilder.append("case ");
@@ -988,6 +995,27 @@ private static int extractLhsComments(
     return lhsEnd;
   }
 
+  /**
+   * Extracts any comments appearing within the switch tree before the first case. Comments are
+   * merged into a single string separated by newlines.
+   */
+  private static Optional<String> extractCommentsBeforeFirstCase(
+      SwitchTree switchTree, ImmutableList<ErrorProneComment> allSwitchComments) {
+    // Indexing relative to the start position of the switch statement
+    int switchStart = getStartPosition(switchTree);
+
+    List<? extends CaseTree> cases = switchTree.getCases();
+    if (cases.isEmpty()) {
+      return Optional.empty();
+    }
+    int firstCaseStartIndex = getStartPosition(cases.get(0)) - switchStart;
+
+    String filteredComments =
+        filterAndRenderComments(
+            allSwitchComments, comment -> comment.getPos() < firstCaseStartIndex);
+    return filteredComments.isEmpty() ? Optional.empty() : Optional.of(filteredComments);
+  }
+
   /**
    * Extracts any comments appearing after the specified {@code caseIndex} but before the subsequent
    * case or end of the {@code switchTree}. Comments are merged into a single string separated by
@@ -1010,21 +1038,30 @@ private static Optional<String> extractCommentsAfterCase(
             : getStartPosition(switchTree.getCases().get(caseIndex + 1)) - switchStart;
 
     String filteredComments =
-        allSwitchComments.stream()
-            // Comments after the end of the current case and before the start of the next case
-            .filter(
-                comment ->
-                    comment.getPos() >= caseEndIndex && comment.getPos() < nextCaseStartIndex)
-            .map(ErrorProneComment::getText)
-            // Remove "fall thru" comments
-            .map(commentText -> removeFallThruLines(commentText))
-            // Remove empty comments
-            .filter(commentText -> !commentText.isEmpty())
-            .collect(joining("\n"));
+        filterAndRenderComments(
+            allSwitchComments,
+            comment -> comment.getPos() >= caseEndIndex && comment.getPos() < nextCaseStartIndex);
 
     return filteredComments.isEmpty() ? Optional.empty() : Optional.of(filteredComments);
   }
 
+  /**
+   * Filters comments according to the supplied predicate ({@code commentFilter}), removes
+   * fall-through and empty comments, and renders them into a single string separated by newlines.
+   */
+  private static String filterAndRenderComments(
+      ImmutableList<ErrorProneComment> comments, Predicate<ErrorProneComment> commentFilter) {
+
+    return comments.stream()
+        .filter(comment -> commentFilter.test(comment))
+        .map(ErrorProneComment::getText)
+        // Remove "fall thru" comments
+        .map(commentText -> removeFallThruLines(commentText))
+        // Remove empty comments
+        .filter(commentText -> !commentText.isEmpty())
+        .collect(joining("\n"));
+  }
+
   /**
    * Finds the position in source corresponding to the end of the code block of the supplied {@code
    * caseIndex} within all {@code cases}.

core/src/test/java/com/google/errorprone/bugpatterns/StatementSwitchToExpressionSwitchTest.java

@@ -2107,6 +2107,86 @@ public int foo(Side side) {
         .doTest(BugCheckerRefactoringTestHelper.TestMode.TEXT_MATCH);
   }
 
+  @Test
+  public void switchByEnum_returnSwitchCommentsBeforeFirstCase_errorAndRetained() {
+    assume().that(Runtime.version().feature()).isAtLeast(14);
+
+    // Check correct generated code
+    refactoringHelper
+        .addInputLines(
+            "Test.java",
+            """
+            class Test {
+              enum Side {
+                HEART,
+                SPADE,
+                DIAMOND,
+                CLUB
+              };
+
+              public Test(int foo) {}
+
+              public int invoke() {
+                return 123;
+              }
+
+              public int foo(Side side) {
+                switch (side) {
+                  // Abracadabra
+                  /* foo */ case HEART:
+                  // Card trick
+                  case DIAMOND:
+                    return invoke();
+                  case SPADE:
+                    throw new RuntimeException();
+                  case CLUB:
+                    throw new NullPointerException();
+                }
+                // This should never happen
+                int z = invoke();
+                z++;
+                throw new RuntimeException("Switch was not exhaustive at runtime " + z);
+              }
+            }
+            """)
+        .addOutputLines(
+            "Test.java",
+            """
+            class Test {
+              enum Side {
+                HEART,
+                SPADE,
+                DIAMOND,
+                CLUB
+              };
+
+              public Test(int foo) {}
+
+              public int invoke() {
+                return 123;
+              }
+
+              public int foo(Side side) {
+                return switch (side) {
+                  case HEART, DIAMOND ->
+                      // Abracadabra
+                      /* foo */
+                      // Card trick
+                      invoke();
+                  case SPADE -> throw new RuntimeException();
+                  case CLUB -> throw new NullPointerException();
+                };
+                // This should never happen
+
+              }
+            }
+            """)
+        .setArgs(
+            ImmutableList.of(
+                "-XepOpt:StatementSwitchToExpressionSwitch:EnableReturnSwitchConversion"))
+        .doTest(BugCheckerRefactoringTestHelper.TestMode.TEXT_MATCH);
+  }
+
   @Test
   public void switchByEnum_switchInReturnSwitchWithShouldNeverHappen_error() {
     assume().that(Runtime.version().feature()).isAtLeast(14);