Support comparison functions in array.sorted (#5627)

Co-authored-by: +merlan #flirora <uruwi@protonmail.com>
Co-authored-by: Laurenz <laurmaedje@gmail.com>

Cargo.lock

@@ -913,6 +913,12 @@ dependencies = [
  "weezl",
 ]
 
+[[package]]
+name = "glidesort"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2e102e6eb644d3e0b186fc161e4460417880a0a0b87d235f2e5b8fb30f2e9e0"
+
 [[package]]
 name = "half"
 version = "2.4.1"
@@ -3052,6 +3058,7 @@ dependencies = [
  "ecow",
  "flate2",
  "fontdb",
+ "glidesort",
  "hayagriva",
  "icu_properties",
  "icu_provider",

Cargo.toml

@@ -59,6 +59,7 @@ fastrand = "2.3"
 flate2 = "1"
 fontdb = { version = "0.23", default-features = false }
 fs_extra = "1.3"
+glidesort = "0.1.2"
 hayagriva = "0.8.1"
 heck = "0.5"
 hypher = "0.1.4"

crates/typst-library/Cargo.toml

@@ -29,6 +29,7 @@ csv = { workspace = true }
 ecow = { workspace = true }
 flate2 = { workspace = true }
 fontdb = { workspace = true }
+glidesort = { workspace = true }
 hayagriva = { workspace = true }
 icu_properties = { workspace = true }
 icu_provider = { workspace = true }

crates/typst-library/src/foundations/array.rs

@@ -808,7 +808,7 @@ impl Array {
     /// function. The sorting algorithm used is stable.
     ///
     /// Returns an error if two values could not be compared or if the key
-    /// function (if given) yields an error.
+    /// or comparison function (if given) yields an error.
     ///
     /// To sort according to multiple criteria at once, e.g. in case of equality
     /// between some criteria, the key function can return an array. The results
@@ -832,33 +832,134 @@ impl Array {
         /// determine the keys to sort by.
         #[named]
         key: Option<Func>,
+        /// If given, uses this function to compare elements in the array.
+        ///
+        /// This function should return a boolean: `{true}` indicates that the
+        /// elements are in order, while `{false}` indicates that they should be
+        /// swapped. To keep the sort stable, if the two elements are equal, the
+        /// function should return `{true}`.
+        ///
+        /// If this function does not order the elements properly (e.g., by
+        /// returning `{false}` for both `{(x, y)}` and `{(y, x)}`, or for
+        /// `{(x, x)}`), the resulting array will be in unspecified order.
+        ///
+        /// When used together with `key`, `by` will be passed the keys instead
+        /// of the elements.
+        ///
+        /// ```example
+        /// #(
+        ///   "sorted",
+        ///   "by",
+        ///   "decreasing",
+        ///   "length",
+        /// ).sorted(
+        ///   key: s => s.len(),
+        ///   by: (l, r) => l >= r,
+        /// )
+        /// ```
+        #[named]
+        by: 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),
-        };
-        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);
+        match by {
+            Some(by) => {
+                let mut are_in_order = |mut x, mut y| {
+                    if let Some(f) = &key {
+                        // We rely on `comemo`'s memoization of function
+                        // evaluation to not excessively reevaluate the key.
+                        x = f.call(engine, context, [x])?;
+                        y = f.call(engine, context, [y])?;
                     }
-                    Ordering::Equal
-                }),
-                (Err(e), _) | (_, Err(e)) => {
-                    if result.is_ok() {
-                        result = Err(e);
+                    match by.call(engine, context, [x, y])? {
+                        Value::Bool(b) => Ok(b),
+                        x => {
+                            bail!(
+                                span,
+                                "expected boolean from `by` function, got {}",
+                                x.ty(),
+                            )
+                        }
                     }
-                    Ordering::Equal
-                }
+                };
+                // If a comparison function is provided, we use `glidesort`
+                // instead of the standard library sorting algorithm to prevent
+                // panics in case the comparison function does not define a
+                // valid order (see https://github.com/typst/typst/pull/5627).
+                let mut result = Ok(());
+                let mut vec = self.0.into_iter().enumerate().collect::<Vec<_>>();
+                glidesort::sort_by(&mut vec, |(i, x), (j, y)| {
+                    // Because we use booleans for the comparison function, in
+                    // order to keep the sort stable, we need to compare in the
+                    // right order.
+                    if i < j {
+                        // If `x` and `y` appear in this order in the original
+                        // array, then we should change their order (i.e.,
+                        // return `Ordering::Greater`) iff `y` is strictly less
+                        // than `x` (i.e., `compare(x, y)` returns `false`).
+                        // Otherwise, we should keep them in the same order
+                        // (i.e., return `Ordering::Less`).
+                        match are_in_order(x.clone(), y.clone()) {
+                            Ok(false) => Ordering::Greater,
+                            Ok(true) => Ordering::Less,
+                            Err(err) => {
+                                if result.is_ok() {
+                                    result = Err(err);
+                                }
+                                Ordering::Equal
+                            }
+                        }
+                    } else {
+                        // If `x` and `y` appear in the opposite order in the
+                        // original array, then we should change their order
+                        // (i.e., return `Ordering::Less`) iff `x` is strictly
+                        // less than `y` (i.e., `compare(y, x)` returns
+                        // `false`). Otherwise, we should keep them in the same
+                        // order (i.e., return `Ordering::Less`).
+                        match are_in_order(y.clone(), x.clone()) {
+                            Ok(false) => Ordering::Less,
+                            Ok(true) => Ordering::Greater,
+                            Err(err) => {
+                                if result.is_ok() {
+                                    result = Err(err);
+                                }
+                                Ordering::Equal
+                            }
+                        }
+                    }
+                });
+                result.map(|()| vec.into_iter().map(|(_, x)| x).collect())
             }
-        });
-        result.map(|_| vec.into())
+
+            None => {
+                let mut key_of = |x: Value| match &key {
+                    // We rely on `comemo`'s memoization of function evaluation
+                    // to not excessively reevaluate the key.
+                    Some(f) => f.call(engine, context, [x]),
+                    None => Ok(x),
+                };
+                // If no comparison function is provided, we know the order is
+                // valid, so we can use the standard library sort and prevent an
+                // extra allocation.
+                let mut result = Ok(());
+                let mut vec = self.0;
+                vec.make_mut().sort_by(|a, b| {
+                    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
+                        }
+                    }
+                });
+                result.map(|()| vec.into())
+            }
+        }
     }
 
     /// Deduplicates all items in the array.

tests/suite/foundations/array.typ

@@ -359,6 +359,12 @@
 #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(by: (x, y) => x.len() < y.len()), ("I", "hi", "the", "text"))
+#test(("I", "the", "hi", "text").sorted(key: x => x.len(), by: (x, y) => y < x), ("text", "the", "hi", "I"))
+
+--- array-sorted-invalid-by-function ---
+// Error: 2-39 expected boolean from `by` function, got string
+#(1, 2, 3).sorted(by: (_, _) => "hmm")
 
 --- array-sorted-key-function-positional-1 ---
 // Error: 12-18 unexpected argument