Add ListsExtension with flatten function (single level only)

Recursive implementation will be added in the future.

PiperOrigin-RevId: 666372924

Author: l46kok

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

@@ -16,6 +16,7 @@ java_library(
     deps = [
         ":bindings",
         ":encoders",
+        ":lists",
         ":math",
         ":protos",
         ":sets",
@@ -148,3 +149,18 @@ java_library(
         "@maven//:com_google_guava_guava",
     ],
 )
+
+java_library(
+    name = "lists",
+    srcs = ["CelListsExtensions.java"],
+    tags = [
+    ],
+    deps = [
+        "//checker:checker_builder",
+        "//common:compiler_common",
+        "//common/types",
+        "//compiler:compiler_builder",
+        "//runtime",
+        "@maven//:com_google_guava_guava",
+    ],
+)

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

@@ -20,6 +20,7 @@
 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;
 
 /**
@@ -34,6 +35,7 @@ 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();
 
   /**
    * Extended functions for string manipulation.
@@ -187,9 +189,9 @@ public static CelSetsExtensions sets() {
    *
    * <p>Refer to README.md for available functions.
    *
-   * <p>This will include all functions denoted in {@link CelSetExtensions.Function}, including any
-   * future additions. To expose only a subset of functions, use {@link
-   * #sets(CelSetExtensions.Function...)} instead.
+   * <p>This will include all functions denoted in {@link CelSetsExtensions.Function}, including any
+   * future additions. To expose only a subset of functions, use {@link #sets(CelOptions,
+   * CelSetsExtensions.Function...)} instead.
    */
   public static CelSetsExtensions sets(CelOptions celOptions) {
     return new CelSetsExtensions(celOptions);
@@ -219,6 +221,43 @@ public static CelSetsExtensions sets(
     return new CelSetsExtensions(celOptions, functions);
   }
 
+  /**
+   * Extended functions for List manipulation.
+   *
+   * <p>Refer to README.md for available functions.
+   *
+   * <p>This will include only the specific functions denoted by {@link
+   * CelListsExtensions.Function}.
+   */
+  public static CelListsExtensions lists() {
+    return LISTS_EXTENSIONS_ALL;
+  }
+
+  /**
+   * Extended functions for List manipulation.
+   *
+   * <p>Refer to README.md for available functions.
+   *
+   * <p>This will include only the specific functions denoted by {@link
+   * CelListsExtensions.Function}.
+   */
+  public static CelListsExtensions lists(CelListsExtensions.Function... functions) {
+    return lists(ImmutableSet.copyOf(functions));
+  }
+
+  /**
+   * Extended functions for List manipulation.
+   *
+   * <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.
+   */
+  public static CelListsExtensions lists(Set<CelListsExtensions.Function> functions) {
+    return new CelListsExtensions(functions);
+  }
+
   /**
    * Retrieves all function names used by every extension libraries.
    *
@@ -235,7 +274,9 @@ public static ImmutableSet<String> getAllFunctionNames() {
             stream(CelSetsExtensions.Function.values())
                 .map(CelSetsExtensions.Function::getFunction),
             stream(CelEncoderExtensions.Function.values())
-                .map(CelEncoderExtensions.Function::getFunction))
+                .map(CelEncoderExtensions.Function::getFunction),
+            stream(CelListsExtensions.Function.values())
+                .map(CelListsExtensions.Function::getFunction))
         .collect(toImmutableSet());
   }
 

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

@@ -0,0 +1,101 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.extensions;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import dev.cel.checker.CelCheckerBuilder;
+import dev.cel.common.CelFunctionDecl;
+import dev.cel.common.CelOverloadDecl;
+import dev.cel.common.types.ListType;
+import dev.cel.common.types.TypeParamType;
+import dev.cel.compiler.CelCompilerLibrary;
+import dev.cel.runtime.CelRuntime;
+import dev.cel.runtime.CelRuntime.CelFunctionBinding;
+import dev.cel.runtime.CelRuntimeBuilder;
+import dev.cel.runtime.CelRuntimeLibrary;
+import java.util.Collection;
+import java.util.Set;
+
+/** Internal implementation of CEL lists extensions. */
+final class CelListsExtensions implements CelCompilerLibrary, CelRuntimeLibrary {
+
+  private static final TypeParamType LIST_PARAM_TYPE = TypeParamType.create("T");
+
+  @SuppressWarnings({"unchecked"}) // Unchecked: Type-checker guarantees casting safety.
+  public enum Function {
+    FLATTEN(
+        CelFunctionDecl.newFunctionDeclaration(
+            "flatten",
+            CelOverloadDecl.newMemberOverload(
+                "list_flatten",
+                "Flattens a list by a single level",
+                ListType.create(LIST_PARAM_TYPE),
+                ListType.create(ListType.create(LIST_PARAM_TYPE)))),
+        // TODO: add list_flatten_list_int
+        CelRuntime.CelFunctionBinding.from(
+            "list_flatten", Collection.class, list -> flatten(list, 1))),
+    ;
+
+    private final CelFunctionDecl functionDecl;
+    private final ImmutableSet<CelFunctionBinding> functionBindings;
+
+    String getFunction() {
+      return functionDecl.name();
+    }
+
+    Function(CelFunctionDecl functionDecl, CelRuntime.CelFunctionBinding... functionBindings) {
+      this.functionDecl = functionDecl;
+      this.functionBindings = ImmutableSet.copyOf(functionBindings);
+    }
+  }
+
+  private final ImmutableSet<Function> functions;
+
+  CelListsExtensions() {
+    this.functions = ImmutableSet.copyOf(Function.values());
+  }
+
+  CelListsExtensions(Set<Function> functions) {
+    this.functions = ImmutableSet.copyOf(functions);
+  }
+
+  @Override
+  public void setCheckerOptions(CelCheckerBuilder checkerBuilder) {
+    functions.forEach(function -> checkerBuilder.addFunctionDeclarations(function.functionDecl));
+  }
+
+  @Override
+  public void setRuntimeOptions(CelRuntimeBuilder runtimeBuilder) {
+    functions.forEach(function -> runtimeBuilder.addFunctionBindings(function.functionBindings));
+  }
+
+  @SuppressWarnings("unchecked")
+  private static ImmutableList<Object> flatten(Collection<Object> list, int level) {
+    Preconditions.checkArgument(level == 1, "recursive flatten is not supported yet.");
+    ImmutableList.Builder<Object> builder = ImmutableList.builder();
+    for (Object element : list) {
+      if (element instanceof Collection) {
+        Collection<Object> listItem = (Collection<Object>) element;
+        builder.addAll(listItem);
+      } else {
+        builder.add(element);
+      }
+    }
+
+    return builder.build();
+  }
+}

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

@@ -406,3 +406,33 @@ sets.intersects([1], [1, 2]) // true
 sets.intersects([[1], [2, 3]], [[1, 2], [2, 3.0]]) // true
 ```
 
+## Lists
+
+Extended functions for list manipulation. As a general note, all indices are
+zero-based.
+
+### Flatten
+
+Flattens a list by one level. Support for flattening to a specified level
+will be provided in the future.
+
+Examples:
+
+```
+[].flatten() // []
+[1,[2,3],[4]].flatten() // [1, 2, 3, 4]
+[1,[2,[3,4]]].flatten() // [1, 2, [3, 4]]
+[1,2,[],[],[3,4]].flatten() // [1, 2, 3, 4]
+```
+
+Note that due to the current limitations of type-checker, a compilation error
+will occur if an already flat list is populated. For time being, you must wrap
+the list in dyn if you anticipate having to deal with a flat list:
+
+```
+[1,2,3].flatten() // error
+dyn([1,2,3]).flatten() // [1,2,3]
+```
+
+This will be addressed once we add the appropriate capabilities in the
+type-checker to handle type-reductions, or union types.

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

@@ -159,6 +159,7 @@ public void getAllFunctionNames() {
             "sets.equivalent",
             "sets.intersects",
             "base64.decode",
-            "base64.encode");
+            "base64.encode",
+            "flatten");
   }
 }

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

@@ -0,0 +1,62 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package dev.cel.extensions;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import com.google.testing.junit.testparameterinjector.TestParameters;
+import dev.cel.bundle.Cel;
+import dev.cel.bundle.CelFactory;
+import dev.cel.common.CelValidationException;
+import dev.cel.parser.CelStandardMacro;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(TestParameterInjector.class)
+public class CelListsExtensionsTest {
+  private static final Cel CEL =
+      CelFactory.standardCelBuilder()
+          .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
+          .addCompilerLibraries(CelExtensions.lists())
+          .addRuntimeLibraries(CelExtensions.lists())
+          .build();
+
+  @Test
+  @TestParameters("{expression: '[].flatten() == []'}")
+  @TestParameters("{expression: '[[1, 2]].flatten().exists(i, i == 1)'}")
+  @TestParameters("{expression: '[[], [[]], [[[]]]].flatten() == [[], [[]]]'}")
+  @TestParameters("{expression: '[1,[2,[3,4]]].flatten() == [1,2,[3,4]]'}")
+  @TestParameters("{expression: '[1,2,[],[],[3,4]].flatten() == [1,2,3,4]'}")
+  @TestParameters("{expression: '[1,[2,3],[[4,5]], [[[6,7]]]].flatten() == [1,2,3,[4,5],[[6,7]]]'}")
+  @TestParameters("{expression: 'dyn([1]).flatten() == [1]'}")
+  @TestParameters("{expression: 'dyn([{1: 2}]).flatten() == [{1: 2}]'}")
+  @TestParameters("{expression: 'dyn([1,2,3,4]).flatten() == [1,2,3,4]'}")
+  public void flattenSingleLevel_success(String expression) throws Exception {
+    boolean result = (boolean) CEL.createProgram(CEL.compile(expression).getAst()).eval();
+
+    assertThat(result).isTrue();
+  }
+
+  @Test
+  @TestParameters("{expression: '[1].flatten()'}")
+  @TestParameters("{expression: '[{1: 2}].flatten()'}")
+  @TestParameters("{expression: '[1,2,3,4].flatten()'}")
+  public void flattenSingleLevel_listIsSingleLevel_throws(String expression) {
+    // Note: Java lacks the capability of conditionally disabling type guards
+    // due to the lack of full-fledged dynamic dispatch.
+    assertThrows(CelValidationException.class, () -> CEL.compile(expression).getAst());
+  }
+}