[StatementSwitchToExpressionSwitch]: for return switch conversion, fix incorrect assumption that last of multiple statements is always a return

PiperOrigin-RevId: 735952238

Author: markhbrady

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

@@ -1675,12 +1675,13 @@ private static String transformReturnOrThrowBlock(
       int offset = transformedBlockBuilder.length();
       transformedBlockBuilder.append(state.getSourceCode(), codeBlockStart, codeBlockEnd);
       transformedBlockBuilder.append("\n}");
-      ReturnTree returnTree = (ReturnTree) getLast(statements);
-      int start = getStartPosition(returnTree);
-      transformedBlockBuilder.replace(
-          offset + start - codeBlockStart,
-          offset + start - codeBlockStart + "return".length(),
-          "yield");
+      int start = getStartPosition(getLast(statements));
+      if (getLast(statements) instanceof ReturnTree) {
+        transformedBlockBuilder.replace(
+            offset + start - codeBlockStart,
+            offset + start - codeBlockStart + "return".length(),
+            "yield");
+      }
     } else if (statements.size() == 1 && statements.get(0).getKind().equals(RETURN)) {
       // For "return x;", we want to take source starting after the "return"
       int unused = extractLhsComments(caseTree, state, transformedBlockBuilder);

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

@@ -2235,6 +2235,76 @@ public int foo(Side side) {
         .doTest();
   }
 
+  @Test
+  public void switchByEnum_multipleStatementsAndTheLastNotReturn_error() {
+
+    // 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) {
+                  case HEART:
+                  case DIAMOND:
+                    return invoke();
+                  case SPADE:
+                    System.out.println("hello");
+                    throw new RuntimeException();
+                  default:
+                    throw new NullPointerException();
+                }
+              }
+            }
+            """)
+        .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 -> invoke();
+                  case SPADE -> {
+                    System.out.println("hello");
+                    throw new RuntimeException();
+                  }
+                  default -> throw new NullPointerException();
+                };
+              }
+            }
+            """)
+        .setArgs("-XepOpt:StatementSwitchToExpressionSwitch:EnableReturnSwitchConversion")
+        .setFixChooser(StatementSwitchToExpressionSwitchTest::assertOneFixAndChoose)
+        .doTest();
+  }
+
   @Test
   public void switchByEnum_returnSwitchWithShouldNeverHappen_error() {