Introduce 'list' extension functions: 'slice'

Add versions to the 'lists' extension to gradually expose new functions.

PiperOrigin-RevId: 780339619

Author: dmitriplotnikov

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

@@ -27,6 +27,9 @@ java_library(
         ":sets_function",
         ":strings",
         "//common:options",
+        "//common/internal:env_visitor",
+        "//compiler",
+        "//compiler:compiler_builder",
         "@maven//:com_google_guava_guava",
     ],
 )

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

@@ -20,6 +20,10 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import dev.cel.common.CelOptions;
+import dev.cel.common.internal.EnvVisitable;
+import dev.cel.compiler.CelCompiler;
+import dev.cel.compiler.CelCompilerFactory;
+import dev.cel.compiler.CelCompilerLibrary;
 import dev.cel.extensions.CelListsExtensions.Function;
 import java.util.Set;
 
@@ -35,7 +39,6 @@ public final class CelExtensions {
   private static final CelProtoExtensions PROTO_EXTENSIONS = new CelProtoExtensions();
   private static final CelBindingsExtensions BINDINGS_EXTENSIONS = new CelBindingsExtensions();
   private static final CelEncoderExtensions ENCODER_EXTENSIONS = new CelEncoderExtensions();
-  private static final CelListsExtensions LISTS_EXTENSIONS_ALL = new CelListsExtensions();
   private static final CelRegexExtensions REGEX_EXTENSIONS = new CelRegexExtensions();
 
   /**
@@ -230,7 +233,16 @@ public static CelSetsExtensions sets(CelOptions celOptions, Set<SetsFunction> fu
    * CelListsExtensions.Function}.
    */
   public static CelListsExtensions lists() {
-    return LISTS_EXTENSIONS_ALL;
+    return new CelListsExtensions(Integer.MAX_VALUE);
+  }
+
+  /**
+   * Extended functions for List manipulation.
+   *
+   * <p>Refer to README.md for functions available in each version.
+   */
+  public static CelListsExtensions lists(int version) {
+    return new CelListsExtensions(version);
   }
 
   /**
@@ -293,5 +305,16 @@ public static ImmutableSet<String> getAllFunctionNames() {
         .collect(toImmutableSet());
   }
 
+  public static ImmutableSet<String> getFunctionNames(CelCompilerLibrary library) {
+    ImmutableSet.Builder<String> builder = ImmutableSet.builder();
+    CelCompiler compiler =
+        CelCompilerFactory.standardCelCompilerBuilder()
+            .setStandardEnvironmentEnabled(false)
+            .addLibraries(library)
+            .build();
+    ((EnvVisitable) compiler).accept((name, decls) -> builder.add(name));
+    return builder.build();
+  }
+
   private CelExtensions() {}
 }

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

@@ -28,6 +28,8 @@
 import dev.cel.runtime.CelRuntimeBuilder;
 import dev.cel.runtime.CelRuntimeLibrary;
 import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 
 /** Internal implementation of CEL lists extensions. */
@@ -37,6 +39,25 @@ final class CelListsExtensions implements CelCompilerLibrary, CelRuntimeLibrary
 
   @SuppressWarnings({"unchecked"}) // Unchecked: Type-checker guarantees casting safety.
   public enum Function {
+    SLICE(
+        CelFunctionDecl.newFunctionDeclaration(
+            "slice",
+            CelOverloadDecl.newMemberOverload(
+                "list_slice",
+                "Returns a new sub-list using the indexes provided",
+                ListType.create(LIST_PARAM_TYPE),
+                ListType.create(LIST_PARAM_TYPE),
+                SimpleType.INT,
+                SimpleType.INT)),
+        CelFunctionBinding.from(
+            "list_slice",
+            ImmutableList.of(Collection.class, Long.class, Long.class),
+            (args) -> {
+              Collection<Object> target = (Collection<Object>) args[0];
+              long from = (Long) args[1];
+              long to = (Long) args[2];
+              return CelListsExtensions.slice(target, from, to);
+            })),
     FLATTEN(
         CelFunctionDecl.newFunctionDeclaration(
             "flatten",
@@ -65,6 +86,31 @@ public enum Function {
                 SimpleType.INT)),
         CelFunctionBinding.from("lists_range", Long.class, CelListsExtensions::genRange));
 
+    private static final ImmutableSet<Function> VERSION_0 = ImmutableSet.of(SLICE);
+
+    private static final ImmutableSet<Function> VERSION_1 =
+        ImmutableSet.<Function>builder().addAll(VERSION_0).add(FLATTEN).build();
+
+    private static final ImmutableSet<Function> VERSION_2 =
+        ImmutableSet.<Function>builder().addAll(VERSION_1).add(RANGE).build();
+
+    private static final ImmutableSet<Function> VERSION_LATEST = VERSION_2;
+
+    private static ImmutableSet<Function> byVersion(int version) {
+      switch (version) {
+        case 0:
+          return Function.VERSION_0;
+        case 1:
+          return Function.VERSION_1;
+        case 2:
+          return Function.VERSION_2;
+        case Integer.MAX_VALUE:
+          return Function.VERSION_LATEST;
+        default:
+          throw new IllegalArgumentException("Unsupported 'lists' extension version " + version);
+      }
+    }
+
     private final CelFunctionDecl functionDecl;
     private final ImmutableSet<CelFunctionBinding> functionBindings;
 
@@ -80,8 +126,8 @@ String getFunction() {
 
   private final ImmutableSet<Function> functions;
 
-  CelListsExtensions() {
-    this.functions = ImmutableSet.copyOf(Function.values());
+  CelListsExtensions(int version) {
+    this(Function.byVersion(version));
   }
 
   CelListsExtensions(Set<Function> functions) {
@@ -98,6 +144,32 @@ public void setRuntimeOptions(CelRuntimeBuilder runtimeBuilder) {
     functions.forEach(function -> runtimeBuilder.addFunctionBindings(function.functionBindings));
   }
 
+  private static ImmutableList<Object> slice(Collection<Object> list, long from, long to) {
+    Preconditions.checkArgument(from >= 0 && to >= 0, "Negative indexes not supported");
+    Preconditions.checkArgument(to >= from, "Start index must be less than or equal to end index");
+    Preconditions.checkArgument(to <= list.size(), "List is length %s", list.size());
+    if (list instanceof List) {
+      List<Object> subList = ((List<Object>) list).subList((int) from, (int) to);
+      if (subList instanceof ImmutableList) {
+        return (ImmutableList<Object>) subList;
+      }
+      return ImmutableList.copyOf(subList);
+    } else {
+      ImmutableList.Builder<Object> builder = ImmutableList.builder();
+      long index = 0;
+      for (Iterator<Object> iterator = list.iterator(); iterator.hasNext(); index++) {
+        Object element = iterator.next();
+        if (index >= to) {
+          break;
+        }
+        if (index >= from) {
+          builder.add(element);
+        }
+      }
+      return builder.build();
+    }
+  }
+
   @SuppressWarnings("unchecked")
   private static ImmutableList<Object> flatten(Collection<Object> list, long depth) {
     Preconditions.checkArgument(depth >= 0, "Level must be non-negative");

extensions/src/main/java/dev/cel/extensions/README.md

@@ -666,8 +666,22 @@ sets.intersects([[1], [2, 3]], [[1, 2], [2, 3.0]]) // true
 Extended functions for list manipulation. As a general note, all indices are
 zero-based.
 
+### Slice
+
+Returns a new sub-list using the indexes provided. The `from` index is
+inclusive, the `to` index is exclusive.
+
+    <list>.slice(<int>, <int>) -> <list>
+
+Examples:
+
+    [1,2,3,4].slice(1, 3) // return [2, 3]
+    [1,2,3,4].slice(2, 4) // return [3, 4]
+
 ### Flatten
 
+Introduced at version: 1
+
 Flattens a list by one level, or to the specified level. Providing a negative level will error.
 
 Examples:
@@ -708,6 +722,8 @@ type-checker to handle type-reductions, or union types.
 
 ### Range
 
+Introduced at version: 2
+
 Given integer size n returns a list of integers from 0 to n-1. If size <= 0
 then return empty list.
 

extensions/src/test/java/dev/cel/extensions/CelListsExtensionsTest.java

@@ -35,6 +35,48 @@ public class CelListsExtensionsTest {
           .addRuntimeLibraries(CelExtensions.lists())
           .build();
 
+  @Test
+  public void functionList_byVersion() {
+    assertThat(CelExtensions.getFunctionNames(CelExtensions.lists(0))).containsExactly("slice");
+    assertThat(CelExtensions.getFunctionNames(CelExtensions.lists(1)))
+        .containsExactly("slice", "flatten");
+    assertThat(CelExtensions.getFunctionNames(CelExtensions.lists(2)))
+        .containsExactly("slice", "flatten", "lists.range");
+  }
+
+  @Test
+  @TestParameters("{expression: '[1,2,3,4].slice(0, 4) == [1,2,3,4]'}")
+  @TestParameters("{expression: '[1,2,3,4].slice(0, 0) == []'}")
+  @TestParameters("{expression: '[1,2,3,4].slice(1, 1) == []'}")
+  @TestParameters("{expression: '[1,2,3,4].slice(4, 4) == []'}")
+  @TestParameters("{expression: '[1,2,3,4].slice(1, 3) == [2, 3]'}")
+  public void slice_success(String expression) throws Exception {
+    boolean result = (boolean) CEL.createProgram(CEL.compile(expression).getAst()).eval();
+
+    assertThat(result).isTrue();
+  }
+
+  @Test
+  @TestParameters(
+      "{expression: '[1,2,3,4].slice(3, 0)', "
+          + "expectedError: 'Start index must be less than or equal to end index'}")
+  @TestParameters("{expression: '[1,2,3,4].slice(0, 10)', expectedError: 'List is length 4'}")
+  @TestParameters(
+      "{expression: '[1,2,3,4].slice(-5, 10)', "
+          + "expectedError: 'Negative indexes not supported'}")
+  @TestParameters(
+      "{expression: '[1,2,3,4].slice(-5, -3)', "
+          + "expectedError: 'Negative indexes not supported'}")
+  public void slice_throws(String expression, String expectedError) throws Exception {
+    assertThat(
+            assertThrows(
+                CelEvaluationException.class,
+                () -> CEL.createProgram(CEL.compile(expression).getAst()).eval()))
+        .hasCauseThat()
+        .hasMessageThat()
+        .contains(expectedError);
+  }
+
   @Test
   @TestParameters("{expression: '[].flatten() == []'}")
   @TestParameters("{expression: '[[1, 2]].flatten().exists(i, i == 1)'}")