Provide an overload to accept a depth level in flatten function

PiperOrigin-RevId: 669346696

Author: l46kok

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

@@ -21,6 +21,7 @@
 import dev.cel.common.CelFunctionDecl;
 import dev.cel.common.CelOverloadDecl;
 import dev.cel.common.types.ListType;
+import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.TypeParamType;
 import dev.cel.compiler.CelCompilerLibrary;
 import dev.cel.runtime.CelRuntime;
@@ -44,11 +45,18 @@ public enum Function {
                 "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
+                ListType.create(ListType.create(LIST_PARAM_TYPE))),
+            CelOverloadDecl.newMemberOverload(
+                "list_flatten_list_int",
+                "Flattens a list to the specified level. A negative depth value flattens the list"
+                    + " recursively to its deepest level.",
+                ListType.create(SimpleType.DYN),
+                ListType.create(SimpleType.DYN),
+                SimpleType.INT)),
+        CelRuntime.CelFunctionBinding.from(
+            "list_flatten", Collection.class, list -> flatten(list, 1)),
         CelRuntime.CelFunctionBinding.from(
-            "list_flatten", Collection.class, list -> flatten(list, 1))),
-    ;
+            "list_flatten_list_int", Collection.class, Long.class, CelListsExtensions::flatten));
 
     private final CelFunctionDecl functionDecl;
     private final ImmutableSet<CelFunctionBinding> functionBindings;
@@ -84,15 +92,15 @@ public void setRuntimeOptions(CelRuntimeBuilder runtimeBuilder) {
   }
 
   @SuppressWarnings("unchecked")
-  private static ImmutableList<Object> flatten(Collection<Object> list, int level) {
-    Preconditions.checkArgument(level == 1, "recursive flatten is not supported yet.");
+  private static ImmutableList<Object> flatten(Collection<Object> list, long depth) {
+    Preconditions.checkArgument(depth >= 0, "Level must be non-negative");
     ImmutableList.Builder<Object> builder = ImmutableList.builder();
     for (Object element : list) {
-      if (element instanceof Collection) {
-        Collection<Object> listItem = (Collection<Object>) element;
-        builder.addAll(listItem);
-      } else {
+      if (!(element instanceof Collection) || depth == 0) {
         builder.add(element);
+      } else {
+        Collection<Object> listItem = (Collection<Object>) element;
+        builder.addAll(flatten(listItem, depth - 1));
       }
     }
 

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

@@ -413,24 +413,38 @@ zero-based.
 
 ### Flatten
 
-Flattens a list by one level. Support for flattening to a specified level
-will be provided in the future.
+Flattens a list by one level, or to the specified level. Providing a negative level will error.
 
 Examples:
 
 ```
+// Single-level flatten:
+
 [].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]
+
+// Recursive flatten
+[1,[2,[3,[4]]]].flatten(2) // return [1, 2, 3, [4]]
+[1,[2,[3,[4]]]].flatten(3) // return [1, 2, 3, 4]
+
+// Error
+[1,[2,[3,[4]]]].flatten(-1)
 ```
 
 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:
+will occur if an already flat list is populated to the argument-less flatten
+function.
+
+For time being, you must explicitly provide 1 as the depth level, or wrap the
+list in dyn if you anticipate having to deal with a flat list:
 
 ```
 [1,2,3].flatten() // error
+
+// But the following will work:
+[1,2,3].flatten(1) // [1,2,3]
 dyn([1,2,3]).flatten() // [1,2,3]
 ```
 

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

@@ -22,6 +22,7 @@
 import dev.cel.bundle.CelFactory;
 import dev.cel.common.CelValidationException;
 import dev.cel.parser.CelStandardMacro;
+import dev.cel.runtime.CelEvaluationException;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -50,6 +51,35 @@ public void flattenSingleLevel_success(String expression) throws Exception {
     assertThat(result).isTrue();
   }
 
+  @Test
+  @TestParameters("{expression: '[1,2,3,4].flatten(1) == [1,2,3,4]'}")
+  @TestParameters("{expression: '[1,[2,[3,[4]]]].flatten(0) == [1,[2,[3,[4]]]]'}")
+  @TestParameters("{expression: '[1,[2,[3,[4]]]].flatten(2) == [1,2,3,[4]]'}")
+  @TestParameters("{expression: '[1,[2,[3,4]]].flatten(2) == [1,2,3,4]'}")
+  @TestParameters("{expression: '[[], [[]], [[[]]]].flatten(2) == [[]]'}")
+  @TestParameters("{expression: '[[], [[]], [[[]]]].flatten(3) == []'}")
+  @TestParameters("{expression: '[[], [[]], [[[]]]].flatten(4) == []'}")
+  // The overload with the depth accepts and returns a List(dyn), so the following is permitted.
+  @TestParameters("{expression: '[1].flatten(1) == [1]'}")
+  public void flatten_withDepthValue_success(String expression) throws Exception {
+    boolean result = (boolean) CEL.createProgram(CEL.compile(expression).getAst()).eval();
+
+    assertThat(result).isTrue();
+  }
+
+  @Test
+  public void flatten_negativeDepth_throws() {
+    CelEvaluationException e =
+        assertThrows(
+            CelEvaluationException.class,
+            () -> CEL.createProgram(CEL.compile("[1,2,3,4].flatten(-1)").getAst()).eval());
+
+    assertThat(e)
+        .hasMessageThat()
+        .contains("evaluation error: Function 'list_flatten_list_int' failed");
+    assertThat(e).hasCauseThat().hasMessageThat().isEqualTo("Level must be non-negative");
+  }
+
   @Test
   @TestParameters("{expression: '[1].flatten()'}")
   @TestParameters("{expression: '[{1: 2}].flatten()'}")