Array sorting by key (#584)

docs/src/reference/types.md

@@ -657,6 +657,8 @@ Combine all items in the array into one.
 ### sorted()
 Return a new array with the same items, but sorted.
 
+- key: function (named)
+  If given, applies this function to the elements in the array to determine the keys to sort by.
 - returns: array
 
 # Dictionary

src/eval/array.rs

@@ -6,6 +6,7 @@ use ecow::{eco_format, EcoString, EcoVec};
 
 use super::{ops, Args, Func, Value, Vm};
 use crate::diag::{At, SourceResult, StrResult};
+use crate::syntax::Span;
 use crate::util::pretty_array_like;
 
 /// Create a new [`Array`] from values.
@@ -276,23 +277,45 @@ impl Array {
         Ok(result)
     }
 
-    /// Return a sorted version of this array.
+    /// Return a sorted version of this array, optionally by a given key function.
     ///
-    /// Returns an error if two values could not be compared.
-    pub fn sorted(&self) -> StrResult<Self> {
+    /// Returns an error if two values could not be compared or if the key function (if given)
+    /// yields an error.
+    pub fn sorted(
+        &self,
+        vm: &mut Vm,
+        span: Span,
+        key: Option<Func>,
+    ) -> SourceResult<Self> {
         let mut result = Ok(());
         let mut vec = self.0.clone();
+        let mut key_of = |x: Value| match &key {
+            // NOTE: We are relying on `comemo`'s memoization of function
+            // evaluation to not excessively reevaluate the `key`.
+            Some(f) => f.call_vm(vm, Args::new(f.span(), [x])),
+            None => Ok(x),
+        };
         vec.make_mut().sort_by(|a, b| {
-            a.partial_cmp(b).unwrap_or_else(|| {
-                if result.is_ok() {
-                    result = Err(eco_format!(
-                        "cannot order {} and {}",
-                        a.type_name(),
-                        b.type_name(),
-                    ));
+            // Until we get `try` blocks :)
+            match (key_of(a.clone()), key_of(b.clone())) {
+                (Ok(a), Ok(b)) => a.partial_cmp(&b).unwrap_or_else(|| {
+                    if result.is_ok() {
+                        result = Err(eco_format!(
+                            "cannot order {} and {}",
+                            a.type_name(),
+                            b.type_name(),
+                        ))
+                        .at(span);
+                    }
+                    Ordering::Equal
+                }),
+                (Err(e), _) | (_, Err(e)) => {
+                    if result.is_ok() {
+                        result = Err(e);
+                    }
+                    Ordering::Equal
                 }
-                Ordering::Equal
-            })
+            }
         });
         result.map(|_| Self::from_vec(vec))
     }

src/eval/methods.rs

@@ -115,7 +115,7 @@ pub fn call(
                 let last = args.named("last")?;
                 array.join(sep, last).at(span)?
             }
-            "sorted" => Value::Array(array.sorted().at(span)?),
+            "sorted" => Value::Array(array.sorted(vm, span, args.named("key")?)?),
             "enumerate" => Value::Array(array.enumerate()),
             _ => return missing(),
         },

tests/typ/compiler/array.typ

@@ -193,9 +193,18 @@
 ---
 // Test the `sorted` method.
 #test(().sorted(), ())
+#test(().sorted(key: x => x), ())
 #test(((true, false) * 10).sorted(), (false,) * 10 + (true,) * 10)
 #test(("it", "the", "hi", "text").sorted(), ("hi", "it", "text", "the"))
+#test(("I", "the", "hi", "text").sorted(key: x => x), ("I", "hi", "text", "the"))
+#test(("I", "the", "hi", "text").sorted(key: x => x.len()), ("I", "hi", "the", "text"))
 #test((2, 1, 3, 10, 5, 8, 6, -7, 2).sorted(), (-7, 1, 2, 2, 3, 5, 6, 8, 10))
+#test((2, 1, 3, -10, -5, 8, 6, -7, 2).sorted(key: x => x), (-10, -7, -5, 1, 2, 2, 3, 6, 8))
+#test((2, 1, 3, -10, -5, 8, 6, -7, 2).sorted(key: x => x * x), (1, 2, 2, 3, -5, 6, -7, 8, -10))
+
+---
+// Error: 32-37 cannot divide by zero
+#(1, 2, 0, 3).sorted(key: x => 5 / x)
 
 ---
 // Error: 2-26 cannot order content and content