From c7216c26f5b3c8e2af5db58ff06c01a0bf76e512 Mon Sep 17 00:00:00 2001
From: +merlan #flirora <uruwi@protonmail.com>
Date: Mon, 4 Nov 2024 15:17:36 -0500
Subject: [PATCH] Add compare argument for array.sorted

---
 crates/typst-library/src/foundations/array.rs | 47 ++++++++++++-------
 tests/suite/foundations/array.typ             |  2 +
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/crates/typst-library/src/foundations/array.rs b/crates/typst-library/src/foundations/array.rs
index 4667ee765..3412056a8 100644
--- a/crates/typst-library/src/foundations/array.rs
+++ b/crates/typst-library/src/foundations/array.rs
@@ -841,31 +841,42 @@ impl Array {
         /// determine the keys to sort by.
         #[named]
         key: Option<Func>,
+        /// If given, uses this function to compare elements in the array to
+        /// determine their relative order.
+        #[named]
+        compare: Option<Func>,
     ) -> SourceResult<Array> {
         let mut result = Ok(());
         let mut vec = self.0;
-        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(engine, context, [x]),
-            None => Ok(x),
+        let mut compare = |x: Value, y: Value| {
+            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(engine, context, [x]),
+                None => Ok(x),
+            };
+            let x = key_of(x)?;
+            let y = key_of(y)?;
+            match &compare {
+                Some(f) => Ok(match f.call(engine, context, [x, y])? {
+                    Value::Int(x) => x.cmp(&0),
+                    x => bail!(
+                        span,
+                        "expected integer from `compare` function; got {}",
+                        x.repr()
+                    ),
+                }),
+                None => ops::compare(&x, &y).at(span),
+            }
         };
         vec.make_mut().sort_by(|a, b| {
             // Until we get `try` blocks :)
-            match (key_of(a.clone()), key_of(b.clone())) {
-                (Ok(a), Ok(b)) => ops::compare(&a, &b).unwrap_or_else(|err| {
-                    if result.is_ok() {
-                        result = Err(err).at(span);
-                    }
-                    Ordering::Equal
-                }),
-                (Err(e), _) | (_, Err(e)) => {
-                    if result.is_ok() {
-                        result = Err(e);
-                    }
-                    Ordering::Equal
+            compare(a.clone(), b.clone()).unwrap_or_else(|err| {
+                if result.is_ok() {
+                    result = Err(err);
                 }
-            }
+                Ordering::Equal
+            })
         });
         result.map(|_| vec.into())
     }
diff --git a/tests/suite/foundations/array.typ b/tests/suite/foundations/array.typ
index 6228f471b..2fbd185d6 100644
--- a/tests/suite/foundations/array.typ
+++ b/tests/suite/foundations/array.typ
@@ -355,6 +355,8 @@
 #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))
+#test(("I", "the", "hi", "text").sorted(compare: (x, y) => x.len() - y.len()), ("I", "hi", "the", "text"))
+#test(("I", "the", "hi", "text").sorted(key: x => x.len(), compare: (x, y) => y - x), ("text", "the", "hi", "I"))
 
 --- array-sorted-key-function-positional-1 ---
 // Error: 12-18 unexpected argument