From bbe01c9392fdcc58adc4cc4bdbc1c0435745a6d5 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 01:05:07 +0200 Subject: [PATCH 1/8] Convert switch cases with returns to returned switch expression --- .../SwitchCaseReturnsToSwitchExpression.java | 268 +++++++++++++++++ .../META-INF/rewrite/java-version-17.yml | 1 + ...itchCaseReturnsToSwitchExpressionTest.java | 271 ++++++++++++++++++ 3 files changed, 540 insertions(+) create mode 100644 src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java create mode 100644 src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java new file mode 100644 index 0000000000..ef66b8f10f --- /dev/null +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -0,0 +1,268 @@ +/* + * Copyright 2025 the original author or authors. + *

+ * 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 + *

+ * https://docs.moderne.io/licensing/moderne-source-available-license + *

+ * 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 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 getVisitor() { + TreeVisitor preconditions = Preconditions.and( + new UsesJavaVersion<>(14), + Preconditions.not(new KotlinFileChecker<>()), + Preconditions.not(new GroovyFileChecker<>()) + ); + return Preconditions.check(preconditions, new JavaIsoVisitor() { + @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() { + @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 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 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 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 caseLabels = padding.getCaseLabels(); + + // Add space after the last case label to create space before arrow + JContainer 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 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; + } + }); + } +} diff --git a/src/main/resources/META-INF/rewrite/java-version-17.yml b/src/main/resources/META-INF/rewrite/java-version-17.yml index 20f9c20bc3..a632c4ff1d 100644 --- a/src/main/resources/META-INF/rewrite/java-version-17.yml +++ b/src/main/resources/META-INF/rewrite/java-version-17.yml @@ -60,6 +60,7 @@ recipeList: newVersion: 1.17.x - org.openrewrite.java.migrate.AddLombokMapstructBinding - org.openrewrite.java.migrate.lang.SwitchCaseAssignmentsToSwitchExpression + - org.openrewrite.java.migrate.lang.SwitchCaseReturnsToSwitchExpression --- type: specs.openrewrite.org/v1beta/recipe diff --git a/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java b/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java new file mode 100644 index 0000000000..39cc4cf827 --- /dev/null +++ b/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java @@ -0,0 +1,271 @@ +/* + * Copyright 2025 the original author or authors. + *

+ * 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 + *

+ * https://docs.moderne.io/licensing/moderne-source-available-license + *

+ * 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 org.junit.jupiter.api.Test; +import org.openrewrite.DocumentExample; +import org.openrewrite.Issue; +import org.openrewrite.test.RecipeSpec; +import org.openrewrite.test.RewriteTest; + +import static org.openrewrite.java.Assertions.java; +import static org.openrewrite.java.Assertions.javaVersion; + +class SwitchCaseReturnsToSwitchExpressionTest implements RewriteTest { + + @Override + public void defaults(RecipeSpec spec) { + spec.recipe(new SwitchCaseReturnsToSwitchExpression()) + .allSources(s -> s.markers(javaVersion(17))); + } + + @DocumentExample + @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/800") + @Test + void convertSimpleSwitchWithReturns() { + rewriteRun( + //language=java + java( + """ + class Test { + String doFormat(String str) { + switch (str) { + case "foo": return "Foo"; + case "bar": return "Bar"; + case null, default: return "Other"; + } + } + } + """, + """ + class Test { + String doFormat(String str) { + return switch (str) { + case "foo" -> "Foo"; + case "bar" -> "Bar"; + case null, default -> "Other"; + }; + } + } + """ + ) + ); + } + + @Test + void convertSwitchWithColonCases() { + rewriteRun( + //language=java + java( + """ + class Test { + int getValue(String str) { + switch (str) { + case "one": + return 1; + case "two": + return 2; + default: + return 0; + } + } + } + """, + """ + class Test { + int getValue(String str) { + return switch (str) { + case "one" -> 1; + case "two" -> 2; + default -> 0; + }; + } + } + """ + ) + ); + } + + @Test + void convertSwitchWithBlocksContainingReturns() { + rewriteRun( + //language=java + java( + """ + class Test { + String process(int value) { + switch (value) { + case 1: { + return "One"; + } + case 2: { + return "Two"; + } + default: { + return "Many"; + } + } + } + } + """, + """ + class Test { + String process(int value) { + return switch (value) { + case 1 -> "One"; + case 2 -> "Two"; + default -> "Many"; + }; + } + } + """ + ) + ); + } + + @Test + void convertSwitchWithArrowCases() { + rewriteRun( + //language=java + java( + """ + class Test { + String format(String str) { + switch (str) { + case "foo" -> { return "Foo"; } + case "bar" -> { return "Bar"; } + default -> { return "Other"; } + } + } + } + """, + """ + class Test { + String format(String str) { + return switch (str) { + case "foo" -> "Foo"; + case "bar" -> "Bar"; + default -> "Other"; + }; + } + } + """ + ) + ); + } + + @Test + void doNotConvertWhenNotAllCasesReturn() { + rewriteRun( + //language=java + java( + """ + class Test { + String process(String str) { + switch (str) { + case "foo": + return "Foo"; + case "bar": + System.out.println("Bar case"); + break; + default: + return "Other"; + } + return "End"; + } + } + """ + ) + ); + } + + @Test + void doNotConvertWhenNoDefaultAndNotExhaustive() { + rewriteRun( + //language=java + java( + """ + class Test { + String format(String str) { + switch (str) { + case "foo": return "Foo"; + case "bar": return "Bar"; + } + return "Not found"; + } + } + """ + ) + ); + } + + @Test + void convertEnumSwitchThatIsExhaustive() { + rewriteRun( + //language=java + java( + """ + class Test { + enum Color { RED, GREEN, BLUE } + + String colorName(Color color) { + switch (color) { + case RED: return "Red"; + case GREEN: return "Green"; + case BLUE: return "Blue"; + } + } + } + """, + """ + class Test { + enum Color { RED, GREEN, BLUE } + + String colorName(Color color) { + return switch (color) { + case RED -> "Red"; + case GREEN -> "Green"; + case BLUE -> "Blue"; + }; + } + } + """ + ) + ); + } + + @Test + void doNotConvertIfNotOnlyStatementInBlock() { + rewriteRun( + //language=java + java( + """ + class Test { + String process(String str) { + switch (str) { + case "foo": + System.out.println("Processing: " + str); + return "Foo"; + case "bar": return "Bar"; + default: return "Other"; + } + } + } + """ + ) + ); + } +} From 4207a6ca2ddb419afeaba7f5bf7a997d71a82818 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 01:17:59 +0200 Subject: [PATCH 2/8] Early bit of polish --- .../SwitchCaseReturnsToSwitchExpression.java | 57 ++++++------ ...itchCaseReturnsToSwitchExpressionTest.java | 90 +++++++++---------- 2 files changed, 70 insertions(+), 77 deletions(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index ef66b8f10f..1a1043ddd2 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -17,6 +17,7 @@ import lombok.EqualsAndHashCode; import lombok.Value; +import org.jspecify.annotations.Nullable; import org.openrewrite.*; import org.openrewrite.internal.ListUtils; import org.openrewrite.java.JavaIsoVisitor; @@ -27,7 +28,6 @@ import org.openrewrite.staticanalysis.kotlin.KotlinFileChecker; import java.util.List; -import java.util.concurrent.atomic.AtomicBoolean; import static org.openrewrite.Tree.randomId; @@ -41,8 +41,7 @@ public String getDisplayName() { @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."; + return "Switch statements where each case returns a value can be converted to a switch expression that returns the value directly."; } @Override @@ -66,25 +65,23 @@ public J.Switch visitSwitch(J.Switch switch_, ExecutionContext ctx) { 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() { - @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); + J.Return returnStatement = new J.Return( + randomId(), + sw.getPrefix(), + Markers.EMPTY, + switchExpression + ); + + // Replace the parent block's content + doAfterVisit(new JavaIsoVisitor() { + @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); + } + }); } } } @@ -93,14 +90,11 @@ public J.Block visitBlock(J.Block block, ExecutionContext ctx) { } private boolean canConvertToSwitchExpression(J.Switch switchStatement) { - AtomicBoolean allCasesReturn = new AtomicBoolean(true); - AtomicBoolean hasDefaultCase = new AtomicBoolean(false); - AtomicBoolean isUsingArrows = new AtomicBoolean(true); + boolean hasDefaultCase = false; for (Statement statement : switchStatement.getCases().getStatements()) { if (!(statement instanceof J.Case)) { - allCasesReturn.set(false); - break; + return false; } J.Case caseStatement = (J.Case) statement; @@ -108,7 +102,7 @@ private boolean canConvertToSwitchExpression(J.Switch switchStatement) { // Check for default case for (J label : caseStatement.getCaseLabels()) { if (label instanceof J.Identifier && "default".equals(((J.Identifier) label).getSimpleName())) { - hasDefaultCase.set(true); + hasDefaultCase = true; } } @@ -119,20 +113,19 @@ private boolean canConvertToSwitchExpression(J.Switch switchStatement) { body = ((J.Block) body).getStatements().get(0); } if (!(body instanceof J.Return)) { - allCasesReturn.set(false); + return false; } } else { // Colon case - isUsingArrows.set(false); List statements = caseStatement.getStatements(); if (statements.isEmpty() || !isReturnCase(statements)) { - allCasesReturn.set(false); + return false; } } } // We need either a default case or the switch to cover all possible values - return allCasesReturn.get() && (hasDefaultCase.get() || SwitchUtils.coversAllPossibleValues(switchStatement)); + return hasDefaultCase || SwitchUtils.coversAllPossibleValues(switchStatement); } private boolean isReturnCase(List statements) { @@ -243,7 +236,7 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { } - private Expression extractReturnExpression(List statements) { + private @Nullable Expression extractReturnExpression(List statements) { if (statements.isEmpty()) { return null; } diff --git a/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java b/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java index 39cc4cf827..d624ccfcff 100644 --- a/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java +++ b/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java @@ -167,51 +167,6 @@ String format(String str) { ); } - @Test - void doNotConvertWhenNotAllCasesReturn() { - rewriteRun( - //language=java - java( - """ - class Test { - String process(String str) { - switch (str) { - case "foo": - return "Foo"; - case "bar": - System.out.println("Bar case"); - break; - default: - return "Other"; - } - return "End"; - } - } - """ - ) - ); - } - - @Test - void doNotConvertWhenNoDefaultAndNotExhaustive() { - rewriteRun( - //language=java - java( - """ - class Test { - String format(String str) { - switch (str) { - case "foo": return "Foo"; - case "bar": return "Bar"; - } - return "Not found"; - } - } - """ - ) - ); - } - @Test void convertEnumSwitchThatIsExhaustive() { rewriteRun( @@ -247,6 +202,51 @@ String colorName(Color color) { ); } + @Test + void doNotConvertWhenNotAllCasesReturn() { + rewriteRun( + //language=java + java( + """ + class Test { + String process(String str) { + switch (str) { + case "foo": + return "Foo"; + case "bar": + System.out.println("Bar case"); + break; + default: + return "Other"; + } + return "End"; + } + } + """ + ) + ); + } + + @Test + void doNotConvertWhenNoDefaultAndNotExhaustive() { + rewriteRun( + //language=java + java( + """ + class Test { + String format(String str) { + switch (str) { + case "foo": return "Foo"; + case "bar": return "Bar"; + } + return "Not found"; + } + } + """ + ) + ); + } + @Test void doNotConvertIfNotOnlyStatementInBlock() { rewriteRun( From 16424e9f4686954eca1dbf8e339ddaaf6f1cf71d Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 01:52:33 +0200 Subject: [PATCH 3/8] Convert switch if it's the last statement, not only statement --- .../SwitchCaseReturnsToSwitchExpression.java | 8 ++--- ...itchCaseReturnsToSwitchExpressionTest.java | 33 +++++++++++++++++++ 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index 1a1043ddd2..961c1da2dc 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -56,13 +56,11 @@ public TreeVisitor getVisitor() { 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 + // Check if this switch is the last 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 (parentBlock.getStatements().get(parentBlock.getStatements().size() - 1) == sw) { if (canConvertToSwitchExpression(sw)) { J.SwitchExpression switchExpression = convertToSwitchExpression(sw); J.Return returnStatement = new J.Return( @@ -77,7 +75,7 @@ public J.Switch visitSwitch(J.Switch switch_, ExecutionContext ctx) { @Override public J.Block visitBlock(J.Block block, ExecutionContext ctx) { if (block == parentBlock) { - return block.withStatements(ListUtils.concat(returnStatement, null)); + return block.withStatements(ListUtils.mapLast(block.getStatements(), last -> returnStatement)); } return super.visitBlock(block, ctx); } diff --git a/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java b/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java index d624ccfcff..d644db1a15 100644 --- a/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java +++ b/src/test/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpressionTest.java @@ -65,6 +65,39 @@ String doFormat(String str) { ); } + @Test + void convertSimpleSwitchWithReturnsAfterOtherStatements() { + rewriteRun( + //language=java + java( + """ + class Test { + String doFormat(String str) { + System.out.println("Formatting: " + str); + switch (str) { + case "foo": return "Foo"; + case "bar": return "Bar"; + case null, default: return "Other"; + } + } + } + """, + """ + class Test { + String doFormat(String str) { + System.out.println("Formatting: " + str); + return switch (str) { + case "foo" -> "Foo"; + case "bar" -> "Bar"; + case null, default -> "Other"; + }; + } + } + """ + ) + ); + } + @Test void convertSwitchWithColonCases() { rewriteRun( From 887ea2ee0c16eb0fb951c66d25a315c6df5d6728 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 01:56:42 +0200 Subject: [PATCH 4/8] Use a more concise `visitBlock` --- .../SwitchCaseReturnsToSwitchExpression.java | 35 ++++++------------- 1 file changed, 10 insertions(+), 25 deletions(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index 961c1da2dc..5c994830ef 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -53,38 +53,23 @@ public TreeVisitor getVisitor() { ); return Preconditions.check(preconditions, new JavaIsoVisitor() { @Override - public J.Switch visitSwitch(J.Switch switch_, ExecutionContext ctx) { - J.Switch sw = super.visitSwitch(switch_, ctx); - - // Check if this switch is the last statement in its parent block - Cursor parentCursor = getCursor().getParentTreeCursor(); - if (parentCursor.getValue() instanceof J.Block) { - J.Block parentBlock = parentCursor.getValue(); - if (parentBlock.getStatements().get(parentBlock.getStatements().size() - 1) == sw) { + public J.Block visitBlock(J.Block block, ExecutionContext ctx) { + J.Block b = super.visitBlock(block, ctx); + return b.withStatements(ListUtils.map(b.getStatements(), statement -> { + if (statement instanceof J.Switch) { + J.Switch sw = (J.Switch) statement; if (canConvertToSwitchExpression(sw)) { J.SwitchExpression switchExpression = convertToSwitchExpression(sw); - J.Return returnStatement = new J.Return( + return new J.Return( randomId(), sw.getPrefix(), Markers.EMPTY, switchExpression ); - - // Replace the parent block's content - doAfterVisit(new JavaIsoVisitor() { - @Override - public J.Block visitBlock(J.Block block, ExecutionContext ctx) { - if (block == parentBlock) { - return block.withStatements(ListUtils.mapLast(block.getStatements(), last -> returnStatement)); - } - return super.visitBlock(block, ctx); - } - }); } } - } - - return sw; + return statement; + })); } private boolean canConvertToSwitchExpression(J.Switch switchStatement) { @@ -207,8 +192,8 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { // Add space after the last case label to create space before arrow JContainer updatedLabels = caseLabels.getPadding().withElements( - ListUtils.mapLast(caseLabels.getPadding().getElements(), - elem -> elem.withAfter(Space.SINGLE_SPACE)) + ListUtils.mapLast(caseLabels.getPadding().getElements(), + elem -> elem.withAfter(Space.SINGLE_SPACE)) ); return caseStatement From dc092de6effe5cd2a3a933173ff87c84184014be Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 02:01:13 +0200 Subject: [PATCH 5/8] Update src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- .../migrate/lang/SwitchCaseReturnsToSwitchExpression.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index 5c994830ef..35957f2874 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -18,7 +18,10 @@ import lombok.EqualsAndHashCode; import lombok.Value; import org.jspecify.annotations.Nullable; -import org.openrewrite.*; +import org.openrewrite.ExecutionContext; +import org.openrewrite.Preconditions; +import org.openrewrite.Recipe; +import org.openrewrite.TreeVisitor; import org.openrewrite.internal.ListUtils; import org.openrewrite.java.JavaIsoVisitor; import org.openrewrite.java.search.UsesJavaVersion; From 64bd095fe6099ee9212dc56c164af1a4fe653a83 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 22:18:44 +0200 Subject: [PATCH 6/8] Some simplifications --- .../SwitchCaseReturnsToSwitchExpression.java | 29 ++----------------- 1 file changed, 3 insertions(+), 26 deletions(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index 35957f2874..7df7b7989a 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -76,22 +76,12 @@ public J.Block visitBlock(J.Block block, ExecutionContext ctx) { } private boolean canConvertToSwitchExpression(J.Switch switchStatement) { - boolean hasDefaultCase = false; - for (Statement statement : switchStatement.getCases().getStatements()) { if (!(statement instanceof J.Case)) { return false; } 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 = true; - } - } - if (caseStatement.getBody() != null) { // Arrow case J body = caseStatement.getBody(); @@ -111,7 +101,7 @@ private boolean canConvertToSwitchExpression(J.Switch switchStatement) { } // We need either a default case or the switch to cover all possible values - return hasDefaultCase || SwitchUtils.coversAllPossibleValues(switchStatement); + return SwitchUtils.coversAllPossibleValues(switchStatement); } private boolean isReturnCase(List statements) { @@ -119,26 +109,13 @@ private boolean isReturnCase(List statements) { 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; - } + return isReturnCase(((J.Block) statements.get(0)).getStatements()); } // 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; + return statements.size() == 1 && statements.get(0) instanceof J.Return; } private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { From 044dda9462e40e77aa1f51ad1420c07060d39685 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 22:33:05 +0200 Subject: [PATCH 7/8] Condense further --- .../SwitchCaseReturnsToSwitchExpression.java | 72 ++++++++----------- 1 file changed, 31 insertions(+), 41 deletions(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index 7df7b7989a..133a03eadd 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -119,34 +119,10 @@ private boolean isReturnCase(List statements) { } 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; - } - } - } + JavaType returnType = extractReturnType(switchStatement); List convertedCases = ListUtils.map(switchStatement.getCases().getStatements(), statement -> { J.Case caseStatement = (J.Case) statement; - if (caseStatement.getBody() != null) { // Arrow case J body = caseStatement.getBody(); @@ -156,9 +132,7 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { 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); + return caseStatement.withBody(ret.getExpression()); } } } else { @@ -166,16 +140,10 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { 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 caseLabels = padding.getCaseLabels(); - - // Add space after the last case label to create space before arrow + JContainer caseLabels = caseStatement.getPadding().getCaseLabels(); JContainer updatedLabels = caseLabels.getPadding().withElements( ListUtils.mapLast(caseLabels.getPadding().getElements(), - elem -> elem.withAfter(Space.SINGLE_SPACE)) - ); - + elem -> elem.withAfter(Space.SINGLE_SPACE))); return caseStatement .withStatements(null) .withBody(returnExpression.withPrefix(Space.SINGLE_SPACE)) @@ -184,10 +152,8 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { .withCaseLabels(updatedLabels); } } - return caseStatement; }); - return new J.SwitchExpression( randomId(), Space.SINGLE_SPACE, @@ -198,14 +164,38 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { ); } + private @Nullable JavaType extractReturnType(J.Switch switchStatement) { + 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) { + return ret.getExpression().getType(); + } + } + } else { + Expression returnExpression = extractReturnExpression(caseStatement.getStatements()); + if (returnExpression != null && returnExpression.getType() != null) { + return returnExpression.getType(); + } + } + } + return null; + } + private @Nullable Expression extractReturnExpression(List statements) { - if (statements.isEmpty()) { + if (statements.size() != 1) { return null; } // Handle block containing a single return - if (statements.size() == 1 && statements.get(0) instanceof J.Block) { + if (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(); @@ -213,7 +203,7 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { } // Direct return statement - if (statements.size() >= 1 && statements.get(0) instanceof J.Return) { + if (statements.get(0) instanceof J.Return) { return ((J.Return) statements.get(0)).getExpression(); } From f9340da5467912309070c2cfa451fbadb7ba3d46 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Sat, 26 Jul 2025 23:17:45 +0200 Subject: [PATCH 8/8] Further condense and reuse --- .../SwitchCaseReturnsToSwitchExpression.java | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java index 133a03eadd..cbff93f8f3 100644 --- a/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java +++ b/src/main/java/org/openrewrite/java/migrate/lang/SwitchCaseReturnsToSwitchExpression.java @@ -85,16 +85,16 @@ private boolean canConvertToSwitchExpression(J.Switch switchStatement) { 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)) { + if (body instanceof J.Block) { + if (!isReturnCase(((J.Block) body).getStatements())) { + return false; + } + } else if (!(body instanceof J.Return)) { return false; } } else { // Colon case - List statements = caseStatement.getStatements(); - if (statements.isEmpty() || !isReturnCase(statements)) { + if (!isReturnCase(caseStatement.getStatements())) { return false; } } @@ -105,17 +105,15 @@ private boolean canConvertToSwitchExpression(J.Switch switchStatement) { } private boolean isReturnCase(List statements) { - if (statements.isEmpty()) { + if (statements.size() != 1) { return false; } - // Handle block containing a single return - if (statements.size() == 1 && statements.get(0) instanceof J.Block) { + if (statements.get(0) instanceof J.Block) { return isReturnCase(((J.Block) statements.get(0)).getStatements()); } - // Direct return statement - return statements.size() == 1 && statements.get(0) instanceof J.Return; + return statements.get(0) instanceof J.Return; } private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { @@ -193,7 +191,6 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { if (statements.size() != 1) { return null; } - // Handle block containing a single return if (statements.get(0) instanceof J.Block) { J.Block block = (J.Block) statements.get(0); @@ -201,12 +198,10 @@ private J.SwitchExpression convertToSwitchExpression(J.Switch switchStatement) { return ((J.Return) block.getStatements().get(0)).getExpression(); } } - // Direct return statement if (statements.get(0) instanceof J.Return) { return ((J.Return) statements.get(0)).getExpression(); } - return null; } });