Add CelStandardDeclaration to allow environment subsetting for type-checker

PiperOrigin-RevId: 672004289

Author: l46kok

checker/BUILD.bazel

@@ -51,3 +51,8 @@ java_library(
     name = "proto_expr_visitor",
     exports = ["//checker/src/main/java/dev/cel/checker:proto_expr_visitor"],
 )
+
+java_library(
+    name = "standard_decl",
+    exports = ["//checker/src/main/java/dev/cel/checker:standard_decl"],
+)

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

@@ -28,7 +28,6 @@ CHECKER_LEGACY_ENV_SOURCES = [
     "ExprChecker.java",
     "ExprVisitor.java",
     "InferenceContext.java",
-    "Standard.java",
     "TypeFormatter.java",
     "TypeProvider.java",
     "Types.java",
@@ -68,6 +67,7 @@ java_library(
         ":checker_builder",
         ":checker_legacy_environment",
         ":proto_type_mask",
+        ":standard_decl",
         ":type_provider_legacy_impl",
         "//:auto_value",
         "//common",
@@ -97,6 +97,7 @@ java_library(
     deps = [
         ":checker_legacy_environment",
         ":proto_type_mask",
+        ":standard_decl",
         "//common",
         "//common:compiler_common",
         "//common:options",
@@ -166,6 +167,7 @@ java_library(
     ],
     deps = [
         ":cel_ident_decl",
+        ":standard_decl",
         "//:auto_value",
         "//common",
         "//common:compiler_common",
@@ -217,3 +219,22 @@ java_library(
         "//common:proto_ast",
     ],
 )
+
+java_library(
+    name = "standard_decl",
+    srcs = [
+        "CelStandardDeclaration.java",
+    ],
+    tags = [
+    ],
+    deps = [
+        ":cel_ident_decl",
+        "//common:compiler_common",
+        "//common/types",
+        "//common/types:cel_types",
+        "//common/types:type_providers",
+        "//parser:operator",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven//:com_google_guava_guava",
+    ],
+)

checker/src/main/java/dev/cel/checker/CelCheckerBuilder.java

@@ -152,6 +152,14 @@ public interface CelCheckerBuilder {
   @CanIgnoreReturnValue
   CelCheckerBuilder setStandardEnvironmentEnabled(boolean value);
 
+  /**
+   * Override the standard declarations for the type-checker. This can be used to subset the
+   * standard environment to only expose the desired declarations to the type-checker. {@link
+   * #setStandardEnvironmentEnabled(boolean)} must be set to false for this to take effect.
+   */
+  @CanIgnoreReturnValue
+  CelCheckerBuilder setStandardDeclarations(CelStandardDeclaration standardDeclarations);
+
   /** Adds one or more libraries for parsing and type-checking. */
   @CanIgnoreReturnValue
   CelCheckerBuilder addLibraries(CelCheckerLibrary... libraries);

checker/src/main/java/dev/cel/checker/CelCheckerLegacyImpl.java

@@ -78,6 +78,8 @@ public final class CelCheckerLegacyImpl implements CelChecker, EnvVisitable {
   private final CelTypeProvider celTypeProvider;
   private final boolean standardEnvironmentEnabled;
 
+  private final CelStandardDeclaration overriddenStandardDeclarations;
+
   // Builder is mutable by design. APIs must make defensive copies in and out of this class.
   @SuppressWarnings("Immutable")
   private final Builder checkerBuilder;
@@ -137,6 +139,8 @@ private Env getEnv(Errors errors) {
     Env env;
     if (standardEnvironmentEnabled) {
       env = Env.standard(errors, typeProvider, celOptions);
+    } else if (overriddenStandardDeclarations != null) {
+      env = Env.standard(overriddenStandardDeclarations, errors, typeProvider, celOptions);
     } else {
       env = Env.unconfigured(errors, typeProvider, celOptions);
     }
@@ -165,6 +169,7 @@ public static final class Builder implements CelCheckerBuilder {
     private TypeProvider customTypeProvider;
     private CelTypeProvider celTypeProvider;
     private boolean standardEnvironmentEnabled;
+    private CelStandardDeclaration standardDeclarations;
 
     @Override
     public CelCheckerBuilder setOptions(CelOptions celOptions) {
@@ -320,7 +325,13 @@ public CelCheckerBuilder addFileTypes(FileDescriptorSet fileDescriptorSet) {
 
     @Override
     public CelCheckerBuilder setStandardEnvironmentEnabled(boolean value) {
-      standardEnvironmentEnabled = value;
+      this.standardEnvironmentEnabled = value;
+      return this;
+    }
+
+    @Override
+    public CelCheckerBuilder setStandardDeclarations(CelStandardDeclaration standardDeclarations) {
+      this.standardDeclarations = checkNotNull(standardDeclarations);
       return this;
     }
 
@@ -372,6 +383,11 @@ ImmutableSet.Builder<CelCheckerLibrary> getCheckerLibraries() {
     @Override
     @CheckReturnValue
     public CelCheckerLegacyImpl build() {
+      if (standardEnvironmentEnabled && standardDeclarations != null) {
+        throw new IllegalArgumentException(
+            "setStandardEnvironmentEnabled must be set to false to override standard"
+                + " declarations.");
+      }
       // Add libraries, such as extensions
       celCheckerLibraries.build().forEach(celLibrary -> celLibrary.setCheckerOptions(this));
 
@@ -430,6 +446,7 @@ public CelCheckerLegacyImpl build() {
           legacyProvider,
           messageTypeProvider,
           standardEnvironmentEnabled,
+          standardDeclarations,
           this);
     }
 
@@ -478,6 +495,7 @@ private CelCheckerLegacyImpl(
       TypeProvider typeProvider,
       CelTypeProvider celTypeProvider,
       boolean standardEnvironmentEnabled,
+      CelStandardDeclaration overriddenStandardDeclarations,
       Builder checkerBuilder) {
     this.celOptions = celOptions;
     this.container = container;
@@ -487,6 +505,7 @@ private CelCheckerLegacyImpl(
     this.typeProvider = typeProvider;
     this.celTypeProvider = celTypeProvider;
     this.standardEnvironmentEnabled = standardEnvironmentEnabled;
+    this.overriddenStandardDeclarations = overriddenStandardDeclarations;
     this.checkerBuilder = new Builder(checkerBuilder);
   }