From 08aa61455246c691d5bd0ed35f067b676d63f05b Mon Sep 17 00:00:00 2001
From: Robert Hensing <robert@roberthensing.nl>
Date: Thu, 7 Dec 2023 19:51:07 +0100
Subject: lib.sortOn: init

A more efficient sort in some cases, and often convenient.

This exposes `lib.lists.sortOn` immediately on `lib`, because it is
a sibling of `sort`, which is already present there.
Omitting it would lead to more confusion, and worse outcomes.
There's no confusion about the types `sort` or `sortOn` operate on.

Haskell agrees about the type for `sortOn`, and it is in its `base`.

(cherry picked from commit 67cc78643d7307b7a4cd783e43c53346fdea18bf)
---
 lib/default.nix    |  2 +-
 lib/lists.nix      | 40 ++++++++++++++++++++++++++++++++++++++++
 lib/tests/misc.nix | 22 ++++++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

(limited to 'lib')

diff --git a/lib/default.nix b/lib/default.nix
index a2958e561cf32..35e31af364d88 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -91,7 +91,7 @@ let
     inherit (self.lists) singleton forEach foldr fold foldl foldl' imap0 imap1
       concatMap flatten remove findSingle findFirst any all count
       optional optionals toList range replicate partition zipListsWith zipLists
-      reverseList listDfs toposort sort naturalSort compareLists take
+      reverseList listDfs toposort sort sortOn naturalSort compareLists take
       drop sublist last init crossLists unique allUnique intersectLists
       subtractLists mutuallyExclusive groupBy groupBy';
     inherit (self.strings) concatStrings concatMapStrings concatImapStrings
diff --git a/lib/lists.nix b/lib/lists.nix
index a56667ec9c388..83821ce759187 100644
--- a/lib/lists.nix
+++ b/lib/lists.nix
@@ -4,6 +4,7 @@ let
   inherit (lib.strings) toInt;
   inherit (lib.trivial) compare min id;
   inherit (lib.attrsets) mapAttrs;
+  inherit (lib.lists) sort;
 in
 rec {
 
@@ -591,6 +592,9 @@ rec {
      the second argument.  The returned list is sorted in an increasing
      order.  The implementation does a quick-sort.
 
+     See also [`sortOn`](#function-library-lib.lists.sortOn), which applies the
+     default comparison on a function-derived property, and may be more efficient.
+
      Example:
        sort (a: b: a < b) [ 5 3 7 ]
        => [ 3 5 7 ]
@@ -612,6 +616,42 @@ rec {
       if len < 2 then list
       else (sort strictLess pivot.left) ++  [ first ] ++  (sort strictLess pivot.right));
 
+  /*
+    Sort a list based on the default comparison of a derived property `b`.
+
+    The items are returned in `b`-increasing order.
+
+    **Performance**:
+
+    The passed function `f` is only evaluated once per item,
+    unlike an unprepared [`sort`](#function-library-lib.lists.sort) using
+    `f p < f q`.
+
+    **Laws**:
+    ```nix
+    sortOn f == sort (p: q: f p < f q)
+    ```
+
+    Example:
+      sortOn stringLength [ "aa" "b" "cccc" ]
+      => [ "b" "aa" "cccc" ]
+
+    Type:
+      sortOn :: (a -> b) -> [a] -> [a], for comparable b
+  */
+  sortOn = f: list:
+    let
+      # Heterogenous list as pair may be ugly, but requires minimal allocations.
+      pairs = map (x: [(f x) x]) list;
+    in
+      map
+        (x: builtins.elemAt x 1)
+        (sort
+          # Compare the first element of the pairs
+          # Do not factor out the `<`, to avoid calls in hot code; duplicate instead.
+          (a: b: head a < head b)
+          pairs);
+
   /* Compare two lists element-by-element.
 
      Example:
diff --git a/lib/tests/misc.nix b/lib/tests/misc.nix
index 9f1fee2ba2341..608af656d02c0 100644
--- a/lib/tests/misc.nix
+++ b/lib/tests/misc.nix
@@ -650,6 +650,28 @@ runTests {
     expected = [2 30 40 42];
   };
 
+  testSortOn = {
+    expr = sortOn stringLength [ "aa" "b" "cccc" ];
+    expected = [ "b" "aa" "cccc" ];
+  };
+
+  testSortOnEmpty = {
+    expr = sortOn (throw "nope") [ ];
+    expected = [ ];
+  };
+
+  testSortOnIncomparable = {
+    expr =
+      map
+        (x: x.f x.ok)
+        (sortOn (x: x.ok) [
+          { ok = 1; f = x: x; }
+          { ok = 3; f = x: x + 3; }
+          { ok = 2; f = x: x; }
+        ]);
+    expected = [ 1 2 6 ];
+  };
+
   testReplicate = {
     expr = replicate 3 "a";
     expected = ["a" "a" "a"];
-- 
cgit 1.4.1