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

bundle/src/main/java/dev/cel/bundle/CelBuilder.java

@@ -39,6 +39,7 @@
 import dev.cel.runtime.CelRuntimeLibrary;
 import dev.cel.runtime.CelStandardFunctions;
 import java.util.function.Function;
+import org.jspecify.annotations.Nullable;
 
 /** Interface for building an instance of Cel. */
 public interface CelBuilder {
@@ -211,6 +212,9 @@ public interface CelBuilder {
   @CanIgnoreReturnValue
   CelBuilder setValueProvider(CelValueProvider celValueProvider);
 
+  /** Returns the configured {@link CelValueProvider}, or null if not set. */
+  @Nullable CelValueProvider getValueProvider();
+
   /**
    * Set the {@code typeProvider} for use with type-checking expressions.
    *

bundle/src/main/java/dev/cel/bundle/CelImpl.java

@@ -59,6 +59,7 @@
 import dev.cel.runtime.CelStandardFunctions;
 import java.util.Arrays;
 import java.util.function.Function;
+import org.jspecify.annotations.Nullable;
 
 /**
  * Implementation of the synchronous CEL stack.
@@ -317,6 +318,11 @@ public CelBuilder setValueProvider(CelValueProvider celValueProvider) {
       return this;
     }
 
+    @Override
+    public @Nullable CelValueProvider getValueProvider() {
+      return runtimeBuilder.getValueProvider();
+    }
+
     @Override
     @Deprecated
     public Builder setTypeProvider(TypeProvider typeProvider) {

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

@@ -31,6 +31,10 @@ java_library(
         "//common/navigation:common",
         "//common/navigation:mutable_navigation",
         "//common/types",
+        "//common/types:type_providers",
+        "//common/values:cel_value",
+        "//common/values:cel_value_provider",
+        "//common/values:proto_message_value",
         "//extensions:optional_library",
         "//optimizer:ast_optimizer",
         "//optimizer:mutable_ast",

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

@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import dev.cel.bundle.Cel;
+import dev.cel.bundle.CelBuilder;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelMutableAst;
 import dev.cel.common.CelSource;
@@ -42,7 +43,13 @@
 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.common.values.CelValue;
+import dev.cel.common.values.CelValueProvider;
+import dev.cel.common.values.ProtoMessageValue;
 import dev.cel.extensions.CelOptionalLibrary.Function;
 import dev.cel.optimizer.AstMutator;
 import dev.cel.optimizer.CelAstOptimizer;
@@ -95,8 +102,15 @@ private static CelMutableExpr newOptionalNoneExpr() {
   @Override
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
       throws CelOptimizationException {
-    // Override the environment's expected type to generally allow all subtrees to be folded.
-    Cel optimizerEnv = cel.toCelBuilder().setResultType(SimpleType.DYN).build();
+    CelBuilder builder = cel.toCelBuilder();
+    CelValueProvider valueProvider;
+    try {
+      valueProvider = builder.getValueProvider();
+    } catch (UnsupportedOperationException e) {
+      // Legacy runtime does not support getValueProvider and may throw.
+      valueProvider = null;
+    }
+    Cel optimizerEnv = builder.setResultType(SimpleType.DYN).build();
 
     CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
     int iterCount = 0;
@@ -123,7 +137,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
         if (!mutatedResult.isPresent()) {
           // Evaluate the call then fold
           try {
-            mutatedResult = maybeFold(optimizerEnv, mutableAst, foldableExpr);
+            mutatedResult = maybeFold(optimizerEnv, valueProvider, mutableAst, foldableExpr);
           } catch (CelEvaluationException e) {
             throw new CelOptimizationException(
                 "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(),
@@ -290,7 +304,10 @@ private static boolean isNestedComprehension(CelNavigableMutableExpr expr) {
   }
 
   private Optional<CelMutableAst> maybeFold(
-      Cel cel, CelMutableAst mutableAst, CelNavigableMutableExpr node)
+      Cel cel,
+      CelValueProvider valueProvider,
+      CelMutableAst mutableAst,
+      CelNavigableMutableExpr node)
       throws CelOptimizationException, CelEvaluationException {
     Object result;
     try {
@@ -305,25 +322,33 @@ 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(), valueProvider, optResult, mutableAst, node.expr());
     }
 
-    CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(result).orElse(null);
+    CelMutableExpr adaptedResult =
+        maybeAdaptEvaluatedResult(cel.getTypeProvider(), valueProvider, 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, CelValueProvider valueProvider, Object result) {
+    if (valueProvider != null && !(result instanceof CelValue)) {
+      result = valueProvider.celValueConverter().toRuntimeValue(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, valueProvider, evaluatedElement).orElse(null);
         if (adaptedExpr == null) {
           return Optional.empty();
         }
@@ -335,11 +360,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, valueProvider, entry.getKey()).orElse(null);
         if (adaptedKey == null) {
           return Optional.empty();
         }
-        CelMutableExpr adaptedValue = maybeAdaptEvaluatedResult(entry.getValue()).orElse(null);
+        CelMutableExpr adaptedValue =
+            maybeAdaptEvaluatedResult(typeProvider, valueProvider, entry.getValue()).orElse(null);
         if (adaptedValue == null) {
           return Optional.empty();
         }
@@ -364,14 +391,42 @@ private Optional<CelMutableExpr> maybeAdaptEvaluatedResult(Object result) {
               CelMutableExpr.ofConstant(CelConstant.ofValue(timestampStrArg)));
 
       return Optional.of(CelMutableExpr.ofCall(timestampCall));
+    } else if (result instanceof ProtoMessageValue) {
+      ProtoMessageValue structValue = (ProtoMessageValue) result;
+      List<CelMutableStruct.Entry> structEntries = new ArrayList<>();
+
+      String typeName = structValue.celType().name();
+      CelType optType = typeProvider.findType(typeName).orElse(null);
+      if (!(optType instanceof StructType)) {
+        return Optional.empty();
+      }
+      StructType structType = (StructType) optType;
+      for (String fieldName : structType.fieldNames()) {
+        Optional<?> fieldOpt = structValue.find(fieldName);
+        if (!fieldOpt.isPresent()) {
+          continue;
+        }
+        CelMutableExpr adaptedFieldExpr =
+            maybeAdaptEvaluatedResult(typeProvider, valueProvider, fieldOpt.get()).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,
+      CelValueProvider valueProvider,
+      Optional<?> optResult,
+      CelMutableAst mutableAst,
+      CelMutableExpr expr) {
     Object unwrappedResult = optResult.orElse(null);
     if (unwrappedResult == null) {
       if (isCallToFunction(expr, Function.OPTIONAL_NONE.getFunction())) {
@@ -387,7 +442,8 @@ private Optional<CelMutableAst> maybeRewriteOptional(
       return Optional.empty();
     }
 
-    CelMutableExpr adaptedResult = maybeAdaptEvaluatedResult(unwrappedResult).orElse(null);
+    CelMutableExpr adaptedResult =
+        maybeAdaptEvaluatedResult(typeProvider, valueProvider, 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/BUILD.bazel

@@ -42,8 +42,10 @@ java_library(
         "@maven//:junit_junit",
         "@maven//:com_google_testparameterinjector_test_parameter_injector",
         "//:java_truth",
+        "@cel_spec//proto/cel/expr/conformance/proto2:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",
         "@maven//:com_google_guava_guava",
+        "@maven//:com_google_protobuf_protobuf_java",
     ],
 )