Merge pull request #199 from l46kok:fix-const-fold-with-list

PiperOrigin-RevId: 592940954
google · Dec 21, 2023 · d76c341 · d76c341
2 parents 95714d0 + 37799bc
commit d76c341
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,7 +2,7 @@
 bazel-bin
 bazel-examples
 bazel-genfiles
-bazel-grpc-java
+bazel-cel-java
 bazel-out
 bazel-testlogs
 
@@ -27,4 +27,4 @@ bin
 target
 
 # Temporary output dir for artifacts
-mvn-artifacts
+mvn-artifacts
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
@@ -299,7 +299,12 @@ private Optional<CelAbstractSyntaxTree> maybePruneBranches(
 
       return Optional.of(replaceSubtree(ast, result, expr.id()));
     } else if (function.equals(Operator.IN.getFunction())) {
-      CelCreateList haystack = call.args().get(1).createList();
+      CelExpr callArg = call.args().get(1);
+      if (!callArg.exprKind().getKind().equals(Kind.CREATE_LIST)) {
+        return Optional.empty();
+      }
+
+      CelCreateList haystack = callArg.createList();
       if (haystack.elements().isEmpty()) {
         return Optional.of(
             replaceSubtree(

diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
@@ -23,6 +23,8 @@
 import dev.cel.bundle.CelFactory;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelOptions;
+import dev.cel.common.types.ListType;
+import dev.cel.common.types.MapType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.extensions.CelOptionalLibrary;
 import dev.cel.optimizer.CelOptimizationException;
@@ -41,6 +43,8 @@ public class ConstantFoldingOptimizerTest {
       CelFactory.standardCelBuilder()
           .addVar("x", SimpleType.DYN)
           .addVar("y", SimpleType.DYN)
+          .addVar("list_var", ListType.create(SimpleType.STRING))
+          .addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.STRING))
           .addMessageTypes(TestAllTypes.getDescriptor())
           .setContainer("dev.cel.testing.testdata.proto3")
           .addCompilerLibraries(CelOptionalLibrary.INSTANCE)
@@ -152,6 +156,8 @@ public class ConstantFoldingOptimizerTest {
   @TestParameters("{source: 'x + dyn([1, 2] + [3, 4])', expected: 'x + [1, 2, 3, 4]'}")
   @TestParameters(
       "{source: '{\"a\": dyn([1, 2]), \"b\": x}', expected: '{\"a\": [1, 2], \"b\": x}'}")
+  @TestParameters("{source: 'map_var[?\"key\"]', expected: 'map_var[?\"key\"]'}")
+  @TestParameters("{source: '\"abc\" in list_var', expected: '\"abc\" in list_var'}")
   // TODO: Support folding lists with mixed types. This requires mutable lists.
   // @TestParameters("{source: 'dyn([1]) + [1.0]'}")
   public void constantFold_success(String source, String expected) throws Exception {