From fe2bead78b1034052eca1d695118997f31826924 Mon Sep 17 00:00:00 2001
From: Gabriel Volpe <volpegabriel@gmail.com>
Date: Mon, 15 Apr 2024 19:16:15 +0200
Subject: [PATCH] lib/attrsets: introduce mapCartesianProduct

---
 lib/attrsets.nix   | 34 ++++++++++++++++++++++++++++++++++
 lib/default.nix    |  2 +-
 lib/tests/misc.nix | 25 +++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/lib/attrsets.nix b/lib/attrsets.nix
index 5c4f735b63da..83f8d0f34186 100644
--- a/lib/attrsets.nix
+++ b/lib/attrsets.nix
@@ -913,6 +913,40 @@ rec {
     ) [{}] (attrNames attrsOfLists);
 
 
+  /**
+    Return the result of function f applied to the cartesian product of attribute set value combinations.
+    Equivalent to using cartesianProduct followed by map.
+
+    # Inputs
+
+    `f`
+
+    : A function, given an attribute set, it returns a new value.
+
+    `attrsOfLists`
+
+    : Attribute set with attributes that are lists of values
+
+    # Type
+
+    ```
+    mapCartesianProduct :: (AttrSet -> a) -> AttrSet -> [a]
+    ```
+
+    # Examples
+    :::{.example}
+    ## `lib.attrsets.mapCartesianProduct` usage example
+
+    ```nix
+    mapCartesianProduct ({a, b}: "${a}-${b}") { a = [ "1" "2" ]; b = [ "3" "4" ]; }
+    => [ "1-3" "1-4" "2-3" "2-4" ]
+    ```
+
+    :::
+
+  */
+  mapCartesianProduct = f: attrsOfLists: map f (cartesianProduct attrsOfLists);
+
   /**
     Utility function that creates a `{name, value}` pair as expected by `builtins.listToAttrs`.
 
diff --git a/lib/default.nix b/lib/default.nix
index 97a81d4763bf..44a2a5216ec7 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -87,7 +87,7 @@ let
       recursiveUpdate matchAttrs mergeAttrsList overrideExisting showAttrPath getOutput
       getBin getLib getDev getMan chooseDevOutputs zipWithNames zip
       recurseIntoAttrs dontRecurseIntoAttrs cartesianProduct cartesianProductOfSets
-      updateManyAttrsByPath;
+      mapCartesianProduct updateManyAttrsByPath;
     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
diff --git a/lib/tests/misc.nix b/lib/tests/misc.nix
index b66f335c74f5..e15e110682fb 100644
--- a/lib/tests/misc.nix
+++ b/lib/tests/misc.nix
@@ -71,6 +71,7 @@ let
     makeIncludePath
     makeOverridable
     mapAttrs
+    mapCartesianProduct
     matchAttrs
     mergeAttrs
     meta
@@ -1753,6 +1754,30 @@ runTests {
     ];
   };
 
+  testMapCartesianProductOfOneSet = {
+    expr = mapCartesianProduct ({a}: a * 2) { a = [ 1 2 3 ]; };
+    expected = [ 2 4 6 ];
+  };
+
+  testMapCartesianProductOfTwoSets = {
+    expr = mapCartesianProduct ({a,b}: a + b) { a = [ 1 ]; b = [ 10 20 ]; };
+    expected = [ 11 21 ];
+  };
+
+  testMapCartesianProcutOfTwoSetsWithOneEmpty = {
+    expr = mapCartesianProduct (x: x.a + x.b) { a = [ ]; b = [ 10 20 ]; };
+    expected = [ ];
+  };
+
+  testMapCartesianProductOfThreeSets = {
+    expr = mapCartesianProduct ({a,b,c}: a + b + c) {
+      a = [ 1 2 3 ];
+      b = [ 10 20 30 ];
+      c = [ 100 200 300 ];
+    };
+    expected = [ 111 211 311 121 221 321 131 231 331 112 212 312 122 222 322 132 232 332 113 213 313 123 223 323 133 233 333 ];
+  };
+
   # The example from the showAttrPath documentation
   testShowAttrPathExample = {
     expr = showAttrPath [ "foo" "10" "bar" ];