Fix CEL Java sortBy macro expansion to preserve element type and avoid heterogeneous list literals.

PiperOrigin-RevId: 979138470

Author: laiyichin

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

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

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

@@ -18,6 +18,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
+import com.google.common.base.Ascii;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -31,9 +32,11 @@
 import dev.cel.common.Operator;
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.internal.ComparisonFunctions;
+import dev.cel.common.types.CelType;
 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;
@@ -54,6 +57,10 @@
 public final class CelListsExtensions
     implements CelCompilerLibrary, CelInternalRuntimeLibrary, CelExtensionLibrary.FeatureSet {
 
+  private static final CelObjectComparator OBJECT_COMPARATOR = new CelObjectComparator();
+  private static final String UNUSED_ITER_VAR = "#unused";
+  private static final String SORT_BY_INPUT_VAR = "@__sortBy_input__";
+
   /** Supported functions for Lists extension library. */
   @SuppressWarnings({"unchecked"}) // Unchecked: Type-checker guarantees casting safety.
   public enum Function {
@@ -131,17 +138,50 @@ public enum Function {
                 ListType.create(TypeParamType.create("T")))),
         CelFunctionBinding.from("list_sort", Collection.class, CelListsExtensions::sort)),
     SORT_BY(
-        CelFunctionDecl.newFunctionDeclaration(
-            "lists.@sortByAssociatedKeys",
-            CelOverloadDecl.newGlobalOverload(
-                "list_sortByAssociatedKeys",
-                "Sorts a list by a key value. Used by the 'sortBy' macro",
+        createSortByFunctionDecl(comparableSortKeyTypes()),
+        createSortByFunctionBindings(comparableSortKeyTypes()));
+
+    private static ImmutableList<CelType> comparableSortKeyTypes() {
+      return ImmutableList.of(
+          SimpleType.INT,
+          SimpleType.UINT,
+          SimpleType.DOUBLE,
+          SimpleType.BOOL,
+          SimpleType.STRING,
+          SimpleType.BYTES,
+          SimpleType.DURATION,
+          SimpleType.TIMESTAMP);
+    }
+
+    private static CelFunctionDecl createSortByFunctionDecl(ImmutableList<CelType> keyTypes) {
+      ImmutableList.Builder<CelOverloadDecl> overloads = ImmutableList.builder();
+      for (CelType type : keyTypes) {
+        String typeName = Ascii.toLowerCase(type.kind().name());
+        overloads.add(
+            CelOverloadDecl.newMemberOverload(
+                String.format("list_%s_sortByAssociatedKeys", typeName),
+                "Sorts a list by an associated list of keys. Used by the 'sortBy' macro",
                 ListType.create(TypeParamType.create("T")),
-                ListType.create(TypeParamType.create("T")))),
-        CelFunctionBinding.from(
-            "list_sortByAssociatedKeys",
-            Collection.class,
-            CelListsExtensions::sortByAssociatedKeys));
+                ListType.create(TypeParamType.create("T")),
+                ListType.create(type)));
+      }
+      return CelFunctionDecl.newFunctionDeclaration("@sortByAssociatedKeys", overloads.build());
+    }
+
+    private static CelFunctionBinding[] createSortByFunctionBindings(
+        ImmutableList<CelType> keyTypes) {
+      return keyTypes.stream()
+          .map(
+              type -> {
+                String typeName = Ascii.toLowerCase(type.kind().name());
+                return CelFunctionBinding.from(
+                    String.format("list_%s_sortByAssociatedKeys", typeName),
+                    Collection.class,
+                    Collection.class,
+                    CelListsExtensions::sortByAssociatedKeys);
+              })
+          .toArray(CelFunctionBinding[]::new);
+    }
 
     private final CelFunctionDecl functionDecl;
     private final ImmutableSet<CelFunctionBinding> functionBindings;
@@ -359,7 +399,15 @@ private static List<Object> reverse(Collection<Object> list) {
   }
 
   private static ImmutableList<Object> sort(Collection<Object> objects) {
-    return ImmutableList.sortedCopyOf(new CelObjectComparator(), objects);
+    if (objects.isEmpty()) {
+      return ImmutableList.of();
+    }
+    if (objects.size() == 1) {
+      Object single = objects.iterator().next();
+      OBJECT_COMPARATOR.compare(single, single);
+      return ImmutableList.of(single);
+    }
+    return ImmutableList.sortedCopyOf(OBJECT_COMPARATOR, objects);
   }
 
   private static class CelObjectComparator implements Comparator<Object> {
@@ -372,6 +420,10 @@ public int compare(Object o1, Object o2) {
       if (o1 instanceof Number && o2 instanceof Number) {
         return ComparisonFunctions.numericCompare((Number) o1, (Number) o2);
       }
+      if (o1 instanceof CelByteString && o2 instanceof CelByteString) {
+        return CelByteString.unsignedLexicographicalComparator()
+            .compare((CelByteString) o1, (CelByteString) o2);
+      }
 
       if (!(o1 instanceof Comparable)) {
         throw new IllegalArgumentException("List elements must be comparable");
@@ -383,6 +435,34 @@ public int compare(Object o1, Object o2) {
     }
   }
 
+  /**
+   * Expands the {@code list.sortBy(var, expr)} receiver macro into a binding expression that sorts
+   * the target list using keys evaluated by mapping {@code expr} over each element.
+   *
+   * <p>For example, given:
+   *
+   * <pre>{@code
+   * myList.sortBy(item, -item.field)
+   * }</pre>
+   *
+   * <p>The macro expands into:
+   *
+   * <pre>{@code
+   * cel.bind(@__sortBy_input__, myList,
+   *     @__sortBy_input__.@sortByAssociatedKeys(
+   *         @__sortBy_input__.map(item, -item.field)
+   *     )
+   * )
+   * }</pre>
+   *
+   * <p>Where:
+   *
+   * <ul>
+   *   <li>{@code @__sortBy_input__.map(item, -item.field)} evaluates the sort key for each element.
+   *   <li>{@code @sortByAssociatedKeys} stably sorts the input list elements based on their
+   *       corresponding sort keys.
+   * </ul>
+   */
   private static Optional<CelExpr> sortByMacro(
       CelMacroExprFactory exprFactory, CelExpr target, ImmutableList<CelExpr> arguments) {
     checkNotNull(exprFactory);
@@ -400,56 +480,86 @@ 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 =
+    // Build map comprehension: @__sortBy_input__.map(varName, sortKeyExpr)
+    CelExpr targetIdent = exprFactory.newIdentifier(SORT_BY_INPUT_VAR);
+    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,
+            targetIdent,
             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);
+    // Build call: @__sortBy_input__.@sortByAssociatedKeys(mapCompr)
+    CelExpr callExpr =
+        exprFactory.newReceiverCall(
+            Function.SORT_BY.getFunction(), exprFactory.newIdentifier(SORT_BY_INPUT_VAR), mapCompr);
+
+    // Build bind: cel.bind(@__sortBy_input__, target, callExpr)
+    CelExpr bindExpr =
+        exprFactory.fold(
+            UNUSED_ITER_VAR,
+            exprFactory.newList(),
+            SORT_BY_INPUT_VAR,
+            target,
+            exprFactory.newBoolLiteral(false),
+            exprFactory.newIdentifier(SORT_BY_INPUT_VAR),
+            callExpr);
+
+    return Optional.of(bindExpr);
   }
 
-  @SuppressWarnings({"unchecked", "rawtypes"})
+  /**
+   * Sorts elements of {@code list} based on the natural order of corresponding elements in {@code
+   * keys}.
+   *
+   * <p>Both {@code list} and {@code keys} must have the exact same size. The sorting is stable
+   * (i.e., preserves the relative order of elements with equal keys).
+   *
+   * @param list The input list to sort
+   * @param keys The associated keys evaluated for each element in {@code list}
+   * @return A new {@link ImmutableList} containing the elements of {@code list} sorted by {@code
+   *     keys}
+   */
   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) {
+    checkArgument(
+        list.size() == keys.size(),
+        "@sortByAssociatedKeys() expected a list of the same size as the associated keys"
+            + " list, but got %s in list and %s in keys",
+        list.size(),
+        keys.size());
+
+    int listSize = list.size();
+    if (listSize == 0) {
+      return ImmutableList.of();
     }
-    return builder.build();
-  }
 
-  private static class CelObjectByKeyComparator implements Comparator<Object> {
-    private final CelObjectComparator keyComparator;
+    Object[] listArray = list.toArray();
+    Object[] keysArray = keys.toArray();
+    if (listSize == 1) {
+      OBJECT_COMPARATOR.compare(keysArray[0], keysArray[0]);
+      return ImmutableList.of(listArray[0]);
+    }
 
-    CelObjectByKeyComparator(CelObjectComparator keyComparator) {
-      this.keyComparator = keyComparator;
+    Integer[] indices = new Integer[listSize];
+    for (int i = 0; i < listSize; i++) {
+      indices[i] = i;
     }
 
-    @SuppressWarnings({"unchecked"})
-    @Override
-    public int compare(Object o1, Object o2) {
-      return keyComparator.compare(((List<Object>) o1).get(0), ((List<Object>) o2).get(0));
+    Arrays.sort(indices, (i1, i2) -> OBJECT_COMPARATOR.compare(keysArray[i1], keysArray[i2]));
+
+    ImmutableList.Builder<Object> builder = ImmutableList.builderWithExpectedSize(listSize);
+    for (int index : indices) {
+      builder.add(listArray[index]);
     }
+    return builder.build();
   }
 }

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

@@ -42,12 +42,14 @@ java_library(
         "//parser:unparser",
         "//runtime",
         "//runtime:function_binding",
-        "//runtime:interpreter_util",
         "//runtime:lite_runtime",
         "//runtime:lite_runtime_factory",
         "//runtime:partial_vars",
         "//runtime:unknown_attributes",
         "//testing:cel_runtime_flavor",
+        "//validator",
+        "//validator:validator_builder",
+        "//validator/validators:homogeneous_literal",
         "@cel_spec//proto/cel/expr/conformance/proto2:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",
         "@cel_spec//proto/cel/expr/conformance/test:simple_java_proto",

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",