Use new ConstantFold.findConstantLiteralValue

Signed-off-by: Jonathan Leitschuh <[email protected]>
openrewrite · Feb 24, 2024 · f336731 · f336731
1 parent c0ae6a6
commit f336731
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 27 deletions.
diff --git a/src/main/java/org/openrewrite/launchdarkly/ChangeVariationDefault.java b/src/main/java/org/openrewrite/launchdarkly/ChangeVariationDefault.java
@@ -18,6 +18,8 @@
 import lombok.EqualsAndHashCode;
 import lombok.Value;
 import org.openrewrite.*;
+import org.openrewrite.analysis.constantfold.ConstantFold;
+import org.openrewrite.analysis.util.CursorUtil;
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.internal.lang.NonNull;
 import org.openrewrite.java.JavaIsoVisitor;
@@ -67,17 +69,23 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
             public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
                 J.MethodInvocation mi = super.visitMethodInvocation(method, ctx);
                 Expression firstArgument = mi.getArguments().get(0);
+                boolean isFirstArgumentFeatureKey =
+                        CursorUtil
+                                .findCursorForTree(getCursor(), firstArgument)
+                                .bind(c -> ConstantFold.findConstantLiteralValue(c, String.class))
+                                .map(featureKey::equals)
+                                .orSome(false);
                 Expression lastArgument = mi.getArguments().get(mi.getArguments().size() - 1);
-                if (BOOL_VARIATION_MATCHER.matches(mi) && J.Literal.isLiteralValue(firstArgument, featureKey)) {
+                if (BOOL_VARIATION_MATCHER.matches(mi) && isFirstArgumentFeatureKey) {
                     return changeValue(mi, lastArgument, new J.Literal(Tree.randomId(), Space.SINGLE_SPACE, Markers.EMPTY, defaultValue, defaultValue, null, JavaType.Primitive.Boolean));
                 }
-                if (STRING_VARIATION_MATCHER.matches(mi) && J.Literal.isLiteralValue(firstArgument, featureKey)) {
+                if (STRING_VARIATION_MATCHER.matches(mi) && isFirstArgumentFeatureKey) {
                     return changeValue(mi, lastArgument, new J.Literal(Tree.randomId(), Space.SINGLE_SPACE, Markers.EMPTY, defaultValue, "\"" + defaultValue + "\"", null, JavaType.Primitive.String));
                 }
-                if (INT_VARIATION_MATCHER.matches(mi) && J.Literal.isLiteralValue(firstArgument, featureKey)) {
+                if (INT_VARIATION_MATCHER.matches(mi) && isFirstArgumentFeatureKey) {
                     return changeValue(mi, lastArgument, new J.Literal(Tree.randomId(), Space.SINGLE_SPACE, Markers.EMPTY, defaultValue, defaultValue, null, JavaType.Primitive.Int));
                 }
-                if (DOUBLE_VARIATION_MATCHER.matches(mi) && J.Literal.isLiteralValue(firstArgument, featureKey)) {
+                if (DOUBLE_VARIATION_MATCHER.matches(mi) && isFirstArgumentFeatureKey) {
                     return changeValue(mi, lastArgument, new J.Literal(Tree.randomId(), Space.SINGLE_SPACE, Markers.EMPTY, defaultValue, defaultValue, null, JavaType.Primitive.Double));
                 }
                 return mi;

diff --git a/src/main/java/org/openrewrite/launchdarkly/RemoveBoolVariation.java b/src/main/java/org/openrewrite/launchdarkly/RemoveBoolVariation.java
@@ -18,6 +18,8 @@
 import lombok.EqualsAndHashCode;
 import lombok.Value;
 import org.openrewrite.*;
+import org.openrewrite.analysis.constantfold.ConstantFold;
+import org.openrewrite.analysis.util.CursorUtil;
 import org.openrewrite.internal.StringUtils;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaVisitor;
@@ -75,7 +77,13 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
             @Override
             public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
                 J.MethodInvocation mi = (J.MethodInvocation) super.visitMethodInvocation(method, ctx);
-                if (methodMatcher.matches(mi) && J.Literal.isLiteralValue(mi.getArguments().get(0), featureKey)) {
+                boolean isFirstArgumentFeatureKey =
+                        CursorUtil
+                                .findCursorForTree(getCursor(), mi.getArguments().get(0))
+                                .bind(c -> ConstantFold.findConstantLiteralValue(c, String.class))
+                                .map(featureKey::equals)
+                                .orSome(false);
+                if (methodMatcher.matches(mi) && isFirstArgumentFeatureKey) {
                     doAfterVisit(new SimplifyConstantIfBranchExecution().getVisitor());
                     doAfterVisit(new RemoveUnusedLocalVariables(null).getVisitor());
                     doAfterVisit(new RemoveUnusedPrivateFields().getVisitor());

diff --git a/src/main/java/org/openrewrite/launchdarkly/search/FindFeatureFlag.java b/src/main/java/org/openrewrite/launchdarkly/search/FindFeatureFlag.java
@@ -22,13 +22,11 @@
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
 import org.openrewrite.analysis.InvocationMatcher;
+import org.openrewrite.analysis.constantfold.ConstantFold;
 import org.openrewrite.analysis.dataflow.DataFlowNode;
 import org.openrewrite.analysis.dataflow.DataFlowSpec;
 import org.openrewrite.analysis.dataflow.Dataflow;
-import org.openrewrite.analysis.trait.expr.Expr;
 import org.openrewrite.analysis.trait.expr.Literal;
-import org.openrewrite.analysis.trait.expr.VarAccess;
-import org.openrewrite.analysis.trait.variable.Variable;
 import org.openrewrite.internal.StringUtils;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaIsoVisitor;
@@ -111,25 +109,9 @@ public Expression visitExpression(Expression expression, ExecutionContext ctx) {
                         .findSinks(new DataFlowSpec() {
                             @Override
                             public boolean isSource(DataFlowNode srcNode) {
-                                return srcNode.asExpr(VarAccess.class)
-                                        .map(VarAccess::getVariable)
-                                        .map(Variable::getAssignedValues)
-                                        .bind(assignedVariables -> {
-                                            if (assignedVariables.size() > 1) {
-                                                return fj.data.Option.none();
-                                            }
-                                            for (Expr e : assignedVariables) {
-                                                if (e instanceof Literal) {
-                                                    Literal l = (Literal) e;
-                                                    return l.getValue()
-                                                            .map(featureKey::equals);
-                                                }
-                                            }
-                                            return fj.data.Option.none();
-                                        }).orElse(() -> srcNode
-                                                .asExpr(Literal.class)
-                                                .bind(Literal::getValue)
-                                                .map(featureKey::equals))
+                                return ConstantFold
+                                        .findConstantLiteralValue(srcNode, String.class)
+                                        .map(featureKey::equals)
                                         .orSome(false);
                             }