Port CEL-Go basic Protobuf constant folding syntax to CEL-Java

PiperOrigin-RevId: 944425297

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;
@@ -305,25 +310,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.getTypeProvider(), 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 +343,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 +374,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) {
+      CelTypeProvider typeProvider, Optional<?> optResult, CelMutableAst mutableAst, CelMutableExpr expr) {
     Object unwrappedResult = optResult.orElse(null);
     if (unwrappedResult == null) {
       if (isCallToFunction(expr, Function.OPTIONAL_NONE.getFunction())) {
@@ -387,7 +423,8 @@ private Optional<CelMutableAst> maybeRewriteOptional(
       return Optional.empty();
     }
 
-    CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(unwrappedResult).orElse(null);
+    CelMutableExpr adaptedResult =
+        maybeAdaptEvaluatedResult(typeProvider, unwrappedResult).orElse(null);
     if (adaptedResult == null) {
       // Evaluated result is not an adaptable constant. Leave the optional as is.
       return Optional.empty();

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,8 @@ 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: '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 +473,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 -> 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();
@@ -561,10 +617,7 @@ public void iterationLimitReached_throws() throws Exception {
     Cel cel =
         runtimeFlavor
             .builder()
-            .setOptions(
-                CelOptions.current()
-                    .enableHeterogeneousNumericComparisons(true)
-                    .build())
+            .setOptions(CelOptions.current().enableHeterogeneousNumericComparisons(true).build())
             .build();
     CelAbstractSyntaxTree ast = cel.compile("1 + 1").getAst();
     CelOptimizer optimizer =