Internal cleanup

PiperOrigin-RevId: 974574709

Author: maskri17

MODULE.bazel

@@ -27,7 +27,7 @@ bazel_dep(name = "rules_java", version = "9.7.0")
 bazel_dep(name = "rules_android", version = "0.7.3")
 bazel_dep(name = "rules_shell", version = "0.8.0")
 bazel_dep(name = "googleapis-java", version = "1.1.5")
-bazel_dep(name = "cel-spec", version = "0.25.2", repo_name = "cel_spec")
+bazel_dep(name = "cel-spec", version = "0.25.3", repo_name = "cel_spec")
 bazel_dep(name = "rules_go", version = "0.62.0")
 
 # Required by cel-spec to satisfy gazelle transitive dependency

conformance/src/test/java/dev/cel/conformance/BUILD.bazel

@@ -89,6 +89,7 @@ _ALL_TESTS = [
     "@cel_spec//tests/simple:testdata/fp_math.textproto",
     "@cel_spec//tests/simple:testdata/integer_math.textproto",
     "@cel_spec//tests/simple:testdata/lists.textproto",
+    "@cel_spec//tests/simple:testdata/lists_ext.textproto",
     "@cel_spec//tests/simple:testdata/logic.textproto",
     "@cel_spec//tests/simple:testdata/macros.textproto",
     "@cel_spec//tests/simple:testdata/macros2.textproto",

conformance/src/test/java/dev/cel/conformance/ConformanceTest.java

@@ -82,7 +82,8 @@ public final class ConformanceTest extends Statement {
           CelExtensions.bindings(),
           CelExtensions.comprehensions(),
           CelExtensions.encoders(OPTIONS),
-          CelExtensions.math(OPTIONS),
+          CelExtensions.lists(),
+          CelExtensions.math(),
           CelExtensions.protos(),
           CelExtensions.sets(OPTIONS),
           CelExtensions.strings(),
@@ -93,7 +94,8 @@ public final class ConformanceTest extends Statement {
       ImmutableList.of(
           CelExtensions.comprehensions(),
           CelExtensions.encoders(OPTIONS),
-          CelExtensions.math(OPTIONS),
+          CelExtensions.lists(),
+          CelExtensions.math(),
           CelExtensions.sets(OPTIONS),
           CelExtensions.strings(),
           CelOptionalLibrary.INSTANCE);

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

@@ -274,6 +274,7 @@ java_library(
         "//common/ast",
         "//common/internal:comparison_functions",
         "//common/types",
+        "//common/values:cel_byte_string",
         "//compiler:compiler_builder",
         "//extensions:extension_library",
         "//parser:macro",

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

@@ -34,6 +34,7 @@
 import dev.cel.common.types.ListType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.TypeParamType;
+import dev.cel.common.values.CelByteString;
 import dev.cel.compiler.CelCompilerLibrary;
 import dev.cel.parser.CelMacro;
 import dev.cel.parser.CelMacroExprFactory;
@@ -42,11 +43,14 @@
 import dev.cel.runtime.CelInternalRuntimeLibrary;
 import dev.cel.runtime.CelRuntimeBuilder;
 import dev.cel.runtime.RuntimeEquality;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
@@ -132,16 +136,21 @@ public enum Function {
         CelFunctionBinding.from("list_sort", Collection.class, CelListsExtensions::sort)),
     SORT_BY(
         CelFunctionDecl.newFunctionDeclaration(
-            "lists.@sortByAssociatedKeys",
-            CelOverloadDecl.newGlobalOverload(
+            "@sortByAssociatedKeys",
+            CelOverloadDecl.newMemberOverload(
                 "list_sortByAssociatedKeys",
-                "Sorts a list by a key value. Used by the 'sortBy' macro",
+                "Sorts a list by associated keys. Used by the 'sortBy' macro",
                 ListType.create(TypeParamType.create("T")),
-                ListType.create(TypeParamType.create("T")))),
+                ListType.create(TypeParamType.create("T")),
+                ListType.create(TypeParamType.create("U")))),
         CelFunctionBinding.from(
             "list_sortByAssociatedKeys",
-            Collection.class,
-            CelListsExtensions::sortByAssociatedKeys));
+            ImmutableList.of(Collection.class, Collection.class),
+            (args) -> {
+              Collection<Object> target = (Collection<Object>) args[0];
+              Collection<Object> keys = (Collection<Object>) args[1];
+              return CelListsExtensions.sortByAssociatedKeys(target, keys);
+            }));
 
     private final CelFunctionDecl functionDecl;
     private final ImmutableSet<CelFunctionBinding> functionBindings;
@@ -222,7 +231,7 @@ public ImmutableSet<CelFunctionDecl> functions() {
 
   @Override
   public ImmutableSet<CelMacro> macros() {
-    if (version >= 2) {
+    if (version >= 2 || (version == -1 && functions.contains(Function.SORT_BY))) {
       return ImmutableSet.of(
           CelMacro.newReceiverMacro("sortBy", 2, CelListsExtensions::sortByMacro));
     }
@@ -300,7 +309,10 @@ private static ImmutableList<Object> flatten(Collection<Object> list, long depth
   }
 
   public static ImmutableList<Long> genRange(long end) {
-    ImmutableList.Builder<Long> builder = ImmutableList.builder();
+    checkArgument(end >= 0, "lists.range: size must be non-negative, got %s", end);
+    checkArgument(end <= 1_000_000, "lists.range: size %s exceeds maximum allowed (1000000)", end);
+
+    ImmutableList.Builder<Long> builder = ImmutableList.builderWithExpectedSize((int) end);
     for (long i = 0; i < end; i++) {
       builder.add(i);
     }
@@ -359,6 +371,17 @@ private static List<Object> reverse(Collection<Object> list) {
   }
 
   private static ImmutableList<Object> sort(Collection<Object> objects) {
+    if (objects.isEmpty()) {
+      return ImmutableList.of();
+    }
+    for (Object element : objects) {
+      if (!isSupportedComparableType(element)) {
+        throw new IllegalArgumentException("List elements must be comparable");
+      }
+    }
+    if (objects.size() < 2) {
+      return ImmutableList.copyOf(objects);
+    }
     return ImmutableList.sortedCopyOf(new CelObjectComparator(), objects);
   }
 
@@ -369,12 +392,14 @@ private static class CelObjectComparator implements Comparator<Object> {
     @SuppressWarnings({"unchecked"})
     @Override
     public int compare(Object o1, Object o2) {
+      if (o1 == null || o2 == null) {
+        throw new IllegalArgumentException("List elements must be comparable");
+      }
       if (o1 instanceof Number && o2 instanceof Number) {
         return ComparisonFunctions.numericCompare((Number) o1, (Number) o2);
       }
-
-      if (!(o1 instanceof Comparable)) {
-        throw new IllegalArgumentException("List elements must be comparable");
+      if (isByteType(o1) && isByteType(o2)) {
+        return compareBytes(o1, o2);
       }
       if (o1.getClass() != o2.getClass()) {
         throw new IllegalArgumentException("List elements must have the same type");
@@ -383,6 +408,45 @@ public int compare(Object o1, Object o2) {
     }
   }
 
+  private static boolean isByteType(Object obj) {
+    return obj instanceof CelByteString || obj instanceof byte[];
+  }
+
+  private static int compareBytes(Object o1, Object o2) {
+    if (o1 instanceof CelByteString && o2 instanceof CelByteString) {
+      return CelByteString.unsignedLexicographicalComparator()
+          .compare((CelByteString) o1, (CelByteString) o2);
+    }
+
+    byte[] b1 = o1 instanceof CelByteString ? ((CelByteString) o1).toByteArray() : (byte[]) o1;
+    byte[] b2 = o2 instanceof CelByteString ? ((CelByteString) o2).toByteArray() : (byte[]) o2;
+
+    int minLength = Math.min(b1.length, b2.length);
+    for (int i = 0; i < minLength; i++) {
+      int result = Integer.compare(Byte.toUnsignedInt(b1[i]), Byte.toUnsignedInt(b2[i]));
+      if (result != 0) {
+        return result;
+      }
+    }
+    return Integer.compare(b1.length, b2.length);
+  }
+
+  private static boolean isSupportedComparableType(Object obj) {
+    if (obj == null) {
+      return false;
+    }
+    if (obj instanceof Number
+        || obj instanceof Boolean
+        || obj instanceof String
+        || obj instanceof CelByteString
+        || obj instanceof byte[]
+        || obj instanceof Duration
+        || obj instanceof Instant) {
+      return true;
+    }
+    return obj instanceof Comparable && !(obj instanceof Collection) && !(obj instanceof Map);
+  }
+
   private static Optional<CelExpr> sortByMacro(
       CelMacroExprFactory exprFactory, CelExpr target, ImmutableList<CelExpr> arguments) {
     checkNotNull(exprFactory);
@@ -400,56 +464,79 @@ private static Optional<CelExpr> sortByMacro(
     String varName = varIdent.ident().name();
     CelExpr sortKeyExpr = checkNotNull(arguments.get(1));
 
-    // Compute the key using the second argument of the `sortBy(e, key)` macro.
-    // Combine the key and the value in a two-element list
-    CelExpr step = exprFactory.newList(sortKeyExpr, varIdent);
-    // Wrap the pair in another list in order to be able to use the `list+list` operator
-    step = exprFactory.newList(step);
-    // Append the key-value pair to the i
-    step =
+    String sortByInputVar = "@__sortBy_input__";
+    CelExpr sortByInputIdent = exprFactory.newIdentifier(sortByInputVar);
+
+    // Map comprehension: target.map(varName, sortKeyExpr)
+    CelExpr mapStep =
         exprFactory.newGlobalCall(
             Operator.ADD.getFunction(),
             exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()),
-            step);
-    // Create an intermediate list and populate it with key-value pairs
-    step =
+            exprFactory.newList(sortKeyExpr));
+    CelExpr mapCompr =
         exprFactory.fold(
             varName,
-            target,
+            sortByInputIdent,
             exprFactory.getAccumulatorVarName(),
             exprFactory.newList(),
-            exprFactory.newBoolLiteral(true), // Include all elements
-            step,
+            exprFactory.newBoolLiteral(true),
+            mapStep,
             exprFactory.newIdentifier(exprFactory.getAccumulatorVarName()));
-    // Finally, sort the list of key-value pairs and map it to a list of values
-    step = exprFactory.newGlobalCall(Function.SORT_BY.getFunction(), step);
 
-    return Optional.of(step);
+    // Receiver call: sortByInputIdent.@sortByAssociatedKeys(mapCompr)
+    CelExpr callExpr =
+        exprFactory.newReceiverCall(
+            Function.SORT_BY.getFunction(),
+            sortByInputIdent,
+            mapCompr);
+
+    // cel.bind(sortByInputVar, target, callExpr)
+    CelExpr bindExpr =
+        exprFactory.fold(
+            "#unused",
+            exprFactory.newList(),
+            sortByInputVar,
+            target,
+            // Loop condition is false because this comprehension simulates a local variable assignment (`bind`), rather than a traditional iteration.
+            exprFactory.newBoolLiteral(false),
+            sortByInputIdent,
+            callExpr);
+
+    return Optional.of(bindExpr);
   }
 
-  @SuppressWarnings({"unchecked", "rawtypes"})
   private static ImmutableList<Object> sortByAssociatedKeys(
-      Collection<List<Object>> keyValuePairs) {
-    List<Object>[] array = keyValuePairs.toArray(new List[0]);
-    Arrays.sort(array, new CelObjectByKeyComparator(new CelObjectComparator()));
-    ImmutableList.Builder<Object> builder = ImmutableList.builderWithExpectedSize(array.length);
-    for (List<Object> pair : array) {
-      builder.add(pair.get(1));
+      Collection<Object> list, Collection<Object> keys) {
+    if (list.size() != keys.size()) {
+      throw new IllegalArgumentException(
+          String.format(
+              "@sortByAssociatedKeys() expected a list of the same size as the associated keys"
+                  + " list, but got %d and %d elements respectively.",
+              list.size(), keys.size()));
     }
-    return builder.build();
-  }
-
-  private static class CelObjectByKeyComparator implements Comparator<Object> {
-    private final CelObjectComparator keyComparator;
-
-    CelObjectByKeyComparator(CelObjectComparator keyComparator) {
-      this.keyComparator = keyComparator;
+    if (list.isEmpty()) {
+      return ImmutableList.of();
     }
-
-    @SuppressWarnings({"unchecked"})
-    @Override
-    public int compare(Object o1, Object o2) {
-      return keyComparator.compare(((List<Object>) o1).get(0), ((List<Object>) o2).get(0));
+    for (Object key : keys) {
+      if (!isSupportedComparableType(key)) {
+        throw new IllegalArgumentException("List elements must be comparable");
+      }
+    }
+    if (list.size() < 2) {
+      return ImmutableList.copyOf(list);
     }
+    Object[] listArray = list.toArray();
+    Object[] keysArray = keys.toArray();
+    Integer[] indices = new Integer[listArray.length];
+    for (int i = 0; i < indices.length; i++) {
+      indices[i] = i;
+    }
+    CelObjectComparator comparator = new CelObjectComparator();
+    Arrays.sort(indices, (i1, i2) -> comparator.compare(keysArray[i1], keysArray[i2]));
+    ImmutableList.Builder<Object> builder = ImmutableList.builderWithExpectedSize(indices.length);
+    for (int idx : indices) {
+      builder.add(listArray[idx]);
+    }
+    return builder.build();
   }
 }

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

@@ -185,7 +185,7 @@ public void getAllFunctionNames() {
             "distinct",
             "reverse",
             "sort",
-            "lists.@sortByAssociatedKeys",
+            "@sortByAssociatedKeys",
             "regex.replace",
             "regex.extract",
             "regex.extractAll",