Add capability to accept a CEL environment in CelCompilerTool

PiperOrigin-RevId: 730675994

Author: l46kok

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

@@ -64,8 +64,10 @@ java_library(
         "//common:source",
         "//common/types",
         "//common/types:type_providers",
+        "//compiler:compiler_builder",
         "//extensions",
         "//extensions:optional_library",
+        "//runtime",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
     ],

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

@@ -19,7 +19,9 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
@@ -35,8 +37,11 @@
 import dev.cel.common.types.OptionalType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.TypeParamType;
+import dev.cel.compiler.CelCompiler;
+import dev.cel.compiler.CelCompilerBuilder;
 import dev.cel.extensions.CelExtensions;
 import dev.cel.extensions.CelOptionalLibrary;
+import dev.cel.runtime.CelRuntimeBuilder;
 import java.util.Arrays;
 import java.util.Optional;
 
@@ -47,6 +52,18 @@
 @AutoValue
 public abstract class CelEnvironment {
 
+  @VisibleForTesting
+  static final ImmutableMap<String, CanonicalCelExtension> CEL_EXTENSION_CONFIG_MAP =
+      ImmutableMap.of(
+          "bindings", CanonicalCelExtension.BINDINGS,
+          "encoders", CanonicalCelExtension.ENCODERS,
+          "lists", CanonicalCelExtension.LISTS,
+          "math", CanonicalCelExtension.MATH,
+          "optional", CanonicalCelExtension.OPTIONAL,
+          "protos", CanonicalCelExtension.PROTOS,
+          "sets", CanonicalCelExtension.SETS,
+          "strings", CanonicalCelExtension.STRINGS);
+
   /** Environment source in textual format (ex: textproto, YAML). */
   public abstract Optional<Source> source();
 
@@ -126,12 +143,14 @@ public static Builder newBuilder() {
         .setFunctions(ImmutableSet.of());
   }
 
-  /** Extends the provided {@code cel} environment with this configuration. */
-  public Cel extend(Cel cel, CelOptions celOptions) throws CelEnvironmentException {
+  /** Extends the provided {@link CelCompiler} environment with this configuration. */
+  public CelCompiler extend(CelCompiler celCompiler, CelOptions celOptions)
+      throws CelEnvironmentException {
     try {
-      CelTypeProvider celTypeProvider = cel.getTypeProvider();
-      CelBuilder celBuilder =
-          cel.toCelBuilder()
+      CelTypeProvider celTypeProvider = celCompiler.getTypeProvider();
+      CelCompilerBuilder compilerBuilder =
+          celCompiler
+              .toCompilerBuilder()
               .setTypeProvider(celTypeProvider)
               .setContainer(container())
               .addVarDeclarations(
@@ -144,53 +163,58 @@ public Cel extend(Cel cel, CelOptions celOptions) throws CelEnvironmentException
                       .collect(toImmutableList()));
 
       if (!container().isEmpty()) {
-        celBuilder.setContainer(container());
+        compilerBuilder.setContainer(container());
       }
 
-      addAllExtensions(celBuilder, celOptions);
+      addAllCompilerExtensions(compilerBuilder, celOptions);
+
+      return compilerBuilder.build();
+    } catch (RuntimeException e) {
+      throw new CelEnvironmentException(e.getMessage(), e);
+    }
+  }
+
+  /** Extends the provided {@link Cel} environment with this configuration. */
+  public Cel extend(Cel cel, CelOptions celOptions) throws CelEnvironmentException {
+    try {
+      // Casting is necessary to only extend the compiler here
+      CelCompiler celCompiler = extend((CelCompiler) cel, celOptions);
+
+      CelRuntimeBuilder celRuntimeBuilder = cel.toRuntimeBuilder();
+      addAllRuntimeExtensions(celRuntimeBuilder, celOptions);
 
-      return celBuilder.build();
+      return CelFactory.combine(celCompiler, celRuntimeBuilder.build());
     } catch (RuntimeException e) {
       throw new CelEnvironmentException(e.getMessage(), e);
     }
   }
 
-  private void addAllExtensions(CelBuilder celBuilder, CelOptions celOptions) {
+  private void addAllCompilerExtensions(
+      CelCompilerBuilder celCompilerBuilder, CelOptions celOptions) {
     // TODO: Add capability to accept user defined exceptions
     for (ExtensionConfig extensionConfig : extensions()) {
-      switch (extensionConfig.name()) {
-        case "bindings":
-          celBuilder.addCompilerLibraries(CelExtensions.bindings());
-          break;
-        case "encoders":
-          celBuilder.addCompilerLibraries(CelExtensions.encoders());
-          celBuilder.addRuntimeLibraries(CelExtensions.encoders());
-          break;
-        case "math":
-          celBuilder.addCompilerLibraries(CelExtensions.math(celOptions));
-          celBuilder.addRuntimeLibraries(CelExtensions.math(celOptions));
-          break;
-        case "optional":
-          celBuilder.addCompilerLibraries(CelOptionalLibrary.INSTANCE);
-          celBuilder.addRuntimeLibraries(CelOptionalLibrary.INSTANCE);
-          break;
-        case "protos":
-          celBuilder.addCompilerLibraries(CelExtensions.protos());
-          break;
-        case "strings":
-          celBuilder.addCompilerLibraries(CelExtensions.strings());
-          celBuilder.addRuntimeLibraries(CelExtensions.strings());
-          break;
-        case "sets":
-          celBuilder.addCompilerLibraries(CelExtensions.sets(celOptions));
-          celBuilder.addRuntimeLibraries(CelExtensions.sets(celOptions));
-          break;
-        default:
-          throw new IllegalArgumentException("Unrecognized extension: " + extensionConfig.name());
-      }
+      CanonicalCelExtension extension = getExtensionOrThrow(extensionConfig.name());
+      extension.addCompilerExtension(celCompilerBuilder, celOptions);
+    }
+  }
+
+  private void addAllRuntimeExtensions(CelRuntimeBuilder celRuntimeBuilder, CelOptions celOptions) {
+    // TODO: Add capability to accept user defined exceptions
+    for (ExtensionConfig extensionConfig : extensions()) {
+      CanonicalCelExtension extension = getExtensionOrThrow(extensionConfig.name());
+      extension.addRuntimeExtension(celRuntimeBuilder, celOptions);
     }
   }
 
+  private static CanonicalCelExtension getExtensionOrThrow(String extensionName) {
+    CanonicalCelExtension extension = CEL_EXTENSION_CONFIG_MAP.get(extensionName);
+    if (extension == null) {
+      throw new IllegalArgumentException("Unrecognized extension: " + extensionName);
+    }
+
+    return extension;
+  }
+
   /** Represents a policy variable declaration. */
   @AutoValue
   public abstract static class VariableDecl {
@@ -529,4 +553,63 @@ public static ExtensionConfig of(String name, int version) {
       return newBuilder().setName(name).setVersion(version).build();
     }
   }
+
+  @VisibleForTesting
+  enum CanonicalCelExtension {
+    BINDINGS((compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.bindings())),
+    PROTOS((compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.protos())),
+    ENCODERS(
+        (compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.encoders()),
+        (runtimeBuilder, options) -> runtimeBuilder.addLibraries(CelExtensions.encoders())),
+    MATH(
+        (compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.math(options)),
+        (runtimeBuilder, options) -> runtimeBuilder.addLibraries(CelExtensions.math(options))),
+    OPTIONAL(
+        (compilerBuilder, options) -> compilerBuilder.addLibraries(CelOptionalLibrary.INSTANCE),
+        (runtimeBuilder, options) -> runtimeBuilder.addLibraries(CelOptionalLibrary.INSTANCE)),
+    STRINGS(
+        (compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.strings()),
+        (runtimeBuilder, options) -> runtimeBuilder.addLibraries(CelExtensions.strings())),
+    SETS(
+        (compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.sets(options)),
+        (runtimeBuilder, options) -> runtimeBuilder.addLibraries(CelExtensions.sets(options))),
+    LISTS(
+        (compilerBuilder, options) -> compilerBuilder.addLibraries(CelExtensions.lists()),
+        (runtimeBuilder, options) -> runtimeBuilder.addLibraries(CelExtensions.lists()));
+
+    @SuppressWarnings("ImmutableEnumChecker")
+    private final CompilerExtensionApplier compilerExtensionApplier;
+
+    @SuppressWarnings("ImmutableEnumChecker")
+    private final RuntimeExtensionApplier runtimeExtensionApplier;
+
+    interface CompilerExtensionApplier {
+      void apply(CelCompilerBuilder compilerBuilder, CelOptions options);
+    }
+
+    interface RuntimeExtensionApplier {
+      void apply(CelRuntimeBuilder runtimeBuilder, CelOptions options);
+    }
+
+    void addCompilerExtension(CelCompilerBuilder compilerBuilder, CelOptions options) {
+      compilerExtensionApplier.apply(compilerBuilder, options);
+    }
+
+    void addRuntimeExtension(CelRuntimeBuilder runtimeBuilder, CelOptions options) {
+      runtimeExtensionApplier.apply(runtimeBuilder, options);
+    }
+
+    CanonicalCelExtension(CompilerExtensionApplier compilerExtensionApplier) {
+      this(
+          compilerExtensionApplier,
+          (runtimeBuilder, options) -> {}); // no-op. Not all extensions augment the runtime.
+    }
+
+    CanonicalCelExtension(
+        CompilerExtensionApplier compilerExtensionApplier,
+        RuntimeExtensionApplier runtimeExtensionApplier) {
+      this.compilerExtensionApplier = compilerExtensionApplier;
+      this.runtimeExtensionApplier = runtimeExtensionApplier;
+    }
+  }
 }

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

@@ -47,7 +47,11 @@
 import org.yaml.snakeyaml.nodes.SequenceNode;
 import org.yaml.snakeyaml.nodes.Tag;
 
-final class CelEnvironmentYamlParser {
+/**
+ * CelEnvironmentYamlParser intakes a YAML document that describes the structure of a CEL
+ * environment, parses it then creates a {@link CelEnvironment}.
+ */
+public final class CelEnvironmentYamlParser {
   // Sentinel values to be returned for various declarations when parsing failure is encountered.
   private static final TypeDecl ERROR_TYPE_DECL = TypeDecl.create(ERROR);
   private static final VariableDecl ERROR_VARIABLE_DECL =

bundle/src/test/java/dev/cel/bundle/BUILD.bazel

@@ -10,7 +10,7 @@ java_library(
     testonly = True,
     srcs = glob(["*Test.java"]),
     resources = [
-        "//testing/environment:environment_yaml_files",
+        "//testing/environment:extended_env",
     ],
     deps = [
         "//:java_truth",

bundle/src/test/java/dev/cel/bundle/CelEnvironmentTest.java

@@ -16,7 +16,12 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableSet;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import dev.cel.bundle.CelEnvironment.CanonicalCelExtension;
+import dev.cel.bundle.CelEnvironment.ExtensionConfig;
+import dev.cel.common.CelAbstractSyntaxTree;
+import dev.cel.common.CelOptions;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -35,4 +40,32 @@ public void newBuilder_defaults() {
     assertThat(environment.variables()).isEmpty();
     assertThat(environment.functions()).isEmpty();
   }
+
+  @Test
+  public void extend_allExtensions() throws Exception {
+    ImmutableSet<ExtensionConfig> extensionConfigs =
+        ImmutableSet.of(
+            ExtensionConfig.of("bindings"),
+            ExtensionConfig.of("encoders"),
+            ExtensionConfig.of("lists"),
+            ExtensionConfig.of("math"),
+            ExtensionConfig.of("optional"),
+            ExtensionConfig.of("protos"),
+            ExtensionConfig.of("sets"),
+            ExtensionConfig.of("strings"));
+    CelEnvironment environment =
+        CelEnvironment.newBuilder().setExtensions(extensionConfigs).build();
+
+    Cel cel = environment.extend(CelFactory.standardCelBuilder().build(), CelOptions.DEFAULT);
+    CelAbstractSyntaxTree ast =
+        cel.compile(
+                "cel.bind(x, 10, math.greatest([1,x])) < int(' 11  '.trim()) &&"
+                    + " optional.none().orValue(true) && [].flatten() == []")
+            .getAst();
+    boolean result = (boolean) cel.createProgram(ast).eval();
+
+    assertThat(extensionConfigs.size()).isEqualTo(CelEnvironment.CEL_EXTENSION_CONFIG_MAP.size());
+    assertThat(extensionConfigs.size()).isEqualTo(CanonicalCelExtension.values().length);
+    assertThat(result).isTrue();
+  }
 }

bundle/src/test/java/dev/cel/bundle/CelEnvironmentYamlParserTest.java

@@ -74,6 +74,7 @@ public void environment_setExtensions() throws Exception {
         "extensions:\n"
             + "  - name: 'bindings'\n"
             + "  - name: 'encoders'\n"
+            + "  - name: 'lists'\n"
             + "  - name: 'math'\n"
             + "  - name: 'optional'\n"
             + "  - name: 'protos'\n"
@@ -91,6 +92,7 @@ public void environment_setExtensions() throws Exception {
                     ImmutableSet.of(
                         ExtensionConfig.of("bindings"),
                         ExtensionConfig.of("encoders"),
+                        ExtensionConfig.of("lists"),
                         ExtensionConfig.of("math"),
                         ExtensionConfig.of("optional"),
                         ExtensionConfig.of("protos"),
@@ -307,6 +309,31 @@ public void environment_setMessageVariable() throws Exception {
     assertThat(environment.extend(CEL_WITH_MESSAGE_TYPES, CelOptions.DEFAULT)).isNotNull();
   }
 
+  @Test
+  public void environment_setContainer() throws Exception {
+    String yamlConfig =
+        "container: google.rpc.context\n"
+            + "variables:\n"
+            + "- name: 'request'\n"
+            + "  type:\n"
+            + "    type_name: 'google.rpc.context.AttributeContext.Request'";
+
+    CelEnvironment environment = ENVIRONMENT_PARSER.parse(yamlConfig);
+
+    assertThat(environment)
+        .isEqualTo(
+            CelEnvironment.newBuilder()
+                .setContainer("google.rpc.context")
+                .setSource(environment.source().get())
+                .setVariables(
+                    ImmutableSet.of(
+                        VariableDecl.create(
+                            "request",
+                            TypeDecl.create("google.rpc.context.AttributeContext.Request"))))
+                .build());
+    assertThat(environment.extend(CEL_WITH_MESSAGE_TYPES, CelOptions.DEFAULT)).isNotNull();
+  }
+
   @Test
   public void environment_withInlinedVariableDecl() throws Exception {
     String yamlConfig =
@@ -675,8 +702,7 @@ private enum EnvironmentYamlResourceTestCase {
     private final String yamlFileContent;
     private final CelEnvironment expectedEnvironment;
 
-    EnvironmentYamlResourceTestCase(String yamlResourcePath, CelEnvironment expectedEnvironment)
-        throws RuntimeException {
+    EnvironmentYamlResourceTestCase(String yamlResourcePath, CelEnvironment expectedEnvironment) {
       try {
         this.yamlFileContent = readFile(yamlResourcePath);
       } catch (IOException e) {

compiler/src/main/java/dev/cel/compiler/tools/BUILD.bazel

@@ -15,14 +15,14 @@ java_binary(
     main_class = "dev.cel.compiler.tools.CelCompilerTool",
     neverlink = 1,
     deps = [
+        "//bundle:environment",
+        "//bundle:environment_exception",
+        "//bundle:environment_yaml_parser",
         "//common",
-        "//common:compiler_common",
         "//common:options",
         "//common:proto_ast",
         "//compiler",
         "//compiler:compiler_builder",
-        "//extensions",
-        "//extensions:optional_library",
         "//parser:macro",
         "@cel_spec//proto/cel/expr:checked_java_proto",
         "@maven//:com_google_guava_guava",