Convert switch expressions to arrows where possible

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

@@ -0,0 +1,269 @@
+/*
+ * Copyright 2025 the original author or authors.
+ * <p>
+ * Licensed under the Moderne Source Available License (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://docs.moderne.io/licensing/moderne-source-available-license
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.migrate.lang;
+
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+import org.openrewrite.*;
+import org.openrewrite.internal.ListUtils;
+import org.openrewrite.java.JavaIsoVisitor;
+import org.openrewrite.java.JavaTemplate;
+import org.openrewrite.java.search.UsesJavaVersion;
+import org.openrewrite.java.tree.*;
+import org.openrewrite.marker.Markers;
+import org.openrewrite.staticanalysis.groovy.GroovyFileChecker;
+import org.openrewrite.staticanalysis.kotlin.KotlinFileChecker;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.openrewrite.Tree.randomId;
+
+@Value
+@EqualsAndHashCode(callSuper = false)
+public class SwitchCaseReturnsToSwitchExpression extends Recipe {
+    @Override
+    public String getDisplayName() {
+        return "Convert switch cases where every case returns into a returned switch expression";
+    }
+
+    @Override
+    public String getDescription() {
+        return "Switch statements where each case returns a value can be converted to a switch expression that returns the value directly. " +
+                "This is only applicable for Java 14 and later.";
+    }
+
+    @Override
+    public TreeVisitor<?, ExecutionContext> getVisitor() {
+        TreeVisitor<?, ExecutionContext> preconditions = Preconditions.and(
+                new UsesJavaVersion<>(14),
+                Preconditions.not(new KotlinFileChecker<>()),
+                Preconditions.not(new GroovyFileChecker<>())
+        );
+        return Preconditions.check(preconditions, new JavaIsoVisitor<ExecutionContext>() {
+            @Override
+            public J.Switch visitSwitch(J.Switch switch_, ExecutionContext ctx) {
+                J.Switch sw = super.visitSwitch(switch_, ctx);
+
+                // Check if this switch is the only statement in its parent block
+                Cursor parentCursor = getCursor().getParentTreeCursor();
+                if (parentCursor.getValue() instanceof J.Block) {
+                    J.Block parentBlock = parentCursor.getValue();
+                    if (parentBlock.getStatements().size() == 1 && 
+                        parentBlock.getStatements().get(0) == sw) {
+
+                        if (canConvertToSwitchExpression(sw)) {
+                            J.SwitchExpression switchExpression = convertToSwitchExpression(sw);
+                            if (switchExpression != null) {
+                                J.Return returnStatement = new J.Return(
+                                        randomId(),
+                                        sw.getPrefix(),
+                                        Markers.EMPTY,
+                                        switchExpression
+                                );
+
+                                // Replace the parent block's content
+                                doAfterVisit(new JavaIsoVisitor<ExecutionContext>() {
+                                    @Override
+                                    public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
+                                        if (block == parentBlock) {
+                                            return block.withStatements(ListUtils.concat(returnStatement, null));
+                                        }
+                                        return super.visitBlock(block, ctx);
+                                    }
+                                });
+                            }
+                        }
+                    }
+                }
+
+                return sw;
+            }
+
+            private boolean canConvertToSwitchExpression(J.Switch switchStatement) {
+                AtomicBoolean allCasesReturn = new AtomicBoolean(true);
+                AtomicBoolean hasDefaultCase = new AtomicBoolean(false);
+                AtomicBoolean isUsingArrows = new AtomicBoolean(true);
+
+                for (Statement statement : switchStatement.getCases().getStatements()) {
+                    if (!(statement instanceof J.Case)) {
+                        allCasesReturn.set(false);
+                        break;
+                    }
+
+                    J.Case caseStatement = (J.Case) statement;
+
+                    // Check for default case
+                    for (J label : caseStatement.getCaseLabels()) {
+                        if (label instanceof J.Identifier && "default".equals(((J.Identifier) label).getSimpleName())) {
+                            hasDefaultCase.set(true);
+                        }
+                    }
+
+                    if (caseStatement.getBody() != null) {
+                        // Arrow case
+                        J body = caseStatement.getBody();
+                        if (body instanceof J.Block && ((J.Block) body).getStatements().size() == 1) {
+                            body = ((J.Block) body).getStatements().get(0);
+                        }
+                        if (!(body instanceof J.Return)) {
+                            allCasesReturn.set(false);
+                        }
+                    } else {
+                        // Colon case
+                        isUsingArrows.set(false);
+                        List<Statement> statements = caseStatement.getStatements();
+                        if (statements.isEmpty() || !isReturnCase(statements)) {
+                            allCasesReturn.set(false);
+                        }
+                    }
+                }
+
+                // We need either a default case or the switch to cover all possible values
+                return allCasesReturn.get() && (hasDefaultCase.get() || SwitchUtils.coversAllPossibleValues(switchStatement));
+            }
+
+            private boolean isReturnCase(List<Statement> statements) {
+                if (statements.isEmpty()) {
+                    return false;
+                }
+
+
+                // Handle block containing a single return
+                if (statements.size() == 1 && statements.get(0) instanceof J.Block) {
+                    J.Block block = (J.Block) statements.get(0);
+                    if (block.getStatements().size() == 1 && block.getStatements().get(0) instanceof J.Return) {
+                        return true;
+                    }
+                }
+
+                // Direct return statement
+                if (statements.size() == 1 && statements.get(0) instanceof J.Return) {
+                    return true;
+                }
+
+                // Return followed by break (unreachable but sometimes present)
+                if (statements.size() == 2 && statements.get(0) instanceof J.Return && statements.get(1) instanceof J.Break) {
+                    return true;
+                }
+
+                return false;
+            }
+
+            private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) {
+                // First pass to determine return type
+                JavaType returnType = null;
+                for (Statement statement : switchStatement.getCases().getStatements()) {
+                    J.Case caseStatement = (J.Case) statement;
+                    if (caseStatement.getBody() != null) {
+                        J body = caseStatement.getBody();
+                        if (body instanceof J.Block && ((J.Block) body).getStatements().size() == 1) {
+                            body = ((J.Block) body).getStatements().get(0);
+                        }
+                        if (body instanceof J.Return) {
+                            J.Return ret = (J.Return) body;
+                            if (ret.getExpression() != null && ret.getExpression().getType() != null) {
+                                returnType = ret.getExpression().getType();
+                                break;
+                            }
+                        }
+                    } else {
+                        Expression returnExpression = extractReturnExpression(caseStatement.getStatements());
+                        if (returnExpression != null && returnExpression.getType() != null) {
+                            returnType = returnExpression.getType();
+                            break;
+                        }
+                    }
+                }
+
+                List<Statement> convertedCases = ListUtils.map(switchStatement.getCases().getStatements(), statement -> {
+                    J.Case caseStatement = (J.Case) statement;
+
+                    if (caseStatement.getBody() != null) {
+                        // Arrow case
+                        J body = caseStatement.getBody();
+                        if (body instanceof J.Block && ((J.Block) body).getStatements().size() == 1) {
+                            body = ((J.Block) body).getStatements().get(0);
+                        }
+                        if (body instanceof J.Return) {
+                            J.Return ret = (J.Return) body;
+                            if (ret.getExpression() != null) {
+                                return caseStatement
+                                        .withBody(ret.getExpression())
+                                        .withType(J.Case.Type.Rule);
+                            }
+                        }
+                    } else {
+                        // Colon case - convert to arrow case
+                        Expression returnExpression = extractReturnExpression(caseStatement.getStatements());
+                        if (returnExpression != null) {
+                            // When converting from colon to arrow syntax, we need to ensure proper spacing
+                            // The space before the arrow is handled by the padding on the case labels
+                            J.Case.Padding padding = caseStatement.getPadding();
+                            JContainer<J> caseLabels = padding.getCaseLabels();
+
+                            // Add space after the last case label to create space before arrow
+                            JContainer<J> updatedLabels = caseLabels.getPadding().withElements(
+                                ListUtils.mapLast(caseLabels.getPadding().getElements(), 
+                                    elem -> elem.withAfter(Space.SINGLE_SPACE))
+                            );
+
+                            return caseStatement
+                                    .withStatements(null)
+                                    .withBody(returnExpression.withPrefix(Space.SINGLE_SPACE))
+                                    .withType(J.Case.Type.Rule)
+                                    .getPadding()
+                                    .withCaseLabels(updatedLabels);
+                        }
+                    }
+
+                    return caseStatement;
+                });
+
+                return new J.SwitchExpression(
+                        randomId(),
+                        Space.SINGLE_SPACE,
+                        Markers.EMPTY,
+                        switchStatement.getSelector(),
+                        switchStatement.getCases().withStatements(convertedCases),
+                        returnType
+                );
+            }
+
+
+            private Expression extractReturnExpression(List<Statement> statements) {
+                if (statements.isEmpty()) {
+                    return null;
+                }
+
+                // Handle block containing a single return
+                if (statements.size() == 1 && statements.get(0) instanceof J.Block) {
+                    J.Block block = (J.Block) statements.get(0);
+                    if (block.getStatements().size() == 1 && block.getStatements().get(0) instanceof J.Return) {
+                        return ((J.Return) block.getStatements().get(0)).getExpression();
+                    }
+                }
+
+                // Direct return statement
+                if (statements.size() >= 1 && statements.get(0) instanceof J.Return) {
+                    return ((J.Return) statements.get(0)).getExpression();
+                }
+
+                return null;
+            }
+        });
+    }
+}

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

@@ -0,0 +1,127 @@
+/*
+ * Copyright 2025 the original author or authors.
+ * <p>
+ * Licensed under the Moderne Source Available License (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://docs.moderne.io/licensing/moderne-source-available-license
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.migrate.lang;
+
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+import org.openrewrite.*;
+import org.openrewrite.internal.ListUtils;
+import org.openrewrite.java.JavaIsoVisitor;
+import org.openrewrite.java.search.UsesJavaVersion;
+import org.openrewrite.java.tree.*;
+import org.openrewrite.marker.Markers;
+import org.openrewrite.staticanalysis.groovy.GroovyFileChecker;
+import org.openrewrite.staticanalysis.kotlin.KotlinFileChecker;
+
+import java.util.List;
+
+import static org.openrewrite.Tree.randomId;
+
+@Value
+@EqualsAndHashCode(callSuper = false)
+public class SwitchExpressionYieldToArrow extends Recipe {
+    @Override
+    public String getDisplayName() {
+        return "Convert switch expression colon case to arrow";
+    }
+
+    @Override
+    public String getDescription() {
+        return "Convert switch expressions with colon cases and yield statements to arrow syntax. " +
+                "This is only applicable for Java 14 and later.";
+    }
+
+    @Override
+    public TreeVisitor<?, ExecutionContext> getVisitor() {
+        TreeVisitor<?, ExecutionContext> preconditions = Preconditions.and(
+                new UsesJavaVersion<>(14),
+                Preconditions.not(new KotlinFileChecker<>()),
+                Preconditions.not(new GroovyFileChecker<>())
+        );
+        return Preconditions.check(preconditions, new JavaIsoVisitor<ExecutionContext>() {
+            @Override
+            public J.SwitchExpression visitSwitchExpression(J.SwitchExpression switchExpression, ExecutionContext ctx) {
+                J.SwitchExpression se = super.visitSwitchExpression(switchExpression, ctx);
+
+                // First check if all cases are either arrow cases or simple yield cases
+                boolean hasColonCases = false;
+                boolean hasArrowCases = false;
+                boolean hasComplexCases = false;
+
+                for (Statement statement : se.getCases().getStatements()) {
+                    if (statement instanceof J.Case) {
+                        J.Case caseStatement = (J.Case) statement;
+                        if (caseStatement.getType() == J.Case.Type.Rule) {
+                            hasArrowCases = true;
+                        } else if (caseStatement.getType() == J.Case.Type.Statement && caseStatement.getBody() == null) {
+                            hasColonCases = true;
+                            List<Statement> statements = caseStatement.getStatements();
+                            // Check if this is a complex case (more than just a yield)
+                            if (statements.size() != 1 || !(statements.get(0) instanceof J.Yield)) {
+                                hasComplexCases = true;
+                            }
+                        }
+                    }
+                }
+
+                // Don't convert if there are no colon cases, has complex cases, or has a mix of arrow and colon
+                if (!hasColonCases || hasComplexCases || (hasArrowCases && hasColonCases)) {
+                    return se;
+                }
+
+                List<Statement> convertedCases = ListUtils.map(se.getCases().getStatements(), statement -> {
+                    if (!(statement instanceof J.Case)) {
+                        return statement;
+                    }
+
+                    J.Case caseStatement = (J.Case) statement;
+
+                    // Only convert colon cases with yield statements
+                    if (caseStatement.getType() == J.Case.Type.Statement && caseStatement.getBody() == null) {
+                        List<Statement> statements = caseStatement.getStatements();
+
+                        // Check if this case has a single yield statement
+                        if (statements.size() == 1 && statements.get(0) instanceof J.Yield) {
+                            J.Yield yieldStatement = (J.Yield) statements.get(0);
+                            Expression value = yieldStatement.getValue();
+
+                            if (value != null) {
+                                // Add space after the last case label to create space before arrow
+                                J.Case.Padding padding = caseStatement.getPadding();
+                                JContainer<J> caseLabels = padding.getCaseLabels();
+                                JContainer<J> updatedLabels = caseLabels.getPadding().withElements(
+                                    ListUtils.mapLast(caseLabels.getPadding().getElements(), 
+                                        elem -> elem.withAfter(Space.SINGLE_SPACE))
+                                );
+
+                                return caseStatement
+                                        .withStatements(null)
+                                        .withBody(value.withPrefix(Space.SINGLE_SPACE))
+                                        .withType(J.Case.Type.Rule)
+                                        .getPadding()
+                                        .withCaseLabels(updatedLabels);
+                            }
+                        }
+                    }
+
+                    return caseStatement;
+                });
+
+                return se.withCases(se.getCases().withStatements(convertedCases));
+            }
+        });
+    }
+}