Adding a check for Void return type and no return expression

openrewrite · Nov 17, 2024 · bcc951a · bcc951a
1 parent 7112d6e
commit bcc951a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/src/main/java/org/openrewrite/staticanalysis/UnnecessaryReturnAsLastStatement.java b/src/main/java/org/openrewrite/staticanalysis/UnnecessaryReturnAsLastStatement.java
@@ -20,7 +20,9 @@
 import org.openrewrite.TreeVisitor;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaType;
 import org.openrewrite.java.tree.Statement;
+import org.openrewrite.java.tree.TypeUtils;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -58,7 +60,7 @@ private J.Block maybeRemoveReturnAsLastStatement(J.Block b) {
                 }
                 Statement lastStatement = statements.get(statements.size() - 1);
                 List<Statement> allButLast = statements.subList(0, statements.size() - 1);
-                if (lastStatement instanceof J.Return) {
+                if (lastStatement instanceof J.Return && ((J.Return) lastStatement).getExpression() == null) {
                     return b.withStatements(allButLast);
                 } else if (lastStatement instanceof J.If) {
                     J.If ifStatement = (J.If) lastStatement;
@@ -77,7 +79,11 @@ private J.Block maybeRemoveReturnAsLastStatement(J.Block b) {
             @Override
             public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
                 J.MethodDeclaration m = super.visitMethodDeclaration(method, ctx);
-                return m.withBody(maybeRemoveReturnAsLastStatement(m.getBody()));
+                if (TypeUtils.asPrimitive(m.getType()) == JavaType.Primitive.Void) {
+                    return m.withBody(maybeRemoveReturnAsLastStatement(m.getBody()));
+                } else {
+                    return m;
+                }
             }
         };
     }

diff --git a/src/test/java/org/openrewrite/staticanalysis/UnnecessaryReturnAsLastStatementTest.java b/src/test/java/org/openrewrite/staticanalysis/UnnecessaryReturnAsLastStatementTest.java
@@ -105,4 +105,19 @@ void world(int i) {
               }
               """));
     }
+
+    @Test
+    void notChangingNonVoidMethods() {
+        //language=java
+        rewriteRun(
+          java(
+            """
+              class Hello {
+                  int world(int i) {
+                      return i + 436;
+                  }
+              }
+              """));
+    }
+
 }