Improve performance for J.Ternary expressions (#63)

* Improve performance for `J.Ternary` expressions * Only add `REMOVE_PARENS` when unary expression has parens
openrewrite · Jan 9, 2024 · 9a7f7b1 · 9a7f7b1
1 parent ff1a566
commit 9a7f7b1
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 15 deletions.
diff --git a/src/main/java/org/openrewrite/java/template/processor/RefasterTemplateProcessor.java b/src/main/java/org/openrewrite/java/template/processor/RefasterTemplateProcessor.java
@@ -81,6 +81,9 @@ protected List<String> computeValue(Class<?> type) {
                 return singletonList("J.MethodInvocation");
             } else if (JCTree.JCFieldAccess.class.isAssignableFrom(type)) {
                 return Arrays.asList("J.FieldAccess", "J.Identifier");
+            } else if (JCTree.JCConditional.class.isAssignableFrom(type)) {
+                // catch all for expressions
+                return singletonList("J.Ternary");
             } else if (JCTree.JCExpression.class.isAssignableFrom(type)) {
                 // catch all for expressions
                 return singletonList("Expression");
@@ -257,7 +260,9 @@ public void visitClassDef(JCTree.JCClassDecl classDecl) {
                             maybeRemoveImports(staticImports, recipe, entry.getValue(), descriptor.afterTemplate);
 
                             List<String> embedOptions = new ArrayList<>();
-                            if (getType(descriptor.afterTemplate) == JCTree.JCParens.class || getType(descriptor.afterTemplate) == JCTree.JCUnary.class) {
+                            JCTree.JCExpression afterReturn = getReturnExpression(descriptor.afterTemplate);
+                            if (afterReturn instanceof JCTree.JCParens ||
+                                afterReturn instanceof JCTree.JCUnary && ((JCTree.JCUnary) afterReturn).getExpression() instanceof JCTree.JCParens) {
                                 embedOptions.add("REMOVE_PARENS");
                             }
                             // TODO check if after template contains type or member references
@@ -621,14 +626,19 @@ public void scan(JCTree jcTree) {
     }
 
     private Class<? extends JCTree> getType(JCTree.JCMethodDecl method) {
-        JCTree.JCStatement statement = method.getBody().getStatements().get(0);
-        Class<? extends JCTree> type = statement.getClass();
+        JCTree.JCExpression returnExpression = getReturnExpression(method);
+        return returnExpression != null ? returnExpression.getClass() : method.getBody().getStatements().last().getClass();
+    }
+
+    @Nullable
+    private JCTree.JCExpression getReturnExpression(JCTree.JCMethodDecl method) {
+        JCTree.JCStatement statement = method.getBody().getStatements().last();
         if (statement instanceof JCTree.JCReturn) {
-            type = ((JCTree.JCReturn) statement).expr.getClass();
+            return ((JCTree.JCReturn) statement).expr;
         } else if (statement instanceof JCTree.JCExpressionStatement) {
-            type = ((JCTree.JCExpressionStatement) statement).expr.getClass();
+            return ((JCTree.JCExpressionStatement) statement).expr;
         }
-        return type;
+        return null;
     }
 
     private String statementType(JCTree.JCMethodDecl method) {

diff --git a/src/test/resources/refaster/ShouldSupportNestedClassesRecipes.java b/src/test/resources/refaster/ShouldSupportNestedClassesRecipes.java
@@ -33,9 +33,15 @@
 
 import static org.openrewrite.java.template.internal.AbstractRefasterJavaVisitor.EmbeddingOption.*;
 
+
+/**
+ * OpenRewrite recipes created for Refaster template {@code foo.ShouldSupportNestedClasses}.
+ */
 @SuppressWarnings("all")
 public class ShouldSupportNestedClassesRecipes extends Recipe {
-
+    /**
+     * Instantiates a new instance.
+     */
     public ShouldSupportNestedClassesRecipes() {}
 
     @Override
@@ -56,10 +62,16 @@ public List<Recipe> getRecipeList() {
         );
     }
 
+    /**
+     * OpenRewrite recipe created for Refaster template {@code ShouldSupportNestedClasses.NestedClass}.
+     */
     @SuppressWarnings("all")
     @NonNullApi
     public static class NestedClassRecipe extends Recipe {
 
+        /**
+         * Instantiates a new instance.
+         */
         public NestedClassRecipe() {}
 
         @Override
@@ -86,7 +98,7 @@ public J visitBinary(J.Binary elem, ExecutionContext ctx) {
                                 after.apply(getCursor(), elem.getCoordinates().replace(), matcher.parameter(0)),
                                 getCursor(),
                                 ctx,
-                                REMOVE_PARENS, SHORTEN_NAMES, SIMPLIFY_BOOLEANS
+                                SHORTEN_NAMES, SIMPLIFY_BOOLEANS
                         );
                     }
                     return super.visitBinary(elem, ctx);
@@ -100,10 +112,16 @@ public J visitBinary(J.Binary elem, ExecutionContext ctx) {
         }
     }
 
+    /**
+     * OpenRewrite recipe created for Refaster template {@code ShouldSupportNestedClasses.AnotherClass}.
+     */
     @SuppressWarnings("all")
     @NonNullApi
     public static class AnotherClassRecipe extends Recipe {
 
+        /**
+         * Instantiates a new instance.
+         */
         public AnotherClassRecipe() {}
 
         @Override

diff --git a/src/test/resources/refaster/SimplifyTernaryRecipes.java b/src/test/resources/refaster/SimplifyTernaryRecipes.java
@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package foo;
 
 import org.openrewrite.ExecutionContext;
@@ -91,7 +90,7 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                 final JavaTemplate after = Semantics.expression(this, "after", (@Primitive Boolean expr) -> expr).build();
 
                 @Override
-                public J visitExpression(Expression elem, ExecutionContext ctx) {
+                public J visitTernary(J.Ternary elem, ExecutionContext ctx) {
                     JavaTemplate.Matcher matcher;
                     if ((matcher = before.matcher(getCursor())).find()) {
                         return embed(
@@ -101,7 +100,7 @@ public J visitExpression(Expression elem, ExecutionContext ctx) {
                                 SHORTEN_NAMES, SIMPLIFY_BOOLEANS
                         );
                     }
-                    return super.visitExpression(elem, ctx);
+                    return super.visitTernary(elem, ctx);
                 }
 
             };
@@ -138,7 +137,7 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                 final JavaTemplate after = Semantics.expression(this, "after", (@Primitive Boolean expr) -> !(expr)).build();
 
                 @Override
-                public J visitExpression(Expression elem, ExecutionContext ctx) {
+                public J visitTernary(J.Ternary elem, ExecutionContext ctx) {
                     JavaTemplate.Matcher matcher;
                     if ((matcher = before.matcher(getCursor())).find()) {
                         return embed(
@@ -148,7 +147,7 @@ public J visitExpression(Expression elem, ExecutionContext ctx) {
                                 REMOVE_PARENS, SHORTEN_NAMES, SIMPLIFY_BOOLEANS
                         );
                     }
-                    return super.visitExpression(elem, ctx);
+                    return super.visitTernary(elem, ctx);
                 }
 
             };

diff --git a/src/test/resources/refaster/UseStringIsEmptyRecipe.java b/src/test/resources/refaster/UseStringIsEmptyRecipe.java
@@ -35,12 +35,15 @@
 
 
 /**
- * OpenRewrite recipe created for Refaster template `UseStringIsEmpty`.
+ * OpenRewrite recipe created for Refaster template {@code UseStringIsEmpty}.
  */
 @SuppressWarnings("all")
 @NonNullApi
 public class UseStringIsEmptyRecipe extends Recipe {
 
+    /**
+     * Instantiates a new instance.
+     */
     public UseStringIsEmptyRecipe() {}
 
     @Override
@@ -67,7 +70,7 @@ public J visitBinary(J.Binary elem, ExecutionContext ctx) {
                             after.apply(getCursor(), elem.getCoordinates().replace(), matcher.parameter(0)),
                             getCursor(),
                             ctx,
-                            REMOVE_PARENS, SHORTEN_NAMES, SIMPLIFY_BOOLEANS
+                            SHORTEN_NAMES, SIMPLIFY_BOOLEANS
                     );
                 }
                 return super.visitBinary(elem, ctx);