Add support for stdlib subsetting via CelEnvironment

PiperOrigin-RevId: 776693221

Author: CEL Dev Team

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

@@ -61,6 +61,7 @@ java_library(
         ":required_fields_checker",
         "//:auto_value",
         "//bundle:cel",
+        "//checker:standard_decl",
         "//common:compiler_common",
         "//common:options",
         "//common:source",
@@ -69,6 +70,7 @@ java_library(
         "//compiler:compiler_builder",
         "//extensions",
         "//extensions:optional_library",
+        "//parser:macro",
         "//runtime",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",

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

@@ -17,6 +17,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.annotations.VisibleForTesting;
@@ -25,6 +26,8 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
+import dev.cel.checker.CelStandardDeclarations;
+import dev.cel.checker.CelStandardDeclarations.StandardFunction;
 import dev.cel.common.CelFunctionDecl;
 import dev.cel.common.CelOptions;
 import dev.cel.common.CelOverloadDecl;
@@ -41,6 +44,7 @@
 import dev.cel.compiler.CelCompilerBuilder;
 import dev.cel.extensions.CelExtensions;
 import dev.cel.extensions.CelOptionalLibrary;
+import dev.cel.parser.CelStandardMacro;
 import dev.cel.runtime.CelRuntimeBuilder;
 import java.util.Arrays;
 import java.util.Optional;
@@ -97,6 +101,9 @@ public abstract class CelEnvironment {
   /** New function declarations to add in the compilation environment. */
   public abstract ImmutableSet<FunctionDecl> functions();
 
+  /** Standard library subset (which macros, functions to include/exclude) */
+  public abstract Optional<LibrarySubset> standardLibrarySubset();
+
   /** Builder for {@link CelEnvironment}. */
   @AutoValue.Builder
   public abstract static class Builder {
@@ -141,6 +148,8 @@ public Builder setFunctions(FunctionDecl... functions) {
 
     public abstract Builder setFunctions(ImmutableSet<FunctionDecl> functions);
 
+    public abstract Builder setStandardLibrarySubset(Optional<LibrarySubset> stdLibrarySubset);
+
     @CheckReturnValue
     public abstract CelEnvironment build();
   }
@@ -180,6 +189,8 @@ public CelCompiler extend(CelCompiler celCompiler, CelOptions celOptions)
 
       addAllCompilerExtensions(compilerBuilder, celOptions);
 
+      applyStandardLibrarySubset(compilerBuilder);
+
       return compilerBuilder.build();
     } catch (RuntimeException e) {
       throw new CelEnvironmentException(e.getMessage(), e);
@@ -218,6 +229,75 @@ private void addAllRuntimeExtensions(CelRuntimeBuilder celRuntimeBuilder, CelOpt
     }
   }
 
+  private void applyStandardLibrarySubset(CelCompilerBuilder compilerBuilder) {
+    if (!standardLibrarySubset().isPresent()) {
+      return;
+    }
+
+    LibrarySubset librarySubset = standardLibrarySubset().get();
+    if (librarySubset.disabled()) {
+      compilerBuilder.setStandardEnvironmentEnabled(false);
+      return;
+    }
+
+    compilerBuilder.setStandardEnvironmentEnabled(true);
+
+    if (librarySubset.macrosDisabled()) {
+      compilerBuilder.setStandardMacros(ImmutableList.of());
+    } else if (!librarySubset.includedMacros().isEmpty()) {
+      compilerBuilder.setStandardMacros(
+          librarySubset.includedMacros().stream()
+              .map(name -> getStandardMacroOrThrow(name))
+              .toList());
+    } else if (!librarySubset.excludedMacros().isEmpty()) {
+      ImmutableSet<CelStandardMacro> set =
+          librarySubset.excludedMacros().stream()
+              .map(name -> getStandardMacroOrThrow(name))
+              .collect(toImmutableSet());
+      compilerBuilder.setStandardMacros(
+          CelStandardMacro.STANDARD_MACROS.stream().filter(macro -> !set.contains(macro)).toList());
+    }
+
+    if (!librarySubset.includedFunctions().isEmpty()) {
+      compilerBuilder.setStandardEnvironmentEnabled(false);
+      compilerBuilder.setStandardDeclarations(
+          CelStandardDeclarations.newBuilder()
+              .includeFunctions(
+                  librarySubset.includedFunctions().stream()
+                      .map(name -> getStandardFunctionOrThrow(name))
+                      .toList())
+              .build());
+    } else if (!librarySubset.excludedFunctions().isEmpty()) {
+      compilerBuilder.setStandardEnvironmentEnabled(false);
+      compilerBuilder.setStandardDeclarations(
+          CelStandardDeclarations.newBuilder()
+              .excludeFunctions(
+                  librarySubset.excludedFunctions().stream()
+                      .map(name -> getStandardFunctionOrThrow(name))
+                      .toList())
+              .build());
+
+    }
+  }
+
+  private static CelStandardMacro getStandardMacroOrThrow(String macroName) {
+    for (CelStandardMacro macro : CelStandardMacro.STANDARD_MACROS) {
+      if (macro.getFunction().equals(macroName)) {
+        return macro;
+      }
+    }
+    throw new IllegalArgumentException("unrecognized standard macro `" + macroName + "'");
+  }
+
+  private static StandardFunction getStandardFunctionOrThrow(String functionName) {
+    for (StandardFunction function : StandardFunction.values()) {
+      if (function.functionName().equals(functionName)) {
+        return function;
+      }
+    }
+    throw new IllegalArgumentException("unrecognized standard function `" + functionName + "'");
+  }
+
   private static CanonicalCelExtension getExtensionOrThrow(String extensionName) {
     CanonicalCelExtension extension = CEL_EXTENSION_CONFIG_MAP.get(extensionName);
     if (extension == null) {
@@ -624,4 +704,76 @@ void addRuntimeExtension(CelRuntimeBuilder runtimeBuilder, CelOptions options) {
       this.runtimeExtensionApplier = runtimeExtensionApplier;
     }
   }
+
+  /**
+   * LibrarySubset indicates a subset of the macros and function supported by a subsettable
+   * library.
+   */
+  @AutoValue
+  public abstract static class LibrarySubset {
+
+    /**
+     * Disabled indicates whether the library has been disabled, typically only used for
+     * default-enabled libraries like stdlib.
+     */
+    public abstract boolean disabled();
+
+    /** DisableMacros disables macros for the given library. */
+    public abstract boolean macrosDisabled();
+
+    /** IncludeMacros specifies a set of macro function names to include in the subset. */
+    public abstract ImmutableSet<String> includedMacros();
+
+    /**
+     * ExcludeMacros specifies a set of macro function names to exclude from the subset.
+     *
+     * <p>Note: if IncludedMacros is non-empty, then ExcludedMacros is ignored.
+     */
+    public abstract ImmutableSet<String> excludedMacros();
+
+    /**
+     * IncludeFunctions specifies a set of functions to include in the subset.
+     *
+     * <p>Note: the overloads specified in the subset need only specify their ID.
+     * <p>Note: if IncludedFunctions is non-empty, then ExcludedFunctions is ignored.
+     */
+    public abstract ImmutableSet<String> includedFunctions();
+
+    /**
+     * ExcludeFunctions specifies the set of functions to exclude from the subset.
+     *
+     * <p>Note: the overloads specified in the subset need only specify their ID.
+     */
+    public abstract ImmutableSet<String> excludedFunctions();
+
+    public static Builder newBuilder() {
+      return new AutoValue_CelEnvironment_LibrarySubset.Builder()
+          .setDisabled(true)
+          .setMacrosDisabled(false)
+          .setIncludedMacros(ImmutableSet.of())
+          .setExcludedMacros(ImmutableSet.of())
+          .setIncludedFunctions(ImmutableSet.of())
+          .setExcludedFunctions(ImmutableSet.of());
+    }
+
+    /** Builder for {@link LibrarySubset}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+      public abstract Builder setDisabled(boolean disabled);
+
+      public abstract Builder setMacrosDisabled(boolean disabled);
+
+      public abstract Builder setIncludedMacros(ImmutableSet<String> includedMacros);
+
+      public abstract Builder setExcludedMacros(ImmutableSet<String> excludedMacros);
+
+      public abstract Builder setIncludedFunctions(ImmutableSet<String> includedFunctions);
+
+      public abstract Builder setExcludedFunctions(ImmutableSet<String> excludedFunctions);
+
+      @CheckReturnValue
+      public abstract LibrarySubset build();
+    }
+  }
 }

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

@@ -29,6 +29,7 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import dev.cel.bundle.CelEnvironment.ExtensionConfig;
 import dev.cel.bundle.CelEnvironment.FunctionDecl;
+import dev.cel.bundle.CelEnvironment.LibrarySubset;
 import dev.cel.bundle.CelEnvironment.OverloadDecl;
 import dev.cel.bundle.CelEnvironment.TypeDecl;
 import dev.cel.bundle.CelEnvironment.VariableDecl;
@@ -38,6 +39,7 @@
 import dev.cel.common.formats.YamlHelper.YamlNodeType;
 import dev.cel.common.formats.YamlParserContextImpl;
 import dev.cel.common.internal.CelCodePointArray;
+import java.util.Optional;
 import org.jspecify.annotations.Nullable;
 import org.yaml.snakeyaml.DumperOptions.FlowStyle;
 import org.yaml.snakeyaml.nodes.MappingNode;
@@ -337,6 +339,79 @@ private static ExtensionConfig parseExtension(ParserContext<Node> ctx, Node node
     return builder.build();
   }
 
+  private static LibrarySubset parseLibrarySubset(
+      ParserContext<Node> ctx, Node node, boolean disabled) {
+    LibrarySubset.Builder builder = LibrarySubset.newBuilder().setDisabled(disabled);
+    MappingNode subsetMap = (MappingNode) node;
+    for (NodeTuple nodeTuple : subsetMap.getValue()) {
+      Node keyNode = nodeTuple.getKeyNode();
+      long keyId = ctx.collectMetadata(keyNode);
+      Node valueNode = nodeTuple.getValueNode();
+      String keyName = ((ScalarNode) keyNode).getValue();
+      System.out.println("keyId = " + keyId + " keyName = " + keyName);
+      switch (keyName) {
+        case "disabled":
+          builder.setDisabled(newBoolean(ctx, valueNode));
+          break;
+        case "disable_macros":
+          builder.setMacrosDisabled(newBoolean(ctx, valueNode));
+          break;
+        case "include_macros":
+          builder.setIncludedMacros(parseNameSet(ctx, keyName, valueNode));
+          break;
+        case "exclude_macros":
+          builder.setExcludedMacros(parseNameSet(ctx, keyName, valueNode));
+          break;
+        case "include_functions":
+          builder.setIncludedFunctions(parseNameSet(ctx, keyName, valueNode));
+          break;
+        case "exclude_functions":
+          builder.setExcludedFunctions(parseNameSet(ctx, keyName, valueNode));
+          break;
+        default:
+          ctx.reportError(keyId, String.format("Unsupported library subset tag: %s", keyName));
+          break;
+      }
+    }
+    return builder.build();
+  }
+
+  private static ImmutableSet<String> parseNameSet(
+      ParserContext<Node> ctx, String listKey, Node node) {
+    long valueId = ctx.collectMetadata(node);
+    if (!assertYamlType(ctx, valueId, node, YamlNodeType.LIST)) {
+      return ImmutableSet.of(ERROR);
+    }
+
+    ImmutableSet.Builder<String> builder = ImmutableSet.builder();
+    SequenceNode nameListNode = (SequenceNode) node;
+    for (Node elementNode : nameListNode.getValue()) {
+      long elementId = ctx.collectMetadata(elementNode);
+      System.out.println("NODE:" + elementNode);
+      if (!assertYamlType(ctx, elementId, elementNode, YamlNodeType.MAP)) {
+        return ImmutableSet.of(ERROR);
+      }
+
+      for (NodeTuple nodeTuple : ((MappingNode) elementNode).getValue()) {
+        Node keyNode = nodeTuple.getKeyNode();
+        long keyId = ctx.collectMetadata(keyNode);
+        Node valueNode = nodeTuple.getValueNode();
+        String keyName = ((ScalarNode) keyNode).getValue();
+        System.out.println("IN LIST keyId = " + keyId + " keyName = " + keyName);
+        switch (keyName) {
+          case "name":
+            builder.add(newString(ctx, valueNode));
+            break;
+          default:
+            ctx.reportError(
+                keyId, String.format("Unsupported %s element tag: %s", listKey, keyName));
+            break;
+        }
+      }
+    }
+    return builder.build();
+  }
+
   private static @Nullable Integer tryParse(String str) {
     try {
       return Integer.parseInt(str);
@@ -477,6 +552,10 @@ private CelEnvironment.Builder parseConfig(ParserContext<Node> ctx, Node node) {
           case "extensions":
             builder.addExtensions(parseExtensions(ctx, valueNode));
             break;
+          case "stdlib":
+            builder.setStandardLibrarySubset(
+                Optional.of(parseLibrarySubset(ctx, valueNode, /* disabled= */ false)));
+            break;
           default:
             ctx.reportError(id, "Unknown config tag: " + fieldName);
             // continue handling the rest of the nodes

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

@@ -15,6 +15,9 @@
 package dev.cel.bundle;
 
 import com.google.common.collect.ImmutableMap;
+import dev.cel.bundle.CelEnvironment.LibrarySubset;
+import java.util.List;
+import java.util.Set;
 import org.yaml.snakeyaml.DumperOptions;
 import org.yaml.snakeyaml.Yaml;
 import org.yaml.snakeyaml.nodes.Node;
@@ -45,6 +48,7 @@ private CelEnvironmentYamlSerializer() {
     this.multiRepresenters.put(CelEnvironment.TypeDecl.class, new RepresentTypeDecl());
     this.multiRepresenters.put(
         CelEnvironment.ExtensionConfig.class, new RepresentExtensionConfig());
+    this.multiRepresenters.put(CelEnvironment.LibrarySubset.class, new RepresetLibrarySubset());
   }
 
   public static String toYaml(CelEnvironment environment) {
@@ -74,6 +78,9 @@ public Node representData(Object data) {
       if (!environment.functions().isEmpty()) {
         configMap.put("functions", environment.functions().asList());
       }
+      if (environment.standardLibrarySubset().isPresent()) {
+        configMap.put("stdlib", environment.standardLibrarySubset().get());
+      }
       return represent(configMap.buildOrThrow());
     }
   }
@@ -143,4 +150,35 @@ public Node representData(Object data) {
       return represent(configMap.buildOrThrow());
     }
   }
+
+  private final class RepresetLibrarySubset implements Represent {
+    @Override
+    public Node representData(Object data) {
+      LibrarySubset librarySubset = (LibrarySubset) data;
+      ImmutableMap.Builder<String, Object> configMap = new ImmutableMap.Builder<>();
+      if (librarySubset.disabled()) {
+        configMap.put("disabled", true);
+      }
+      if (librarySubset.macrosDisabled()) {
+        configMap.put("disable_macros", true);
+      }
+      if (!librarySubset.includedMacros().isEmpty()) {
+        configMap.put("include_macros", toListOfNames(librarySubset.includedMacros()));
+      }
+      if (!librarySubset.excludedMacros().isEmpty()) {
+        configMap.put("exclude_macros", toListOfNames(librarySubset.excludedMacros()));
+      }
+      if (!librarySubset.includedFunctions().isEmpty()) {
+        configMap.put("include_functions", toListOfNames(librarySubset.includedFunctions()));
+      }
+      if (!librarySubset.excludedFunctions().isEmpty()) {
+        configMap.put("exclude_functions", toListOfNames(librarySubset.excludedFunctions()));
+      }
+      return represent(configMap.buildOrThrow());
+    }
+
+    private List<ImmutableMap<String, String>> toListOfNames(Set<String> names) {
+      return names.stream().sorted().map(name -> ImmutableMap.of("name", name)).toList();
+    }
+  }
 }