From 5f0406503f04aad703b7de4d17e09385a50a702d Mon Sep 17 00:00:00 2001
From: Daniel Olsen <daniel.olsen99@gmail.com>
Date: Fri, 21 Oct 2022 03:43:30 +0200
Subject: [PATCH] lib/attrsets: Add more ways to map over/to attrsets

These functions have been particularly useful while developing nixos modules

Co-authored-by: h7x4 <h7x4@nani.wtf>
---
 lib/attrsets.nix | 34 +++++++++++++++++++++++++++++++++-
 lib/default.nix  |  5 +++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/lib/attrsets.nix b/lib/attrsets.nix
index d9a6eab0603ed..7d4ab2eb797ad 100644
--- a/lib/attrsets.nix
+++ b/lib/attrsets.nix
@@ -5,7 +5,7 @@ let
   inherit (builtins) head tail length;
   inherit (lib.trivial) flip id mergeAttrs pipe;
   inherit (lib.strings) concatStringsSep concatMapStringsSep escapeNixIdentifier sanitizeDerivationName;
-  inherit (lib.lists) foldr foldl' concatMap concatLists elemAt all partition groupBy take foldl;
+  inherit (lib.lists) foldr foldl' concatMap concatLists elemAt all partition groupBy take foldl imap0;
 in
 
 rec {
@@ -518,6 +518,27 @@ rec {
     set:
     listToAttrs (map (attr: f attr set.${attr}) (attrNames set));
 
+  /* Like mapAttrs, but additionally taking a variable that
+     will be automatically incremented, starting from 0.
+
+     Example:
+       imapAttrs (i: name: value: value + i)
+          { a = 1; b = 1; c = 1; }
+       => { a = 1; b = 2; c = 3; }
+  */
+  imapAttrs = f: set:
+    listToAttrs (imap0 (i: attr: nameValuePair attr (f i attr set.${attr})) (attrNames set));
+
+  /* Like mapAttrs', but additionally taking a variable that
+     will be automatically incremented, starting from 0.
+
+     Example:
+       imapAttrs' (i: name: value: nameValuePair "${name}${toString i}" value + i)
+          { a = 2; b = 3; }
+       => { a0 = 2; b1 = 4 }
+  */
+  imapAttrs' = f: set:
+    listToAttrs (imap0 (i: attr: f i attr set.${attr}) (attrNames set));
 
   /* Call a function for each attribute in the given set and return
      the result in a list.
@@ -617,6 +638,17 @@ rec {
     f:
     listToAttrs (map (n: nameValuePair n (f n)) names);
 
+  /* Generate an attribute set by mapping over a list of items
+     Taking two functions, one to apply for the name and one to
+     apply for the value of the attribute set.
+
+     Example:
+       genAttrs' [ {a = 2; b.c = "a" } {a = 1; b.c = "b"} ]
+          (i: i.b.c) (i: i.a)
+       => { a = 2; b = 1; }
+  */
+  genAttrs' = items: f: g:
+    listToAttrs (map (i: nameValuePair (f i) (g i)) items);
 
   /* Check whether the argument is a derivation. Any set with
      `{ type = "derivation"; }` counts as a derivation.
diff --git a/lib/default.nix b/lib/default.nix
index 0424db36b2e99..d744df2ca1b1a 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -79,8 +79,9 @@ let
     inherit (self.attrsets) attrByPath hasAttrByPath setAttrByPath
       getAttrFromPath attrVals attrValues getAttrs catAttrs filterAttrs
       filterAttrsRecursive foldlAttrs foldAttrs collect nameValuePair mapAttrs
-      mapAttrs' mapAttrsToList concatMapAttrs mapAttrsRecursive mapAttrsRecursiveCond
-      genAttrs isDerivation toDerivation optionalAttrs
+      mapAttrs' imapAttrs imapAttrs' mapAttrsToList
+      concatMapAttrs mapAttrsRecursive mapAttrsRecursiveCond
+      genAttrs genAttrs' isDerivation toDerivation optionalAttrs
       zipAttrsWithNames zipAttrsWith zipAttrs recursiveUpdateUntil
       recursiveUpdate matchAttrs overrideExisting showAttrPath getOutput getBin
       getLib getDev getMan chooseDevOutputs zipWithNames zip