lib.attrsets: Introduce updateManyAttrsByPath

3 files changed, 214 insertions, 2 deletions

diff --git a/lib/attrsets.nix b/lib/attrsets.nix
index 7a3df9059c0..28e6cd09f85 100644
--- a/lib/attrsets.nix
+++ b/lib/attrsets.nix
@@ -5,7 +5,7 @@ let
   inherit (builtins) head tail length;
   inherit (lib.trivial) id;
   inherit (lib.strings) concatStringsSep concatMapStringsSep escapeNixIdentifier sanitizeDerivationName;
-  inherit (lib.lists) foldr foldl' concatMap concatLists elemAt all;
+  inherit (lib.lists) foldr foldl' concatMap concatLists elemAt all partition groupBy take foldl;
 in
 
 rec {
@@ -78,6 +78,103 @@ rec {
     in attrByPath attrPath (abort errorMsg);
 
 
+  /* Update or set specific paths of an attribute set.
+
+     Takes a list of updates to apply and an attribute set to apply them to,
+     and returns the attribute set with the updates applied. Updates are
+     represented as { path = ...; update = ...; } values, where `path` is a
+     list of strings representing the attribute path that should be updated,
+     and `update` is a function that takes the old value at that attribute path
+     as an argument and returns the new
+     value it should be.
+
+     Properties:
+     - Updates to deeper attribute paths are applied before updates to more
+       shallow attribute paths
+     - Multiple updates to the same attribute path are applied in the order
+       they appear in the update list
+     - If any but the last `path` element leads into a value that is not an
+       attribute set, an error is thrown
+     - If there is an update for an attribute path that doesn't exist,
+       accessing the argument in the update function causes an error, but
+       intermediate attribute sets are implicitly created as needed
+
+     Example:
+       updateManyAttrsByPath [
+         {
+           path = [ "a" "b" ];
+           update = old: { d = old.c; };
+         }
+         {
+           path = [ "a" "b" "c" ];
+           update = old: old + 1;
+         }
+         {
+           path = [ "x" "y" ];
+           update = old: "xy";
+         }
+       ] { a.b.c = 0; }
+       => { a = { b = { d = 1; }; }; x = { y = "xy"; }; }
+  */
+  updateManyAttrsByPath = let
+    # When recursing into attributes, instead of updating the `path` of each
+    # update using `tail`, which needs to allocate an entirely new list,
+    # we just pass a prefix length to use and make sure to only look at the
+    # path without the prefix length, so that we can reuse the original list
+    # entries.
+    go = prefixLength: hasValue: value: updates:
+      let
+        # Splits updates into ones on this level (split.right)
+        # And ones on levels further down (split.wrong)
+        split = partition (el: length el.path == prefixLength) updates;
+
+        # Groups updates on further down levels into the attributes they modify
+        nested = groupBy (el: elemAt el.path prefixLength) split.wrong;
+
+        # Applies only nested modification to the input value
+        withNestedMods =
+          # Return the value directly if we don't have any nested modifications
+          if split.wrong == [] then
+            if hasValue then value
+            else
+              # Throw an error if there is no value. This `head` call here is
+              # safe, but only in this branch since `go` could only be called
+              # with `hasValue == false` for nested updates, in which case
+              # it's also always called with at least one update
+              let updatePath = (head split.right).path; in
+              throw
+              ( "updateManyAttrsByPath: Path '${showAttrPath updatePath}' does "
+              + "not exist in the given value, but the first update to this "
+              + "path tries to access the existing value.")
+          else
+            # If there are nested modifications, try to apply them to the value
+            if ! hasValue then
+              # But if we don't have a value, just use an empty attribute set
+              # as the value, but simplify the code a bit
+              mapAttrs (name: go (prefixLength + 1) false null) nested
+            else if isAttrs value then
+              # If we do have a value and it's an attribute set, override it
+              # with the nested modifications
+              value //
+              mapAttrs (name: go (prefixLength + 1) (value ? ${name}) value.${name}) nested
+            else
+              # However if it's not an attribute set, we can't apply the nested
+              # modifications, throw an error
+              let updatePath = (head split.wrong).path; in
+              throw
+              ( "updateManyAttrsByPath: Path '${showAttrPath updatePath}' needs to "
+              + "be updated, but path '${showAttrPath (take prefixLength updatePath)}' "
+              + "of the given value is not an attribute set, so we can't "
+              + "update an attribute inside of it.");
+
+        # We get the final result by applying all the updates on this level
+        # after having applied all the nested updates
+        # We use foldl instead of foldl' so that in case of multiple updates,
+        # intermediate values aren't evaluated if not needed
+      in foldl (acc: el: el.update acc) withNestedMods split.right;
+
+  in updates: value: go 0 true value updates;
+
   /* Return the specified attributes from a set.
 
      Example:
diff --git a/lib/default.nix b/lib/default.nix
index 6e5465bc003..f8ab51c6579 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -80,7 +80,8 @@ let
       zipAttrsWithNames zipAttrsWith zipAttrs recursiveUpdateUntil
       recursiveUpdate matchAttrs overrideExisting showAttrPath getOutput getBin
       getLib getDev getMan chooseDevOutputs zipWithNames zip
-      recurseIntoAttrs dontRecurseIntoAttrs cartesianProductOfSets;
+      recurseIntoAttrs dontRecurseIntoAttrs cartesianProductOfSets
+      updateManyAttrsByPath;
     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/tests/misc.nix b/lib/tests/misc.nix
index c2e76f18d95..27111903139 100644
--- a/lib/tests/misc.nix
+++ b/lib/tests/misc.nix
@@ -799,4 +799,118 @@ runTests {
     expr = groupBy' builtins.add 0 (x: boolToString (x > 2)) [ 5 1 2 3 4 ];
     expected = { false = 3; true = 12; };
   };
+
+  # The example from the updateManyAttrsByPath documentation
+  testUpdateManyAttrsByPathExample = {
+    expr = updateManyAttrsByPath [
+      {
+        path = [ "a" "b" ];
+        update = old: { d = old.c; };
+      }
+      {
+        path = [ "a" "b" "c" ];
+        update = old: old + 1;
+      }
+      {
+        path = [ "x" "y" ];
+        update = old: "xy";
+      }
+    ] { a.b.c = 0; };
+    expected = { a = { b = { d = 1; }; }; x = { y = "xy"; }; };
+  };
+
+  # If there are no updates, the value is passed through
+  testUpdateManyAttrsByPathNone = {
+    expr = updateManyAttrsByPath [] "something";
+    expected = "something";
+  };
+
+  # A single update to the root path is just like applying the function directly
+  testUpdateManyAttrsByPathSingleIncrement = {
+    expr = updateManyAttrsByPath [
+      {
+        path = [ ];
+        update = old: old + 1;
+      }
+    ] 0;
+    expected = 1;
+  };
+
+  # Multiple updates can be applied are done in order
+  testUpdateManyAttrsByPathMultipleIncrements = {
+    expr = updateManyAttrsByPath [
+      {
+        path = [ ];
+        update = old: old + "a";
+      }
+      {
+        path = [ ];
+        update = old: old + "b";
+      }
+      {
+        path = [ ];
+        update = old: old + "c";
+      }
+    ] "";
+    expected = "abc";
+  };
+
+  # If an update doesn't use the value, all previous updates are not evaluated
+  testUpdateManyAttrsByPathLazy = {
+    expr = updateManyAttrsByPath [
+      {
+        path = [ ];
+        update = old: old + throw "nope";
+      }
+      {
+        path = [ ];
+        update = old: "untainted";
+      }
+    ] (throw "start");
+    expected = "untainted";
+  };
+
+  # Deeply nested attributes can be updated without affecting others
+  testUpdateManyAttrsByPathDeep = {
+    expr = updateManyAttrsByPath [
+      {
+        path = [ "a" "b" "c" ];
+        update = old: old + 1;
+      }
+    ] {
+      a.b.c = 0;
+
+      a.b.z = 0;
+      a.y.z = 0;
+      x.y.z = 0;
+    };
+    expected = {
+      a.b.c = 1;
+
+      a.b.z = 0;
+      a.y.z = 0;
+      x.y.z = 0;
+    };
+  };
+
+  # Nested attributes are updated first
+  testUpdateManyAttrsByPathNestedBeforehand = {
+    expr = updateManyAttrsByPath [
+      {
+        path = [ "a" ];
+        update = old: old // { x = old.b; };
+      }
+      {
+        path = [ "a" "b" ];
+        update = old: old + 1;
+      }
+    ] {
+      a.b = 0;
+    };
+    expected = {
+      a.b = 1;
+      a.x = 1;
+    };
+  };
+
 }