Add key and reverse parameters to builtin sorted function

Implement:bazelbuild/starlark#20 (comment)

Closes #8881.

PiperOrigin-RevId: 258401453

src/main/java/com/google/devtools/build/lib/syntax/MethodLibrary.java

@@ -36,6 +36,8 @@
 import com.google.devtools.build.lib.syntax.SkylarkList.Tuple;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -102,6 +104,22 @@ private static Object findExtreme(
     }
   }
 
+  private static Object evalKeyFunc(
+      Object obj, Object key, Location loc, Environment env, FuncallExpression ast)
+      throws EvalException, InterruptedException {
+    checkValidKeyFunc(key, loc);
+    return ((StarlarkFunction) key)
+        .call(Collections.singletonList(obj), Collections.emptyMap(), ast, env);
+  }
+
+  private static void checkValidKeyFunc(Object key, Location loc) throws EvalException {
+    if (key instanceof StarlarkFunction) {
+      return;
+    }
+    throw new EvalException(
+        loc, Printer.format("%r object is not callable", EvalUtils.getDataTypeName(key)));
+  }
+
   @SkylarkCallable(
       name = "all",
       doc =
@@ -170,17 +188,77 @@ private static boolean hasElementWithBooleanValue(
             type = Object.class,
             doc = "This collection.",
             // TODO(cparsons): This parameter should be positional-only.
-            legacyNamed = true)
+            legacyNamed = true),
+        @Param(
+            name = "key",
+            doc = "An optional function applied to each element before comparison.",
+            named = true,
+            defaultValue = "None",
+            positional = false,
+            noneable = true),
+        @Param(
+            name = "reverse",
+            type = Boolean.class,
+            doc = "Return results in descending order.",
+            named = true,
+            defaultValue = "False",
+            positional = false,
+            noneable = true)
       },
       useLocation = true,
       useEnvironment = true)
-  public MutableList<?> sorted(Object self, Location loc, Environment env) throws EvalException {
-    try {
-      return MutableList.copyOf(
-          env, EvalUtils.SKYLARK_COMPARATOR.sortedCopy(EvalUtils.toCollection(self, loc, env)));
-    } catch (EvalUtils.ComparisonException e) {
-      throw new EvalException(loc, e);
+  public MutableList<?> sorted(
+      Object self, final Object key, Boolean reverse, final Location loc, final Environment env)
+      throws EvalException, InterruptedException {
+
+    ArrayList list = new ArrayList(EvalUtils.toCollection(self, loc, env));
+    if (key == Runtime.NONE) {
+      try {
+        Collections.sort(list, EvalUtils.SKYLARK_COMPARATOR);
+      } catch (EvalUtils.ComparisonException e) {
+        throw new EvalException(loc, e);
+      }
+    } else {
+      checkValidKeyFunc(key, loc);
+
+      final FuncallExpression ast = new FuncallExpression(Identifier.of(""), ImmutableList.of());
+
+      class KeyComparator implements Comparator<Object> {
+        Exception e;
+
+        @Override
+        public int compare(Object x, Object y) {
+          try {
+            return EvalUtils.SKYLARK_COMPARATOR.compare(
+                evalKeyFunc(x, key, loc, env, ast), evalKeyFunc(y, key, loc, env, ast));
+          } catch (InterruptedException | EvalException e) {
+            if (this.e == null) {
+              this.e = e;
+            }
+            return 0;
+          }
+        }
+      }
+
+      KeyComparator comp = new KeyComparator();
+      try {
+        Collections.sort(list, comp);
+      } catch (EvalUtils.ComparisonException e) {
+        throw new EvalException(loc, e);
+      }
+
+      if (comp.e != null) {
+        if (comp.e instanceof InterruptedException) {
+          throw (InterruptedException) comp.e;
+        }
+        throw (EvalException) comp.e;
+      }
+    }
+
+    if (reverse) {
+      Collections.reverse(list);
     }
+    return MutableList.wrapUnsafe(env, list);
   }
 
   @SkylarkCallable(

src/test/java/com/google/devtools/build/lib/syntax/MethodLibraryTest.java

@@ -249,6 +249,9 @@ public void testListSort() throws Exception {
         .testEval("sorted([True, False, True])", "[False, True, True]")
         .testEval("sorted(['a','x','b','z'])", "[\"a\", \"b\", \"x\", \"z\"]")
         .testEval("sorted({1: True, 5: True, 4: False})", "[1, 4, 5]")
+        .testEval("sorted([3, 2, 1, 0], reverse=True)", "[3, 2, 1, 0]")
+        .testEval("sorted([[1], [], [1, 2]], key=len, reverse=True)", "[[1, 2], [1], []]")
+        .testEval("sorted([[0, 5], [4, 1], [1, 7]], key=max)", "[[4, 1], [0, 5], [1, 7]]")
         .testIfExactError("Cannot compare function with function", "sorted([sorted, sorted])");
   }
 
@@ -712,7 +715,7 @@ public void testLegacyNamed() throws Exception {
         // Parameters which may be specified by keyword but are not explicitly 'named'.
         .testStatement("all(elements=[True, True])", Boolean.TRUE)
         .testStatement("any(elements=[True, False])", Boolean.TRUE)
-        .testEval("sorted(self=[3, 0, 2])", "[0, 2, 3]")
+        .testEval("sorted(self=[3, 0, 2], key=None, reverse=False)", "[0, 2, 3]")
         .testEval("reversed(sequence=[3, 2, 0])", "[0, 2, 3]")
         .testEval("tuple(x=[1, 2])", "(1, 2)")
         .testEval("list(x=(1, 2))", "[1, 2]")

src/test/starlark/testdata/sorted.sky

@@ -0,0 +1,37 @@
+assert_eq(sorted([42, 123, 3]), [3, 42, 123])
+assert_eq(sorted(["wiz", "foo", "bar"]), ["bar", "foo", "wiz"])
+assert_eq(sorted([42, 123, 3], reverse=True), [123, 42, 3])
+assert_eq(sorted(["wiz", "foo", "bar"], reverse=True), ["wiz", "foo", "bar"])
+assert_eq(sorted(list({"a": 1, "b": 2})), ['a', 'b'])
+---
+
+def f(x):
+  return x[0]
+pairs = [(4, 0), (3, 1), (4, 2), (2, 3), (3, 4), (1, 5), (2, 6), (3, 7)]
+
+assert_eq(sorted(pairs, key=f),
+         [(1, 5),
+          (2, 3), (2, 6),
+          (3, 1), (3, 4), (3, 7),
+          (4, 0), (4, 2)])
+
+---
+assert_eq(sorted(["two", "three", "four"], key=len),
+         ["two", "four", "three"])
+assert_eq(sorted(["two", "three", "four"], key=len, reverse=True),
+         ["three", "four", "two"])
+assert_eq(sorted([[1, 5], [0, 10], [4]], key=max),
+         [[4], [1, 5], [0, 10]])
+assert_eq(sorted([[1, 5], [0, 10], [4]], key=min, reverse=True),
+         [[4], [1, 5], [0, 10]])
+assert_eq(sorted([[2, 6, 1], [5, 2, 1], [1, 4, 2]], key=sorted),
+         [[1, 4, 2], [5, 2, 1], [2, 6, 1]])
+
+---
+sorted(1) ### type 'int' is not a collection
+---
+sorted([1, 2, None, 3]) ### Cannot compare NoneType with int
+---
+sorted([1, "one"]) ### Cannot compare string with int
+---
+sorted([1, 2, 3], key=1) ### "int" object is not callable