Refactor CelExtensionLibrary to centralize version definitions.

This allows to have a better support when the extension library has multiple versions. Now the knowledge of all supported extension versions and
which version is the latest is encapsulated inside the corresponding Cel*Extension class.

PiperOrigin-RevId: 786864243

Author: dmitriplotnikov

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

@@ -23,7 +23,6 @@
 import dev.cel.expr.Decl.FunctionDecl.Overload;
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ArrayListMultimap;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ListMultimap;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -97,7 +96,8 @@ public abstract class CelEnvironmentExporter {
    */
   abstract int maxExcludedStandardFunctionOverloads();
 
-  abstract ImmutableSet<CelExtensionLibrary> extensionLibraries();
+  abstract ImmutableSet<CelExtensionLibrary<? extends CelExtensionLibrary.FeatureSet>>
+      extensionLibraries();
 
   /** Builder for {@link CelEnvironmentExporter}. */
   @AutoValue.Builder
@@ -107,22 +107,21 @@ public abstract static class Builder {
 
     public abstract Builder setMaxExcludedStandardFunctionOverloads(int count);
 
-    abstract ImmutableSet.Builder<CelExtensionLibrary> extensionLibrariesBuilder();
+    abstract ImmutableSet.Builder<CelExtensionLibrary<? extends CelExtensionLibrary.FeatureSet>>
+        extensionLibrariesBuilder();
 
     @CanIgnoreReturnValue
     public Builder addStandardExtensions(CelOptions options) {
       addExtensionLibraries(
-          CelExtensions.math(options, 0),
-          CelExtensions.math(options, 1),
-          CelExtensions.math(options, 2),
-          CelExtensions.lists(0),
-          CelExtensions.lists(1),
-          CelExtensions.lists(2));
+          CelExtensions.getExtensionLibrary("math", options),
+          CelExtensions.getExtensionLibrary("lists", options));
+      // TODO: add support for remaining standard extensions
       return this;
     }
 
     @CanIgnoreReturnValue
-    public Builder addExtensionLibraries(CelExtensionLibrary... libraries) {
+    public Builder addExtensionLibraries(
+        CelExtensionLibrary<? extends CelExtensionLibrary.FeatureSet>... libraries) {
       extensionLibrariesBuilder().add(libraries);
       return this;
     }
@@ -216,52 +215,59 @@ public void visitMacro(CelMacro macro) {
    */
   private void addExtensionConfigsAndRemoveFromInventory(
       CelEnvironment.Builder envBuilder, Set<Object> inventory) {
-    ImmutableList<CelExtensionLibrary> libraries =
-        ImmutableList.sortedCopyOf(
-            Comparator.comparing(CelExtensionLibrary::getName)
-                .thenComparing(CelExtensionLibrary::getVersion)
-                .reversed(),
-            extensionLibraries());
+    ArrayList<NamedFeatureSet> featureSets = new ArrayList<>();
+
+    for (CelExtensionLibrary<? extends CelExtensionLibrary.FeatureSet> extensionLibrary :
+        extensionLibraries()) {
+      for (CelExtensionLibrary.FeatureSet featureSet : extensionLibrary.versions()) {
+        featureSets.add(NamedFeatureSet.create(extensionLibrary.name(), featureSet));
+      }
+    }
+
+    featureSets.sort(
+        Comparator.comparing(NamedFeatureSet::name)
+            .thenComparing(nfs -> nfs.featureSet().version())
+            .reversed());
 
     Set<String> includedExtensions = new HashSet<>();
-    for (CelExtensionLibrary library : libraries) {
-      if (includedExtensions.contains(library.getName())) {
+    for (NamedFeatureSet lib : featureSets) {
+      if (includedExtensions.contains(lib.name())) {
         // We only need to infer the highest version library, so we can skip lower versions
         continue;
       }
 
-      if (checkIfExtensionIsIncludedAndRemoveFromInventory(inventory, library)) {
-        envBuilder.addExtensions(ExtensionConfig.of(library.getName(), library.getVersion()));
-        includedExtensions.add(library.getName());
+      if (checkIfExtensionIsIncludedAndRemoveFromInventory(inventory, lib.featureSet())) {
+        envBuilder.addExtensions(ExtensionConfig.of(lib.name(), lib.featureSet().version()));
+        includedExtensions.add(lib.name());
       }
     }
   }
 
   private boolean checkIfExtensionIsIncludedAndRemoveFromInventory(
-      Set<Object> inventory, CelExtensionLibrary library) {
-    ImmutableSet<CelFunctionDecl> functions = library.getFunctions();
-    ArrayList<Object> includedItems = new ArrayList<>(functions.size());
+      Set<Object> inventory, CelExtensionLibrary.FeatureSet featureSet) {
+    ImmutableSet<CelFunctionDecl> functions = featureSet.functions();
+    ArrayList<Object> includedFeatures = new ArrayList<>(functions.size());
     for (CelFunctionDecl function : functions) {
       for (CelOverloadDecl overload : function.overloads()) {
-        NamedOverload item = NamedOverload.create(function.name(), overload);
-        if (!inventory.contains(item)) {
+        NamedOverload feature = NamedOverload.create(function.name(), overload);
+        if (!inventory.contains(feature)) {
           return false;
         }
-        includedItems.add(item);
+        includedFeatures.add(feature);
       }
     }
 
-    ImmutableSet<CelMacro> macros = library.getMacros();
+    ImmutableSet<CelMacro> macros = featureSet.macros();
     for (CelMacro macro : macros) {
       if (!inventory.contains(macro)) {
         return false;
       }
-      includedItems.add(macro);
+      includedFeatures.add(macro);
     }
 
     // TODO - Add checks for variables.
 
-    inventory.removeAll(includedItems);
+    inventory.removeAll(includedFeatures);
     return true;
   }
 
@@ -424,5 +430,18 @@ static NamedOverload create(String functionName, CelOverloadDecl overload) {
       return new AutoValue_CelEnvironmentExporter_NamedOverload(functionName, overload);
     }
   }
-}
 
+  /**
+   * Wrapper for CelExtensionLibrary.FeatureSet, associating it with the corresponding library name.
+   */
+  @AutoValue
+  abstract static class NamedFeatureSet {
+    abstract String name();
+
+    abstract CelExtensionLibrary.FeatureSet featureSet();
+
+    static NamedFeatureSet create(String name, CelExtensionLibrary.FeatureSet featureSet) {
+      return new AutoValue_CelEnvironmentExporter_NamedFeatureSet(name, featureSet);
+    }
+  }
+}

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

@@ -39,6 +39,7 @@ java_library(
         ":sets_function",
         ":strings",
         "//common:options",
+        "//extensions:extension_library",
         "@maven//:com_google_guava_guava",
     ],
 )

extensions/src/main/java/dev/cel/extensions/CelExtensionLibrary.java

@@ -17,26 +17,51 @@
 import com.google.common.collect.ImmutableSet;
 import dev.cel.common.CelFunctionDecl;
 import dev.cel.parser.CelMacro;
+import java.util.Comparator;
 
 /**
  * Interface for defining CEL extension libraries.
  *
  * <p>An extension library is a collection of CEL functions, variables, and macros that can be added
  * to a CEL environment to provide additional functionality.
  */
-public interface CelExtensionLibrary {
+public interface CelExtensionLibrary<T extends CelExtensionLibrary.FeatureSet> {
 
   /** Returns the name of the extension library. */
-  String getName();
+  String name();
 
-  /** Returns the extension library version or -1 if unspecified. */
-  int getVersion();
+  ImmutableSet<T> versions();
 
-  /** Returns the set of function declarations defined by this extension library. */
-  ImmutableSet<CelFunctionDecl> getFunctions();
+  default T latest() {
+    return versions().stream().max(Comparator.comparing(FeatureSet::version)).get();
+  }
 
-  /** Returns the set of macros defined by this extension library. */
-  ImmutableSet<CelMacro> getMacros();
+  default T version(int version) {
+    if (version == Integer.MAX_VALUE) {
+      return latest();
+    }
 
-  // TODO - Add a method for variables.
+    for (T v : versions()) {
+      if (v.version() == version) {
+        return v;
+      }
+    }
+    throw new IllegalArgumentException("Unsupported '" + name() + "' extension version " + version);
+  }
+
+  /**
+   * Interface for defining a version of a CEL extension library.
+   */
+  interface FeatureSet {
+    /** Returns the extension library version or -1 if unspecified. */
+    int version();
+
+    /** Returns the set of function declarations defined by this extension library. */
+    ImmutableSet<CelFunctionDecl> functions();
+
+    /** Returns the set of macros defined by this extension library. */
+    ImmutableSet<CelMacro> macros();
+
+    // TODO - Add a method for variables.
+  }
 }

extensions/src/main/java/dev/cel/extensions/CelExtensions.java

@@ -20,7 +20,6 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import dev.cel.common.CelOptions;
-import dev.cel.extensions.CelListsExtensions.Function;
 import java.util.Set;
 
 /**
@@ -108,7 +107,7 @@ public static CelProtoExtensions protos() {
    *     options object used to configure the compilation/runtime environments.
    */
   public static CelMathExtensions math(CelOptions celOptions) {
-    return new CelMathExtensions(celOptions, Integer.MAX_VALUE);
+    return CelMathExtensions.library(celOptions).latest();
   }
 
   /**
@@ -117,7 +116,7 @@ public static CelMathExtensions math(CelOptions celOptions) {
    * <p>Refer to README.md for functions available in each version.
    */
   public static CelMathExtensions math(CelOptions celOptions, int version) {
-    return new CelMathExtensions(celOptions, version);
+    return CelMathExtensions.library(celOptions).version(version);
   }
 
   /**
@@ -229,7 +228,7 @@ public static CelSetsExtensions sets(CelOptions celOptions, Set<SetsFunction> fu
    * CelListsExtensions.Function}.
    */
   public static CelListsExtensions lists() {
-    return new CelListsExtensions(Integer.MAX_VALUE);
+    return CelListsExtensions.library().latest();
   }
 
   /**
@@ -238,7 +237,7 @@ public static CelListsExtensions lists() {
    * <p>Refer to README.md for functions available in each version.
    */
   public static CelListsExtensions lists(int version) {
-    return new CelListsExtensions(version);
+    return CelListsExtensions.library().version(version);
   }
 
   /**
@@ -259,8 +258,8 @@ public static CelListsExtensions lists(CelListsExtensions.Function... functions)
    * <p>Refer to README.md for available functions.
    *
    * <p>This will include all functions denoted in {@link CelListsExtensions.Function}, including
-   * any future additions. To expose only a subset of functions, use {@link #lists(Function...)}
-   * instead.
+   * any future additions. To expose only a subset of functions, use {@link
+   * #lists(CelListsExtensions.Function...)} instead.
    */
   public static CelListsExtensions lists(Set<CelListsExtensions.Function> functions) {
     return new CelListsExtensions(functions);
@@ -301,5 +300,18 @@ public static ImmutableSet<String> getAllFunctionNames() {
         .collect(toImmutableSet());
   }
 
+  public static CelExtensionLibrary<? extends CelExtensionLibrary.FeatureSet> getExtensionLibrary(
+      String name, CelOptions options) {
+    switch (name) {
+      case "math":
+        return CelMathExtensions.library(options);
+      case "lists":
+        return CelListsExtensions.library();
+      // TODO: add support for remaining standard extensions
+      default:
+        throw new IllegalArgumentException("Unknown standard extension '" + name + "'");
+    }
+  }
+
   private CelExtensions() {}
 }