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

PiperOrigin-RevId: 975736930

Author: laiyichin

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

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

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

@@ -132,15 +132,17 @@ 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 an associated list of 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,
+            Collection.class,
             CelListsExtensions::sortByAssociatedKeys));
 
     private final CelFunctionDecl functionDecl;
@@ -358,8 +360,18 @@ private static List<Object> reverse(Collection<Object> list) {
     }
   }
 
+  private static final CelObjectComparator OBJECT_COMPARATOR = new CelObjectComparator();
+
   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.copyOf(objects);
+    }
+    return ImmutableList.sortedCopyOf(OBJECT_COMPARATOR, objects);
   }
 
   private static class CelObjectComparator implements Comparator<Object> {
@@ -383,6 +395,37 @@ public int compare(Object o1, Object o2) {
     }
   }
 
+  private static final String UNUSED_ITER_VAR = "#unused";
+  private static final String SORT_BY_INPUT_VAR = "@__sortBy_input__";
+
+  /**
+   * 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 +443,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.copyOf(list);
+    }
 
-    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",

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

@@ -22,6 +22,7 @@
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import com.google.testing.junit.testparameterinjector.TestParameters;
 import dev.cel.bundle.Cel;
+import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelContainer;
 import dev.cel.common.CelValidationException;
 import dev.cel.common.CelValidationResult;
@@ -30,6 +31,9 @@
 import dev.cel.parser.CelStandardMacro;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.testing.CelRuntimeFlavor;
+import dev.cel.validator.CelValidator;
+import dev.cel.validator.CelValidatorFactory;
+import dev.cel.validator.validators.HomogeneousLiteralValidator;
 import org.junit.Assume;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -64,7 +68,7 @@ public void functionList_byVersion() {
             "distinct",
             "reverse",
             "sort",
-            "lists.@sortByAssociatedKeys");
+            "@sortByAssociatedKeys");
   }
 
   @Test
@@ -257,6 +261,9 @@ public void sort_success_heterogeneousNumbers(String expression, String expected
   @TestParameters(
       "{expression: '[SimpleTest{name: \"a\"}, SimpleTest{name: \"b\"}].sort()', "
           + "expectedError: 'List elements must be comparable'}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"a\"}].sort()', "
+          + "expectedError: 'List elements must be comparable'}")
   public void sort_throws(String expression, String expectedError) throws Exception {
     assertThat(assertThrows(CelEvaluationException.class, () -> eval(cel, expression)))
         .hasCauseThat()
@@ -283,6 +290,11 @@ public void sort_throws(String expression, String expectedError) throws Exceptio
           + "expected: '[SimpleTest{name: \"bar\"},"
           + " SimpleTest{name: \"baz\"},"
           + " SimpleTest{name: \"foo\"}]'}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"baz\"},"
+          + " SimpleTest{name: \"foo\"},"
+          + " SimpleTest{name: \"bar\"}].sortBy(e, e.name)[0].name', "
+          + "expected: '\"bar\"'}")
   public void sortBy_success(String expression, String expected) throws Exception {
     Object result = eval(cel, expression);
 
@@ -311,6 +323,9 @@ public void sortBy_throws_validationException(String expression, String expected
   @TestParameters(
       "{expression: '[SimpleTest{name: \"a\"}, SimpleTest{name: \"b\"}].sortBy(e, e)', "
           + "expectedError: 'List elements must be comparable'}")
+  @TestParameters(
+      "{expression: '[SimpleTest{name: \"a\"}].sortBy(e, e)', "
+          + "expectedError: 'List elements must be comparable'}")
   public void sortBy_throws_evaluationException(String expression, String expectedError)
       throws Exception {
     assertThat(assertThrows(CelEvaluationException.class, () -> eval(cel, expression)))
@@ -319,5 +334,21 @@ public void sortBy_throws_evaluationException(String expression, String expected
         .contains(expectedError);
   }
 
-
+  @Test
+  public void sortBy_withHomogeneousLiteralValidator_success() throws Exception {
+    CelValidator validator =
+        CelValidatorFactory.standardCelValidatorBuilder(cel)
+            .addAstValidators(HomogeneousLiteralValidator.newInstance())
+            .build();
+
+    CelAbstractSyntaxTree ast =
+        cel.compile(
+                "[SimpleTest{name: 'baz'}, SimpleTest{name: 'foo'}, SimpleTest{name: 'bar'}]"
+                    + ".sortBy(e, e.name)[0].name")
+            .getAst();
+    CelValidationResult result = validator.validate(ast);
+
+    assertThat(result.hasError()).isFalse();
+    assertThat(cel.createProgram(ast).eval()).isEqualTo("bar");
+  }
 }