Porting range() to Java Stack

PiperOrigin-RevId: 753777198

Author: CEL Dev Team

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

@@ -54,7 +54,16 @@ public enum Function {
                 SimpleType.INT)),
         CelFunctionBinding.from("list_flatten", Collection.class, list -> flatten(list, 1)),
         CelFunctionBinding.from(
-            "list_flatten_list_int", Collection.class, Long.class, CelListsExtensions::flatten));
+            "list_flatten_list_int", Collection.class, Long.class, CelListsExtensions::flatten)),
+    RANGE(
+        CelFunctionDecl.newFunctionDeclaration(
+            "lists.range",
+            CelOverloadDecl.newGlobalOverload(
+                "lists_range",
+                "Returns a list of integers from 0 to n-1.",
+                ListType.create(SimpleType.INT),
+                SimpleType.INT)),
+        CelFunctionBinding.from("lists_range", Long.class, CelListsExtensions::genRange));
 
     private final CelFunctionDecl functionDecl;
     private final ImmutableSet<CelFunctionBinding> functionBindings;
@@ -104,4 +113,12 @@ private static ImmutableList<Object> flatten(Collection<Object> list, long depth
 
     return builder.build();
   }
+
+  public static ImmutableList<Long> genRange(long end) {
+    ImmutableList.Builder<Long> builder = ImmutableList.builder();
+    for (long i = 0; i < end; i++) {
+      builder.add(i);
+    }
+    return builder.build();
+  }
 }

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

@@ -688,3 +688,19 @@ 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.
+
+### Range
+
+Given integer size n returns a list of integers from 0 to n-1. If size <= 0
+then return empty list.
+
+```
+lists.range(int) -> list(int)
+```
+
+Examples:
+
+```
+lists.range(5) -> [0, 1, 2, 3, 4]
+lists.range(0) -> []
+```

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

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

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

@@ -89,4 +89,14 @@ public void flattenSingleLevel_listIsSingleLevel_throws(String expression) {
     // due to the lack of full-fledged dynamic dispatch.
     assertThrows(CelValidationException.class, () -> CEL.compile(expression).getAst());
   }
+
+  @Test
+  @TestParameters("{expression: 'lists.range(9) == [0,1,2,3,4,5,6,7,8]'}")
+  @TestParameters("{expression: 'lists.range(0) == []'}")
+  @TestParameters("{expression: 'lists.range(-1) == []'}")
+  public void range_success(String expression) throws Exception {
+    boolean result = (boolean) CEL.createProgram(CEL.compile(expression).getAst()).eval();
+
+    assertThat(result).isTrue();
+  }
 }