From a8087a9dbb8b24b451bfaa5d31e07e9b5574226d Mon Sep 17 00:00:00 2001
From: Daniel Csillag <dccsillag@gmail.com>
Date: Tue, 11 Apr 2023 07:48:17 -0300
Subject: Array sorting by key (#584)

---
 src/eval/array.rs   | 47 +++++++++++++++++++++++++++++++++++------------
 src/eval/methods.rs |  2 +-
 2 files changed, 36 insertions(+), 13 deletions(-)

(limited to 'src')

diff --git a/src/eval/array.rs b/src/eval/array.rs
index 394191ea..6bd2eb47 100644
--- a/src/eval/array.rs
+++ b/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))
     }
diff --git a/src/eval/methods.rs b/src/eval/methods.rs
index 8b364fcb..452b90da 100644
--- a/src/eval/methods.rs
+++ b/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(),
         },
-- 
cgit v1.2.3