Simplify ifs with literal conditions even if they jump

src/main/java/org/openrewrite/staticanalysis/RemoveUnreachableCodeVisitor.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2023 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (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://www.apache.org/licenses/LICENSE-2.0
+ * <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.staticanalysis;
+
+import org.openrewrite.ExecutionContext;
+import org.openrewrite.internal.ListUtils;
+import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.Statement;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+
+class RemoveUnreachableCodeVisitor extends JavaVisitor<ExecutionContext> {
+
+  @Override
+  public J visitBlock(J.Block block, ExecutionContext executionContext) {
+    block = (J.Block) super.visitBlock(block, executionContext);
+
+    List<Statement> statements = block.getStatements();
+    Optional<Integer> maybeFirstJumpIndex = findFirstJump(statements);
+    if (!maybeFirstJumpIndex.isPresent()) {
+      return block;
+    }
+    int firstJumpIndex = maybeFirstJumpIndex.get();
+
+    if (firstJumpIndex == statements.size() - 1) {
+      // Jump is at the end of the block, so nothing to do
+      return block;
+    }
+
+    List<Statement> newStatements =
+        ListUtils.flatMap(
+            block.getStatements(),
+            (index, statement) -> {
+              if (index <= firstJumpIndex) {
+                return statement;
+              }
+              return Collections.emptyList();
+            }
+        );
+
+    return block.withStatements(newStatements);
+  }
+
+  private Optional<Integer> findFirstJump(List<Statement> statements) {
+    for (int i = 0; i < statements.size(); i++) {
+      Statement statement = statements.get(i);
+      if (
+          statement instanceof J.Return ||
+          statement instanceof J.Throw ||
+          statement instanceof J.Break ||
+          statement instanceof J.Continue
+      ) {
+        return Optional.of(i);
+      }
+    }
+    return Optional.empty();
+  }
+}

src/main/java/org/openrewrite/staticanalysis/SimplifyConstantIfBranchExecution.java

@@ -20,7 +20,6 @@
 import org.openrewrite.SourceFile;
 import org.openrewrite.TreeVisitor;
 import org.openrewrite.internal.lang.Nullable;
-import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.cleanup.SimplifyBooleanExpressionVisitor;
 import org.openrewrite.java.cleanup.UnnecessaryParenthesesVisitor;
@@ -32,7 +31,6 @@
 import org.openrewrite.java.tree.Statement;
 
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 public class SimplifyConstantIfBranchExecution extends Recipe {
 
@@ -93,10 +91,6 @@ public J visitIf(J.If if_, ExecutionContext context) {
             J.ControlParentheses<Expression> cp = cleanupBooleanExpression(if__.getIfCondition(), context);
             if__ = if__.withIfCondition(cp);
 
-            if (visitsKeyWord(if__)) {
-                return if__;
-            }
-
             // The compile-time constant value of the if condition control parentheses.
             final Optional<Boolean> compileTimeConstantBoolean;
             if (isLiteralTrue(cp.getTree())) {
@@ -115,6 +109,7 @@ public J visitIf(J.If if_, ExecutionContext context) {
                 // True branch
                 // Only keep the `then` branch, and remove the `else` branch.
                 Statement s = if__.getThenPart().withPrefix(if__.getPrefix());
+                doAfterVisit(new RemoveUnreachableCodeVisitor());
                 return maybeAutoFormat(
                         if__,
                         s,
@@ -126,6 +121,7 @@ public J visitIf(J.If if_, ExecutionContext context) {
                 if (if__.getElsePart() != null) {
                     // The `else` part needs to be kept
                     Statement s = if__.getElsePart().getBody().withPrefix(if__.getPrefix());
+                    doAfterVisit(new RemoveUnreachableCodeVisitor());
                     return maybeAutoFormat(
                             if__,
                             s,
@@ -151,41 +147,6 @@ public J visitIf(J.If if_, ExecutionContext context) {
             }
         }
 
-        private boolean visitsKeyWord(J.If iff) {
-            if (isLiteralFalse(iff.getIfCondition().getTree())) {
-                return false;
-            }
-
-            AtomicBoolean visitedCFKeyword = new AtomicBoolean(false);
-            // if there is a return, break, continue, throws in _then, then set visitedKeyword to true
-            new JavaIsoVisitor<AtomicBoolean>() {
-                @Override
-                public J.Return visitReturn(J.Return _return, AtomicBoolean atomicBoolean) {
-                    atomicBoolean.set(true);
-                    return _return;
-                }
-
-                @Override
-                public J.Continue visitContinue(J.Continue continueStatement, AtomicBoolean atomicBoolean) {
-                    atomicBoolean.set(true);
-                    return continueStatement;
-                }
-
-                @Override
-                public J.Break visitBreak(J.Break breakStatement, AtomicBoolean atomicBoolean) {
-                    atomicBoolean.set(true);
-                    return breakStatement;
-                }
-
-                @Override
-                public J.Throw visitThrow(J.Throw thrown, AtomicBoolean atomicBoolean) {
-                    atomicBoolean.set(true);
-                    return thrown;
-                }
-            }.visit(iff.getThenPart(), visitedCFKeyword);
-            return visitedCFKeyword.get();
-        }
-
         private static boolean isLiteralTrue(@Nullable Expression expression) {
             return J.Literal.isLiteralValue(expression, Boolean.TRUE);
         }