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

CEL Dev Team · copybara-github · commit eca997c150dd · 2026-07-09T02:53:00.000-07:00
PiperOrigin-RevId: 944425297
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel b/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",
     ],
 )
 
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/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,25 @@ 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 +340,11 @@ 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 +369,39 @@ 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 +417,7 @@ 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();
diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel b/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",
     ],
 )
 
diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
@@ -195,6 +195,27 @@ 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 +249,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 +472,33 @@ 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();

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ java_library(`
`44`	`44`	`"//:java_truth",`
`45`	`45`	`"@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",`
`46`	`46`	`"@maven//:com_google_guava_guava",`
	`47`	`+ "@maven//:com_google_protobuf_protobuf_java",`
`47`	`48`	`],`
`48`	`49`	`)`
`49`	`50`