20
20
import org .openrewrite .TreeVisitor ;
21
21
import org .openrewrite .java .JavaIsoVisitor ;
22
22
import org .openrewrite .java .tree .J ;
23
+ import org .openrewrite .java .tree .JavaType ;
23
24
import org .openrewrite .java .tree .Statement ;
25
+ import org .openrewrite .java .tree .TypeUtils ;
24
26
25
27
import java .util .ArrayList ;
26
28
import java .util .List ;
@@ -58,7 +60,7 @@ private J.Block maybeRemoveReturnAsLastStatement(J.Block b) {
58
60
}
59
61
Statement lastStatement = statements .get (statements .size () - 1 );
60
62
List <Statement > allButLast = statements .subList (0 , statements .size () - 1 );
61
- if (lastStatement instanceof J .Return ) {
63
+ if (lastStatement instanceof J .Return && (( J . Return ) lastStatement ). getExpression () == null ) {
62
64
return b .withStatements (allButLast );
63
65
} else if (lastStatement instanceof J .If ) {
64
66
J .If ifStatement = (J .If ) lastStatement ;
@@ -77,7 +79,11 @@ private J.Block maybeRemoveReturnAsLastStatement(J.Block b) {
77
79
@ Override
78
80
public J .MethodDeclaration visitMethodDeclaration (J .MethodDeclaration method , ExecutionContext ctx ) {
79
81
J .MethodDeclaration m = super .visitMethodDeclaration (method , ctx );
80
- return m .withBody (maybeRemoveReturnAsLastStatement (m .getBody ()));
82
+ if (TypeUtils .asPrimitive (m .getType ()) == JavaType .Primitive .Void ) {
83
+ return m .withBody (maybeRemoveReturnAsLastStatement (m .getBody ()));
84
+ } else {
85
+ return m ;
86
+ }
81
87
}
82
88
};
83
89
}
0 commit comments