Allow expected result type to be configured per AST validator instance

PiperOrigin-RevId: 831021405

Author: l46kok

validator/src/main/java/dev/cel/validator/BUILD.bazel

@@ -72,6 +72,8 @@ java_library(
         "//common:compiler_common",
         "//common:source_location",
         "//common/navigation",
+        "//common/types",
+        "//common/types:type_providers",
         "@maven//:com_google_guava_guava",
     ],
 )

validator/src/main/java/dev/cel/validator/CelAstValidator.java

@@ -21,15 +21,22 @@
 import dev.cel.common.CelSource;
 import dev.cel.common.CelSourceLocation;
 import dev.cel.common.navigation.CelNavigableAst;
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.SimpleType;
 import java.util.Optional;
 
 /** Public interface for performing a single, custom validation on an AST. */
 public interface CelAstValidator {
 
   void validate(CelNavigableAst navigableAst, Cel cel, IssuesFactory issuesFactory);
 
+  /** Enforces a specific expected result type during validation, if set. */
+  default CelType expectedResultType() {
+    return SimpleType.DYN;
+  }
+
   /** Factory for populating issues while performing AST validation. */
-  public final class IssuesFactory {
+  final class IssuesFactory {
     private final ImmutableList.Builder<CelIssue> issuesBuilder;
     private final CelNavigableAst navigableAst;
 

validator/src/main/java/dev/cel/validator/CelValidatorImpl.java

@@ -48,9 +48,11 @@ public CelValidationResult validate(CelAbstractSyntaxTree ast) {
     ImmutableList.Builder<CelIssue> issueBuilder = ImmutableList.builder();
 
     for (CelAstValidator validator : astValidators) {
+      Cel celEnv = this.cel.toCelBuilder().setResultType(validator.expectedResultType()).build();
+
       CelNavigableAst navigableAst = CelNavigableAst.fromAst(ast);
       IssuesFactory issuesFactory = new IssuesFactory(navigableAst);
-      validator.validate(navigableAst, cel, issuesFactory);
+      validator.validate(navigableAst, celEnv, issuesFactory);
       issueBuilder.addAll(issuesFactory.getIssues());
     }
 

validator/src/main/java/dev/cel/validator/validators/BUILD.bazel

@@ -17,7 +17,8 @@ java_library(
     ],
     deps = [
         ":literal_validator",
-        "@maven//:com_google_protobuf_protobuf_java",
+        "//common/types",
+        "//common/types:type_providers",
     ],
 )
 
@@ -30,7 +31,8 @@ java_library(
     ],
     deps = [
         ":literal_validator",
-        "@maven//:com_google_protobuf_protobuf_java",
+        "//common/types",
+        "//common/types:type_providers",
     ],
 )
 
@@ -116,6 +118,7 @@ java_library(
         "//common/ast",
         "//common/ast:expr_factory",
         "//common/navigation",
+        "//common/types:type_providers",
         "//runtime",
         "//validator:ast_validator",
         "@maven//:com_google_errorprone_error_prone_annotations",

validator/src/main/java/dev/cel/validator/validators/DurationLiteralValidator.java

@@ -14,14 +14,17 @@
 
 package dev.cel.validator.validators;
 
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.SimpleType;
 import java.time.Duration;
 
 /** DurationLiteralValidator ensures that duration literal arguments are valid. */
 public final class DurationLiteralValidator extends LiteralValidator {
   public static final DurationLiteralValidator INSTANCE =
-      new DurationLiteralValidator("duration", Duration.class);
+      new DurationLiteralValidator("duration", Duration.class, SimpleType.DURATION);
 
-  private DurationLiteralValidator(String functionName, Class<?> expectedResultType) {
-    super(functionName, expectedResultType);
+  private DurationLiteralValidator(
+      String functionName, Class<?> expectedJavaType, CelType expectedResultType) {
+    super(functionName, expectedJavaType, expectedResultType);
   }
 }

validator/src/main/java/dev/cel/validator/validators/HomogeneousLiteralValidator.java

@@ -38,16 +38,16 @@ public final class HomogeneousLiteralValidator implements CelAstValidator {
   private final ImmutableSet<String> exemptFunctions;
 
   /**
-   * Construct a new instance of {@link HomogeneousLiteralValidator}. This validator will not for
-   * functions in {@code exemptFunctions}.
+   * Construct a new instance of {@link HomogeneousLiteralValidator}. This validator will not run
+   * for functions in {@code exemptFunctions}.
    */
   public static HomogeneousLiteralValidator newInstance(Iterable<String> exemptFunctions) {
     return new HomogeneousLiteralValidator(exemptFunctions);
   }
 
   /**
-   * Construct a new instance of {@link HomogeneousLiteralValidator}. This validator will not for
-   * functions in {@code exemptFunctions}.
+   * Construct a new instance of {@link HomogeneousLiteralValidator}. This validator will not run
+   * for functions in {@code exemptFunctions}.
    */
   public static HomogeneousLiteralValidator newInstance(String... exemptFunctions) {
     return newInstance(Arrays.asList(exemptFunctions));

validator/src/main/java/dev/cel/validator/validators/LiteralValidator.java

@@ -24,6 +24,7 @@
 import dev.cel.common.ast.CelExprFactory;
 import dev.cel.common.navigation.CelNavigableAst;
 import dev.cel.common.navigation.CelNavigableExpr;
+import dev.cel.common.types.CelType;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.validator.CelAstValidator;
 
@@ -34,13 +35,21 @@
  */
 public abstract class LiteralValidator implements CelAstValidator {
   private final String functionName;
-  private final Class<?> expectedResultType;
+  private final Class<?> expectedJavaType;
+  private final CelType expectedResultType;
 
-  protected LiteralValidator(String functionName, Class<?> expectedResultType) {
+  protected LiteralValidator(
+      String functionName, Class<?> expectedJavaType, CelType expectedResultType) {
     this.functionName = functionName;
+    this.expectedJavaType = expectedJavaType;
     this.expectedResultType = expectedResultType;
   }
 
+  @Override
+  public CelType expectedResultType() {
+    return expectedResultType;
+  }
+
   @Override
   public void validate(CelNavigableAst navigableAst, Cel cel, IssuesFactory issuesFactory) {
     CelExprFactory exprFactory = CelExprFactory.newInstance();
@@ -61,7 +70,7 @@ public void validate(CelNavigableAst navigableAst, Cel cel, IssuesFactory issues
               CelExpr callExpr =
                   exprFactory.newGlobalCall(functionName, exprFactory.newConstant(expr.constant()));
               try {
-                evaluateExpr(cel, callExpr, expectedResultType);
+                evaluateExpr(cel, callExpr, expectedJavaType);
               } catch (Exception e) {
                 issuesFactory.addError(
                     expr.id(),
@@ -72,18 +81,18 @@ public void validate(CelNavigableAst navigableAst, Cel cel, IssuesFactory issues
   }
 
   @CanIgnoreReturnValue
-  private static Object evaluateExpr(Cel cel, CelExpr expr, Class<?> expectedResultType)
+  private static Object evaluateExpr(Cel cel, CelExpr expr, Class<?> expectedJavaType)
       throws CelValidationException, CelEvaluationException {
     CelAbstractSyntaxTree ast =
         CelAbstractSyntaxTree.newParsedAst(expr, CelSource.newBuilder().build());
     ast = cel.check(ast).getAst();
     Object result = cel.createProgram(ast).eval();
 
-    if (!expectedResultType.isInstance(result)) {
+    if (!expectedJavaType.isInstance(result)) {
       throw new IllegalStateException(
           String.format(
               "Expected %s type but got %s instead",
-              expectedResultType.getName(), result.getClass().getName()));
+              expectedJavaType.getName(), result.getClass().getName()));
     }
     return result;
   }

validator/src/main/java/dev/cel/validator/validators/TimestampLiteralValidator.java

@@ -14,14 +14,17 @@
 
 package dev.cel.validator.validators;
 
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.SimpleType;
 import java.time.Instant;
 
 /** TimestampLiteralValidator ensures that timestamp literal arguments are valid. */
 public final class TimestampLiteralValidator extends LiteralValidator {
   public static final TimestampLiteralValidator INSTANCE =
-      new TimestampLiteralValidator("timestamp", Instant.class);
+      new TimestampLiteralValidator("timestamp", Instant.class, SimpleType.TIMESTAMP);
 
-  private TimestampLiteralValidator(String functionName, Class<?> expectedResultType) {
-    super(functionName, expectedResultType);
+  private TimestampLiteralValidator(
+      String functionName, Class<?> expectedJavaType, CelType expectedResultType) {
+    super(functionName, expectedJavaType, expectedResultType);
   }
 }

validator/src/test/java/dev/cel/validator/validators/TimestampLiteralValidatorTest.java

@@ -200,4 +200,23 @@ public void parentIsNotCallExpr_doesNotThrow(String source) throws Exception {
     assertThat(result.hasError()).isFalse();
     assertThat(result.getAllIssues()).isEmpty();
   }
+
+  @Test
+  public void env_withSetResultType_success() throws Exception {
+    Cel cel =
+        CelFactory.standardCelBuilder()
+            .setOptions(CelOptions.current().enableTimestampEpoch(true).build())
+            .setResultType(SimpleType.BOOL)
+            .build();
+    CelValidator validator =
+        CelValidatorFactory.standardCelValidatorBuilder(cel)
+            .addAstValidators(TimestampLiteralValidator.INSTANCE)
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("timestamp(123) == timestamp(123)").getAst();
+
+    CelValidationResult result = validator.validate(ast);
+
+    assertThat(result.hasError()).isFalse();
+    assertThat(result.getAllIssues()).isEmpty();
+  }
 }