Merge pull request #47 from openrewrite/feat/JLL/constant-folding

Add new `ConstantFold` utility

src/main/java/org/openrewrite/analysis/constantfold/ConstantFold.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2024 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.analysis.constantfold;
+
+import fj.data.Option;
+import lombok.AllArgsConstructor;
+import org.openrewrite.Cursor;
+import org.openrewrite.Incubating;
+import org.openrewrite.Tree;
+import org.openrewrite.analysis.dataflow.DataFlowNode;
+import org.openrewrite.analysis.trait.Top;
+import org.openrewrite.analysis.trait.expr.Expr;
+import org.openrewrite.analysis.trait.expr.VarAccess;
+import org.openrewrite.analysis.trait.variable.Variable;
+import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.tree.J;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+@Incubating(since="2.4.0")
+public class ConstantFold {
+
+    /**
+     * Find the constant value that is being assigned to a variable, if any.
+     * Otherwise, return the {@link J} itself.
+     */
+    public static Option<J> findConstantValue(Cursor cursor) {
+        return findConstantExpr(cursor)
+                .bind(expr -> {
+                    TopFinderVisitor topFinder = new TopFinderVisitor(expr);
+                    AtomicReference<J> found = new AtomicReference<>();
+                    topFinder.visit(
+                            cursor.dropParentUntil(J.CompilationUnit.class::isInstance).<Tree>getValue(),
+                            found
+                    );
+                    return Option.some(found.get());
+                });
+
+
+    }
+
+    /**
+     * Find the constant expression that is being assigned to a variable, if any.
+     * Otherwise, return the expression itself.
+     */
+    public static Option<Expr> findConstantExpr(Cursor cursor) {
+        return DataFlowNode
+                .of(cursor)
+                .bind(n -> n
+                        // If the DataFlow node is a VarAccess
+                        .asExpr(VarAccess.class)
+                        // Get the variable that is being accessed
+                        .map(VarAccess::getVariable)
+                        // Get the assigned values to that variable
+                        .map(Variable::getAssignedValues)
+                        // If there is more than one assigned values,
+                        // we can't determine which one is the one we are looking for
+                        .filter(values -> values.size() == 1)
+                        // Get the single value
+                        .map(values -> values.iterator().next())
+                        // If the value is an expression other than a var access, just return it
+                        .orElse(() -> n.asExpr(Expr.class)));
+    }
+
+    @AllArgsConstructor
+    private static final class TopFinderVisitor extends JavaVisitor<AtomicReference<J>> {
+        private final Top top;
+
+        @Override
+        public J preVisit(J tree, AtomicReference<J> p) {
+            if (top.getId().equals(tree.getId())) {
+                stopAfterPreVisit();
+                if (p.get() != null) {
+                    throw new IllegalStateException("Multiple top-level trees found for " + top);
+                }
+                p.set(tree);
+            }
+            return super.preVisit(tree, p);
+        }
+    }
+}

src/main/java/org/openrewrite/analysis/constantfold/package-info.java

@@ -0,0 +1,19 @@
+/*
+ * 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.
+ */
+@NonNullApi
+package org.openrewrite.analysis.constantfold;
+
+import org.openrewrite.internal.lang.NonNullApi;

src/main/java/org/openrewrite/analysis/trait/expr/VarAccess.java

@@ -39,7 +39,7 @@
 import java.util.function.Predicate;
 
 /**
- * A variable access is a (possibly qualified) reference to a field,
+ * Variable access is a (possibly qualified) reference to a field,
  * parameter or local variable.
  */
 public interface VarAccess extends Expr {

src/test/java/org/openrewrite/analysis/constantfold/ConstantFoldTest.java

@@ -0,0 +1,157 @@
+/*
+ * Copyright 2024 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.analysis.constantfold;
+
+import fj.data.Option;
+import org.junit.jupiter.api.Test;
+import org.openrewrite.DocumentExample;
+import org.openrewrite.analysis.InvocationMatcher;
+import org.openrewrite.java.JavaIsoVisitor;
+import org.openrewrite.java.tree.Expression;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.marker.SearchResult;
+import org.openrewrite.test.RecipeSpec;
+import org.openrewrite.test.RewriteTest;
+
+import static org.openrewrite.java.Assertions.java;
+
+public class ConstantFoldTest implements RewriteTest {
+
+    @Override
+    public void defaults(RecipeSpec spec) {
+        spec.recipe(RewriteTest.toRecipe(ConstantFoldVisitor::new));
+    }
+
+    static class ConstantFoldVisitor<P> extends JavaIsoVisitor<P> {
+        private static InvocationMatcher ALL_MATCHER = t -> true;
+
+        @Override
+        public Expression visitExpression(Expression expression, P p) {
+            if (ALL_MATCHER.advanced().isFirstArgument(getCursor())) {
+                Option<J> constantValue = ConstantFold.findConstantValue(getCursor());
+                if (constantValue.isSome() && constantValue.map(j -> j.getMarkers().findFirst(SearchResult.class).isEmpty()).orSome(false)) {
+                    return SearchResult.found(expression, constantValue.some().print(getCursor()).trim());
+                }
+            }
+            return super.visitExpression(expression, p);
+        }
+    }
+
+    @Test
+    public void constantFoldSingleElementString() {
+        rewriteRun(
+          java(
+            """
+            class Test {
+                private static final String FOO_IDENT = "FOO";
+                void test() {
+                    any(FOO_IDENT);
+                }
+                
+                void any(Object arg) {
+                    // No-op
+                }
+            }
+            """,
+            """
+            class Test {
+                private static final String FOO_IDENT = "FOO";
+                void test() {
+                    any(/*~~("FOO")~~>*/FOO_IDENT);
+                }
+                
+                void any(Object arg) {
+                    // No-op
+                }
+            }
+            """
+          )
+        );
+    }
+
+    @Test
+    @DocumentExample
+    public void constantFoldTagsElements() {
+        rewriteRun(
+          java(
+            """
+            class Test {
+                private static final String FOO_IDENT = "FOO";
+                private static final String BAR_IDENT = "BAR";
+                void test() {
+                    any(FOO_IDENT);
+                    any(BAR_IDENT);
+                    any("BAZ");
+                }
+                
+                void any(Object arg) {
+                    // No-op
+                }
+            }
+            """,
+            """
+            class Test {
+                private static final String FOO_IDENT = "FOO";
+                private static final String BAR_IDENT = "BAR";
+                void test() {
+                    any(/*~~("FOO")~~>*/FOO_IDENT);
+                    any(/*~~("BAR")~~>*/BAR_IDENT);
+                    any(/*~~("BAZ")~~>*/"BAZ");
+                }
+                
+                void any(Object arg) {
+                    // No-op
+                }
+            }
+            """
+          )
+        );
+    }
+
+    @Test
+    public void constantFoldTagsArray() {
+        rewriteRun(
+          java(
+            """
+            class Test {
+                private static final String[] FOO_IDENT = new String[] {"FOO"};
+                void test() {
+                    any(FOO_IDENT);
+                    any(new String[] {"BAR"});
+                }
+                
+                void any(Object arg) {
+                    // No-op
+                }
+            }
+            """,
+            """
+            class Test {
+                private static final String[] FOO_IDENT = new String[] {"FOO"};
+                void test() {
+                    any(/*~~(new String[] {"FOO"})~~>*/FOO_IDENT);
+                    any(/*~~(new String[] {"BAR"})~~>*/new String[] {"BAR"});
+                }
+                
+                void any(Object arg) {
+                    // No-op
+                }
+            }
+            """
+          )
+        );
+    }
+}