Drop statements after the first converted switch expression

Fixes #828

Author: timtebeek

src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java

@@ -31,6 +31,7 @@
 import org.openrewrite.staticanalysis.kotlin.KotlinFileChecker;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.openrewrite.Tree.randomId;
 
@@ -59,10 +60,15 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
             @Override
             public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
                 J.Block b = super.visitBlock(block, ctx);
+                AtomicReference<Boolean> newReturn = new AtomicReference<>(false);
                 return b.withStatements(ListUtils.map(b.getStatements(), statement -> {
+                    if (newReturn.get()) {
+                        return null; // Drop statements after the first converted switch expression
+                    }
                     if (statement instanceof J.Switch) {
                         J.Switch sw = (J.Switch) statement;
                         if (canConvertToSwitchExpression(sw)) {
+                            newReturn.set(true);
                             J.SwitchExpression switchExpression = convertToSwitchExpression(sw);
                             return new J.Return(randomId(), sw.getPrefix(), Markers.EMPTY, switchExpression);
                         }

src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java

@@ -330,4 +330,41 @@ String doFormat(String str) {
             )
         );
     }
+
+    @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/828")
+    @Test
+    void removeUnreachableReturnAfterExhaustiveEnumSwitchWithMultipleValues() {
+        rewriteRun(
+            //language=java
+            java(
+                """
+                class Test {
+                    enum Color { RED, GREEN, BLUE }
+
+                    public String getName(Color color) {
+                        switch (color) {
+                            case RED: return "Red";
+                            case GREEN: return "Green";
+                            case BLUE: return "Blue";
+                        }
+                        return "Unknown";
+                    }
+                }
+                """,
+                """
+                class Test {
+                    enum Color { RED, GREEN, BLUE }
+
+                    public String getName(Color color) {
+                        return switch (color) {
+                            case RED -> "Red";
+                            case GREEN -> "Green";
+                            case BLUE -> "Blue";
+                        };
+                    }
+                }
+                """
+            )
+        );
+    }
 }