Null check as switch case

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

@@ -0,0 +1,148 @@
+/*
+ * 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.Value;
+import org.jspecify.annotations.Nullable;
+import org.openrewrite.Cursor;
+import org.openrewrite.java.JavaIsoVisitor;
+import org.openrewrite.java.search.SemanticallyEqual;
+import org.openrewrite.java.tree.Expression;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaType;
+import org.openrewrite.java.tree.Statement;
+import org.openrewrite.trait.SimpleTraitMatcher;
+import org.openrewrite.trait.Trait;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.openrewrite.java.tree.J.Binary.Type.Equal;
+
+@Value
+public class NullCheck implements Trait<J.If> {
+    Cursor cursor;
+    Expression nullCheckedParameter;
+
+    public Statement whenNull() {
+        return getTree().getThenPart();
+    }
+
+    public @Nullable Statement whenNotNull() {
+        J.If.Else else_ = getTree().getElsePart();
+        return else_ == null ? null : else_.getBody();
+    }
+
+    public boolean returns() {
+        Statement statement = whenNull();
+        if (statement instanceof J.Block) {
+            for (Statement s : ((J.Block) statement).getStatements()) {
+                if (s instanceof J.Return) {
+                    return true;
+                }
+            }
+            return false;
+        }
+        return statement instanceof J.Return;
+    }
+
+    /**
+     * Calculates few potential cases where the null checked variable gets reassigned and only returns false if these cases DO NOT match.
+     * In any other case this returns true as we do not know that particular situation yet -> no harm -> assume it could be altered in the block.
+     * @return false only if we are 100% sure the block does not reassigns/changes the null checked variable.
+     */
+    public boolean couldModifyNullCheckedValue() {
+        Statement statement = whenNull();
+        if (statement instanceof J.Block || statement instanceof Expression || statement instanceof J.Throw) {
+            return couldModifyNullCheckedValue(statement, nullCheckedParameter);
+        }
+        // Cautious by default
+        return true;
+    }
+    private static boolean couldModifyNullCheckedValue(J expression, Expression nullChecked) {
+        if (nullChecked instanceof J.FieldAccess && couldModifyNullCheckedValue(expression, ((J.FieldAccess) nullChecked).getTarget())) {
+            return true;
+        }
+        if (nullChecked instanceof J.MethodInvocation &&
+                ((J.MethodInvocation) nullChecked).getSelect() != null &&
+                couldModifyNullCheckedValue(expression, ((J.MethodInvocation) nullChecked).getSelect())) {
+            return true;
+        }
+        return new JavaIsoVisitor<AtomicBoolean>() {
+
+            private final boolean isCertainlyImmutable = nullChecked.getType() != null && JavaType.Primitive.fromClassName(nullChecked.getType().toString()) != null;
+
+            @Override
+            public J.Identifier visitIdentifier(J.Identifier identifier, AtomicBoolean couldModifyValue) {
+                J.Identifier id = super.visitIdentifier(identifier, couldModifyValue);
+                if (!isCertainlyImmutable && SemanticallyEqual.areEqual(id, nullChecked)) {
+                    couldModifyValue.set(true);
+                }
+                return id;
+            }
+            @Override
+            public J.Assignment visitAssignment(J.Assignment assignment, AtomicBoolean couldModifyValue) {
+                J.Assignment as = super.visitAssignment(assignment, couldModifyValue);
+                if (SemanticallyEqual.areEqual(as.getVariable(), nullChecked)) {
+                    couldModifyValue.set(true);
+                }
+                return as;
+            }
+            @Override
+            public J.FieldAccess visitFieldAccess(J.FieldAccess fieldAccess, AtomicBoolean couldModifyValue) {
+                J.FieldAccess fa = super.visitFieldAccess(fieldAccess, couldModifyValue);
+                if (SemanticallyEqual.areEqual(fa, nullChecked) ||
+                        SemanticallyEqual.areEqual(fa.getTarget(), nullChecked)) {
+                    couldModifyValue.set(true);
+                }
+                return fa;
+            }
+            @Override
+            public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, AtomicBoolean couldModifyValue) {
+                J.MethodInvocation mi = super.visitMethodInvocation(method, couldModifyValue);
+                if (SemanticallyEqual.areEqual(mi, nullChecked) ||
+                        ((mi.getSelect() != null) && SemanticallyEqual.areEqual(mi.getSelect(), nullChecked))) {
+                    couldModifyValue.set(true);
+                }
+                return mi;
+            }
+        }.reduce(expression, new AtomicBoolean(false)).get();
+    }
+
+    public static class Matcher extends SimpleTraitMatcher<NullCheck> {
+
+        public static Matcher nullCheck() {
+            return new Matcher();
+        }
+
+        @Override
+        protected @Nullable NullCheck test(Cursor cursor) {
+            if (cursor.getValue() instanceof J.If) {
+                J.If if_ = cursor.getValue();
+                if (if_.getIfCondition().getTree() instanceof J.Binary) {
+                    J.Binary binary = (J.Binary) if_.getIfCondition().getTree();
+                    if (binary.getOperator() == Equal) {
+                        if (J.Literal.isLiteralValue(binary.getLeft(), null)) {
+                            return new NullCheck(cursor, binary.getRight());
+                        } else if (J.Literal.isLiteralValue(binary.getRight(), null)) {
+                            return new NullCheck(cursor, binary.getLeft());
+                        }
+                    }
+                }
+            }
+            return null;
+        }
+    }
+}

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

@@ -0,0 +1,124 @@
+/*
+ * 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.jspecify.annotations.Nullable;
+import org.openrewrite.*;
+import org.openrewrite.internal.ListUtils;
+import org.openrewrite.java.JavaTemplate;
+import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.search.SemanticallyEqual;
+import org.openrewrite.java.tree.Expression;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.Space;
+import org.openrewrite.java.tree.Statement;
+import org.openrewrite.staticanalysis.kotlin.KotlinFileChecker;
+
+import java.time.Duration;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static java.util.Objects.requireNonNull;
+import static org.openrewrite.java.migrate.lang.NullCheck.Matcher.nullCheck;
+
+@Value
+@EqualsAndHashCode(callSuper = false)
+public class NullCheckAsSwitchCase extends Recipe {
+    @Override
+    public String getDisplayName() {
+        return "Add null check to existing switch cases";
+    }
+
+    @Override
+    public String getDescription() {
+        return "In later Java 21+, null checks are valid in switch cases. " +
+                "This recipe will only add null checks to existing switch cases if there are no other statements in between them " +
+                "or if the block in the if statement is not impacting the flow of the switch.";
+    }
+
+    @Override
+    public Duration getEstimatedEffortPerOccurrence() {
+        return Duration.ofMinutes(3);
+    }
+
+    @Override
+    public TreeVisitor<?, ExecutionContext> getVisitor() {
+        return Preconditions.check(Preconditions.not(new KotlinFileChecker<>()), new JavaVisitor<ExecutionContext>() {
+            @Override
+            public J visitBlock(J.Block block, ExecutionContext ctx) {
+                AtomicReference<@Nullable NullCheck> nullCheck = new AtomicReference<>();
+                J.Block b = block.withStatements(ListUtils.map(block.getStatements(), (index, statement) -> {
+                    // Maybe remove a null check preceding a switch statement
+                    Optional<NullCheck> nullCheckOpt = nullCheck().get(statement, getCursor());
+                    if (nullCheckOpt.isPresent()) {
+                        NullCheck check = nullCheckOpt.get();
+                        J nextStatement = index + 1 < block.getStatements().size() ? block.getStatements().get(index + 1) : null;
+                        if (!(nextStatement instanceof J.Switch) ||
+                                hasNullCase((J.Switch) nextStatement) ||
+                                !SemanticallyEqual.areEqual(((J.Switch) nextStatement).getSelector().getTree(), check.getNullCheckedParameter()) ||
+                                check.returns() ||
+                                check.couldModifyNullCheckedValue()) {
+                            return statement;
+                        }
+                        nullCheck.set(check);
+                        return null;
+                    }
+
+                    // Update the switch following a removed null check
+                    NullCheck check = nullCheck.getAndSet(null);
+                    if (check != null && statement instanceof J.Switch) {
+                        J.Switch aSwitch = (J.Switch) statement;
+                        J.Case nullCase = createNullCase(aSwitch, check.whenNull());
+                        return aSwitch.withCases(aSwitch.getCases().withStatements(
+                                ListUtils.insert(aSwitch.getCases().getStatements(), nullCase, 0)));
+                    }
+                    return statement;
+                }));
+                return super.visitBlock(b, ctx);
+            }
+
+            private boolean hasNullCase(J.Switch switch_) {
+                for (Statement c : switch_.getCases().getStatements()) {
+                    if (c instanceof J.Case) {
+                        for (J j : ((J.Case) c).getCaseLabels()) {
+                            if (j instanceof Expression && J.Literal.isLiteralValue((Expression) j, null)) {
+                                return true;
+                            }
+                        }
+                    }
+                }
+                return false;
+            }
+
+            private J.Case createNullCase(J.Switch aSwitch, Statement whenNull) {
+                if (whenNull instanceof J.Block && ((J.Block) whenNull).getStatements().size() == 1) {
+                    whenNull = ((J.Block) whenNull).getStatements().get(0);
+                }
+                String semicolon = whenNull instanceof J.Block ? "" : ";";
+                J.Switch switchWithNullCase = JavaTemplate.apply(
+                        "switch(#{any()}) { case null -> #{any()}" + semicolon + " }",
+                        new Cursor(getCursor(), aSwitch),
+                        aSwitch.getCoordinates().replace(),
+                        aSwitch.getSelector().getTree(),
+                        whenNull);
+                J.Case nullCase = (J.Case) switchWithNullCase.getCases().getStatements().get(0);
+                return nullCase.withBody(requireNonNull(nullCase.getBody()).withPrefix(Space.SINGLE_SPACE));
+            }
+        });
+    }
+}

src/main/resources/META-INF/rewrite/java-version-21.yml

@@ -141,5 +141,6 @@ description: ->-
 tags:
   - java21
 recipeList:
+  - org.openrewrite.java.migrate.lang.NullCheckAsSwitchCase
   - org.openrewrite.java.migrate.lang.RefineSwitchCases
   - org.openrewrite.java.migrate.lang.SwitchCaseEnumGuardToLabel