Merge pull request #110787 from tfc/cartesian-product

lib/attrsets: add cartesianProductOfSets function

4 files changed, 89 insertions, 3 deletions

diff --git a/lib/attrsets.nix b/lib/attrsets.nix
index d91d7a0cd47e..0ce3aaeca452 100644
--- a/lib/attrsets.nix
+++ b/lib/attrsets.nix
@@ -183,6 +183,24 @@ rec {
     else
       [];
 
+  /* Return the cartesian product of attribute set value combinations.
+
+    Example:
+      cartesianProductOfSets { a = [ 1 2 ]; b = [ 10 20 ]; }
+      => [
+           { a = 1; b = 10; }
+           { a = 1; b = 20; }
+           { a = 2; b = 10; }
+           { a = 2; b = 20; }
+         ]
+  */
+  cartesianProductOfSets = attrsOfLists:
+    lib.foldl' (listOfAttrs: attrName:
+      concatMap (attrs:
+        map (listValue: attrs // { ${attrName} = listValue; }) attrsOfLists.${attrName}
+      ) listOfAttrs
+    ) [{}] (attrNames attrsOfLists);
+
 
   /* Utility function that creates a {name, value} pair as expected by
      builtins.listToAttrs.
@@ -493,5 +511,4 @@ rec {
   zipWithNames = zipAttrsWithNames;
   zip = builtins.trace
     "lib.zip is deprecated, use lib.zipAttrsWith instead" zipAttrsWith;
-
 }
diff --git a/lib/default.nix b/lib/default.nix
index 803f1f765647..50320669e280 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -78,7 +78,7 @@ let
       zipAttrsWithNames zipAttrsWith zipAttrs recursiveUpdateUntil
       recursiveUpdate matchAttrs overrideExisting getOutput getBin
       getLib getDev getMan chooseDevOutputs zipWithNames zip
-      recurseIntoAttrs dontRecurseIntoAttrs;
+      recurseIntoAttrs dontRecurseIntoAttrs cartesianProductOfSets;
     inherit (self.lists) singleton forEach foldr fold foldl foldl' imap0 imap1
       concatMap flatten remove findSingle findFirst any all count
       optional optionals toList range partition zipListsWith zipLists
diff --git a/lib/lists.nix b/lib/lists.nix
index 06cee2eb112a..56af4d9daa18 100644
--- a/lib/lists.nix
+++ b/lib/lists.nix
@@ -629,7 +629,9 @@ rec {
       crossLists (x:y: "${toString x}${toString y}") [[1 2] [3 4]]
       => [ "13" "14" "23" "24" ]
   */
-  crossLists = f: foldl (fs: args: concatMap (f: map f args) fs) [f];
+  crossLists = builtins.trace
+    "lib.crossLists is deprecated, use lib.cartesianProductOfSets instead"
+    (f: foldl (fs: args: concatMap (f: map f args) fs) [f]);
 
 
   /* Remove duplicate elements from the list. O(n^2) complexity.
diff --git a/lib/tests/misc.nix b/lib/tests/misc.nix
index 35a5801c724f..0d249968402d 100644
--- a/lib/tests/misc.nix
+++ b/lib/tests/misc.nix
@@ -660,4 +660,71 @@ runTests {
     expected = [ [ "foo" ] [ "foo" "<name>" "bar" ] [ "foo" "bar" ] ];
   };
 
+  testCartesianProductOfEmptySet = {
+    expr = cartesianProductOfSets {};
+    expected = [ {} ];
+  };
+
+  testCartesianProductOfOneSet = {
+    expr = cartesianProductOfSets { a = [ 1 2 3 ]; };
+    expected = [ { a = 1; } { a = 2; } { a = 3; } ];
+  };
+
+  testCartesianProductOfTwoSets = {
+    expr = cartesianProductOfSets { a = [ 1 ]; b = [ 10 20 ]; };
+    expected = [
+      { a = 1; b = 10; }
+      { a = 1; b = 20; }
+    ];
+  };
+
+  testCartesianProductOfTwoSetsWithOneEmpty = {
+    expr = cartesianProductOfSets { a = [ ]; b = [ 10 20 ]; };
+    expected = [ ];
+  };
+
+  testCartesianProductOfThreeSets = {
+    expr = cartesianProductOfSets {
+      a = [   1   2   3 ];
+      b = [  10  20  30 ];
+      c = [ 100 200 300 ];
+    };
+    expected = [
+      { a = 1; b = 10; c = 100; }
+      { a = 1; b = 10; c = 200; }
+      { a = 1; b = 10; c = 300; }
+
+      { a = 1; b = 20; c = 100; }
+      { a = 1; b = 20; c = 200; }
+      { a = 1; b = 20; c = 300; }
+
+      { a = 1; b = 30; c = 100; }
+      { a = 1; b = 30; c = 200; }
+      { a = 1; b = 30; c = 300; }
+
+      { a = 2; b = 10; c = 100; }
+      { a = 2; b = 10; c = 200; }
+      { a = 2; b = 10; c = 300; }
+
+      { a = 2; b = 20; c = 100; }
+      { a = 2; b = 20; c = 200; }
+      { a = 2; b = 20; c = 300; }
+
+      { a = 2; b = 30; c = 100; }
+      { a = 2; b = 30; c = 200; }
+      { a = 2; b = 30; c = 300; }
+
+      { a = 3; b = 10; c = 100; }
+      { a = 3; b = 10; c = 200; }
+      { a = 3; b = 10; c = 300; }
+
+      { a = 3; b = 20; c = 100; }
+      { a = 3; b = 20; c = 200; }
+      { a = 3; b = 20; c = 300; }
+
+      { a = 3; b = 30; c = 100; }
+      { a = 3; b = 30; c = 200; }
+      { a = 3; b = 30; c = 300; }
+    ];
+  };
 }