Port CEL-Go comprehension and basic Protobuf constant folding to CEL-Java

PiperOrigin-RevId: 944425297

Author: CEL Dev Team

optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel

@@ -40,6 +40,7 @@ java_library(
         "//runtime:unknown_attributes",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
+        "@maven//:com_google_protobuf_protobuf_java",
     ],
 )
 

optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java

@@ -23,6 +23,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import com.google.protobuf.Descriptors.FieldDescriptor;
+import com.google.protobuf.Message;
 import dev.cel.bundle.Cel;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelMutableAst;
@@ -276,6 +278,38 @@ private static boolean areChildrenArgConstant(CelNavigableMutableExpr expr) {
     return false;
   }
 
+  /**
+   * Checks whether the provided identifier node is a loop variable bound by a surrounding
+   * comprehension. Identifiers existing strictly within loop initialization bounds (e.g. iterRange)
+   * are not considered loop variables.
+   */
+  private static boolean isComprehensionBoundLoopVariable(CelNavigableMutableExpr identNode) {
+    String identName = identNode.expr().ident().name();
+    CelNavigableMutableExpr child = identNode;
+    CelNavigableMutableExpr parent = identNode.parent().orElse(null);
+
+    while (parent != null) {
+      if (parent.getKind().equals(Kind.COMPREHENSION)) {
+        CelMutableComprehension comprehension = parent.expr().comprehension();
+        long childId = child.expr().id();
+
+        boolean isInitializationNode =
+            childId == comprehension.iterRange().id() || childId == comprehension.accuInit().id();
+
+        if (!isInitializationNode) {
+          if (identName.equals(comprehension.accuVar())
+              || identName.equals(comprehension.iterVar())
+              || identName.equals(comprehension.iterVar2())) {
+            return true;
+          }
+        }
+      }
+      child = parent;
+      parent = parent.parent().orElse(null);
+    }
+    return false;
+  }
+
   private static boolean isNestedComprehension(CelNavigableMutableExpr expr) {
     Optional<CelNavigableMutableExpr> maybeParent = expr.parent();
     while (maybeParent.isPresent()) {
@@ -364,6 +398,21 @@ private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(Object result) {
               CelMutableExpr.ofConstant(CelConstant.ofValue(timestampStrArg)));
 
       return Optional.of(CelMutableExpr.ofCall(timestampCall));
+    } else if (result instanceof Message) {
+      Message message = (Message) result;
+      List<CelMutableStruct.Entry> structEntries = new ArrayList<>();
+      for (Map.Entry<FieldDescriptor, Object> entry : message.getAllFields().entrySet()) {
+        CelMutableExpr adaptedFieldExpr = maybeAdaptEvaluatedResult(entry.getValue()).orElse(null);
+        if (adaptedFieldExpr == null) {
+          return Optional.empty();
+        }
+        structEntries.add(
+            CelMutableStruct.Entry.create(0, entry.getKey().getName(), adaptedFieldExpr));
+      }
+      return Optional.of(
+          CelMutableExpr.ofStruct(
+              CelMutableStruct.create(
+                  message.getDescriptorForType().getFullName(), structEntries)));
     }
 
     // Evaluated result cannot be folded (e.g: unknowns)
@@ -733,6 +782,7 @@ private static Object evaluateExpr(Cel cel, CelNavigableMutableExpr navigableMut
         navigableMutableExpr
             .allNodes()
             .filter(node -> node.getKind().equals(Kind.IDENT))
+            .filter(node -> !isComprehensionBoundLoopVariable(node))
             .map(node -> node.expr().ident().name())
             .map(CelAttributePattern::fromQualifiedIdentifier)
             .collect(toImmutableList());

optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java

@@ -228,6 +228,17 @@ private static Cel setupEnv(CelBuilder celBuilder) {
   @TestParameters("{source: 'true == false', expected: 'false'}")
   @TestParameters("{source: 'true == true', expected: 'true'}")
   @TestParameters("{source: 'false == true', expected: 'false'}")
+  @TestParameters("{source: '[1, 2, 3].map(item, item * 2)', expected: '[2, 4, 6]'}")
+  @TestParameters("{source: '[1, 2, 3].filter(item, item > 1)', expected: '[2, 3]'}")
+  @TestParameters("{source: '[1, 2, 3].exists(item, item > 1)', expected: 'true'}")
+  @TestParameters("{source: '[1, 2, 3].all(item, item > 1)', expected: 'false'}")
+  @TestParameters(
+      "{source: '[1, 2, x].map(item, item * 2)', expected: '[1, 2, x].map(item, item * 2)'}")
+  @TestParameters(
+      "{source: '[(x - 1 > 3) ? (x - 1) : 5].exists(x, x - 1 > 3)', expected: '[(x - 1 > 3) ? (x -"
+          + " 1) : 5].exists(x, x - 1 > 3)'}")
+  @TestParameters("{source: 'TestAllTypes{single_int64: 1 + 2 + 3}.single_int64', expected: '6'}")
+  @TestParameters("{source: 'TestAllTypes{single_int64: 3}.single_int64', expected: '3'}")
   @TestParameters("{source: 'false == false', expected: 'true'}")
   @TestParameters("{source: '10 == 42', expected: 'false'}")
   @TestParameters("{source: '42 == 42', expected: 'true'}")