Port CEL-Go comprehension loop variable shadowing to CEL-Java

PiperOrigin-RevId: 944969366

Author: CEL Dev Team

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

@@ -31,6 +31,7 @@ java_library(
         "//common/navigation:common",
         "//common/navigation:mutable_navigation",
         "//common/types",
+        "//common/types:type_providers",
         "//extensions:optional_library",
         "//optimizer:ast_optimizer",
         "//optimizer:mutable_ast",
@@ -40,6 +41,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;
@@ -42,7 +44,10 @@
 import dev.cel.common.navigation.CelNavigableMutableAst;
 import dev.cel.common.navigation.CelNavigableMutableExpr;
 import dev.cel.common.navigation.TraversalOrder;
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.CelTypeProvider;
 import dev.cel.common.types.SimpleType;
+import dev.cel.common.types.StructType;
 import dev.cel.extensions.CelOptionalLibrary.Function;
 import dev.cel.optimizer.AstMutator;
 import dev.cel.optimizer.CelAstOptimizer;
@@ -276,6 +281,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()) {
@@ -305,25 +342,28 @@ private Optional<CelMutableAst> maybeFold(
     // ex2: optional.ofNonZeroValue(5) -> optional.of(5)
     if (result instanceof Optional<?>) {
       Optional<?> optResult = ((Optional<?>) result);
-      return maybeRewriteOptional(optResult, mutableAst, node.expr());
+      return maybeRewriteOptional(cel, optResult, mutableAst, node.expr());
     }
 
-    CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(result).orElse(null);
+    CelMutableExpr adaptedResult =
+        maybeAdaptEvaluatedResult(cel.getTypeProvider(), result).orElse(null);
     if (adaptedResult == null) {
       return Optional.empty();
     }
 
     return Optional.of(astMutator.replaceSubtree(mutableAst, adaptedResult, node.id()));
   }
 
-  private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(Object result) {
+  private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(
+      CelTypeProvider typeProvider, Object result) {
     if (CelConstant.isConstantValue(result)) {
       return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofObjectValue(result)));
     } else if (result instanceof Collection<?>) {
       Collection<?> collection = (Collection<?>) result;
       List<CelMutableExpr> listElements = new ArrayList<>();
       for (Object evaluatedElement : collection) {
-        CelMutableExpr adaptedExpr = maybeAdaptEvaluatedResult(evaluatedElement).orElse(null);
+        CelMutableExpr adaptedExpr =
+            maybeAdaptEvaluatedResult(typeProvider, evaluatedElement).orElse(null);
         if (adaptedExpr == null) {
           return Optional.empty();
         }
@@ -335,11 +375,13 @@ private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(Object result) {
       Map<?, ?> map = (Map<?, ?>) result;
       List<CelMutableMap.Entry> mapEntries = new ArrayList<>();
       for (Map.Entry<?, ?> entry : map.entrySet()) {
-        CelMutableExpr adaptedKey = maybeAdaptEvaluatedResult(entry.getKey()).orElse(null);
+        CelMutableExpr adaptedKey =
+            maybeAdaptEvaluatedResult(typeProvider, entry.getKey()).orElse(null);
         if (adaptedKey == null) {
           return Optional.empty();
         }
-        CelMutableExpr adaptedValue = maybeAdaptEvaluatedResult(entry.getValue()).orElse(null);
+        CelMutableExpr adaptedValue =
+            maybeAdaptEvaluatedResult(typeProvider, entry.getValue()).orElse(null);
         if (adaptedValue == null) {
           return Optional.empty();
         }
@@ -364,14 +406,40 @@ 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;
+      Optional<CelType> typeOpt =
+          typeProvider.findType(message.getDescriptorForType().getFullName());
+      if (!typeOpt.isPresent() || !(typeOpt.get() instanceof StructType)) {
+        return Optional.empty();
+      }
+      StructType structType = (StructType) typeOpt.get();
+      List<CelMutableStruct.Entry> structEntries = new ArrayList<>();
+      for (String fieldName : structType.fieldNames()) {
+        FieldDescriptor fieldDescriptor = message.getDescriptorForType().findFieldByName(fieldName);
+        if (fieldDescriptor == null) {
+          continue;
+        }
+        if (!message.getAllFields().containsKey(fieldDescriptor)) {
+          continue;
+        }
+        CelMutableExpr adaptedFieldExpr =
+            maybeAdaptEvaluatedResult(typeProvider, message.getField(fieldDescriptor)).orElse(null);
+        if (adaptedFieldExpr == null) {
+          return Optional.empty();
+        }
+        structEntries.add(CelMutableStruct.Entry.create(0, fieldName, adaptedFieldExpr));
+      }
+      return Optional.of(
+          CelMutableExpr.ofStruct(CelMutableStruct.create(structType.name(), structEntries)));
     }
 
     // Evaluated result cannot be folded (e.g: unknowns)
     return Optional.empty();
   }
 
   private Optional<CelMutableAst> maybeRewriteOptional(
-      Optional<?> optResult, CelMutableAst mutableAst, CelMutableExpr expr) {
+      Cel cel, Optional<?> optResult, CelMutableAst mutableAst, CelMutableExpr expr) {
     Object unwrappedResult = optResult.orElse(null);
     if (unwrappedResult == null) {
       if (isCallToFunction(expr, Function.OPTIONAL_NONE.getFunction())) {
@@ -387,7 +455,8 @@ private Optional<CelMutableAst> maybeRewriteOptional(
       return Optional.empty();
     }
 
-    CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(unwrappedResult).orElse(null);
+    CelMutableExpr adaptedResult =
+        maybeAdaptEvaluatedResult(cel.getTypeProvider(), unwrappedResult).orElse(null);
     if (adaptedResult == null) {
       // Evaluated result is not an adaptable constant. Leave the optional as is.
       return Optional.empty();
@@ -733,6 +802,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/BUILD.bazel

@@ -44,6 +44,7 @@ java_library(
         "//:java_truth",
         "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",
         "@maven//:com_google_guava_guava",
+        "@maven//:com_google_protobuf_protobuf_java",
     ],
 )
 

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

@@ -195,6 +195,28 @@ private static Cel setupEnv(CelBuilder celBuilder) {
   @TestParameters(
       "{source: 'TestAllTypes{single_nested_message: TestAllTypes.NestedMessage{bb:"
           + " 42}}.single_nested_message.bb', expected: '42'}")
+  @TestParameters(
+      "{source: 'TestAllTypes{single_int32: 3}', "
+          + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_int32: 3}'}")
+  @TestParameters(
+      "{source: 'TestAllTypes{single_float: 1.5}', "
+          + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_float: 1.5}'}")
+  @TestParameters(
+      "{source: 'TestAllTypes{single_nested_message: TestAllTypes.NestedMessage{bb: 42}}', "
+          + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_nested_message:"
+          + " cel.expr.conformance.proto3.TestAllTypes.NestedMessage{bb: 42}}'}")
+  @TestParameters(
+      "{source: 'TestAllTypes{repeated_int32: [1, 2, 3]}', "
+          + " expected: 'cel.expr.conformance.proto3.TestAllTypes{repeated_int32: [1, 2, 3]}'}")
+  @TestParameters(
+      "{source: 'TestAllTypes{map_int32_int64: {1: 2}}', "
+          + " expected: 'cel.expr.conformance.proto3.TestAllTypes{map_int32_int64: {1: 2}}'}")
+  @TestParameters(
+      "{source: 'TestAllTypes{single_any: google.protobuf.Any{type_url:"
+          + " \"type.googleapis.com/google.protobuf.Int32Value\", value: b\"\\010\\001\"}}', "
+          + " expected: 'cel.expr.conformance.proto3.TestAllTypes{single_any:"
+          + " google.protobuf.Any{type_url:"
+          + " \"type.googleapis.com/google.protobuf.Int32Value\", value: b\"\\010\\001\"}}'}")
   @TestParameters("{source: '{\"a\": 1}[\"a\"]', expected: '1'}")
   @TestParameters("{source: '{\"a\": {\"b\": 2}}[\"a\"][\"b\"]', expected: '2'}")
   @TestParameters("{source: '{\"hello\": \"world\"}.hello == x', expected: '\"world\" == x'}")
@@ -228,6 +250,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'}")
@@ -449,6 +482,38 @@ public void constantFold_addFoldableFunction_success() throws Exception {
     assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("true");
   }
 
+  @Test
+  public void constantFold_functionReturningUnregisteredMessage_doesNotFold() throws Exception {
+    Cel customCel =
+        runtimeFlavor
+            .builder()
+            .addVar("x", SimpleType.DYN)
+            .addFunctionDeclarations(
+                CelFunctionDecl.newFunctionDeclaration(
+                    "get_unregistered_message",
+                    CelOverloadDecl.newGlobalOverload(
+                        "get_unregistered_message_overload", SimpleType.ANY)))
+            .addFunctionBindings(
+                CelFunctionBinding.from(
+                    "get_unregistered_message_overload",
+                    ImmutableList.of(),
+                    unused -> dev.cel.expr.conformance.proto3.TestAllTypes.getDefaultInstance()))
+            .build();
+    ConstantFoldingOptions options =
+        ConstantFoldingOptions.newBuilder()
+            .addFoldableFunctions("get_unregistered_message")
+            .build();
+    CelOptimizer optimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(customCel)
+            .addAstOptimizers(ConstantFoldingOptimizer.newInstance(options))
+            .build();
+    CelAbstractSyntaxTree ast = customCel.compile("get_unregistered_message()").getAst();
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("get_unregistered_message()");
+  }
+
   @Test
   public void constantFold_withExpectedResultTypeSet_success() throws Exception {
     Cel cel = runtimeFlavor.builder().setResultType(SimpleType.STRING).build();