From 327fe921d520a5896d18ecf1e2c017d47daf0b18 Mon Sep 17 00:00:00 2001
From: Johannes Kirschbauer <hsjobeki@gmail.com>
Date: Sat, 9 Nov 2024 11:22:25 +0100
Subject: [PATCH] lib/types: init {types.attrsWith}

---
 lib/tests/modules.sh                 |  3 +
 lib/tests/modules/lazy-attrsWith.nix | 45 ++++++++++++++
 lib/types.nix                        | 93 ++++++++++++++++++----------
 3 files changed, 108 insertions(+), 33 deletions(-)
 create mode 100644 lib/tests/modules/lazy-attrsWith.nix

diff --git a/lib/tests/modules.sh b/lib/tests/modules.sh
index 2a48d88cbf732f..4d1a7fd8516e24 100755
--- a/lib/tests/modules.sh
+++ b/lib/tests/modules.sh
@@ -569,6 +569,9 @@ checkConfigOutput '^38|27$' options.submoduleLine38.declarationPositions.1.line
 # nested options work
 checkConfigOutput '^34$' options.nested.nestedLine34.declarationPositions.0.line ./declaration-positions.nix
 
+# AttrsWith tests
+checkConfigOutput '^11$' config.result ./lazy-attrsWith.nix
+
 cat <<EOF
 ====== module tests ======
 $pass Pass
diff --git a/lib/tests/modules/lazy-attrsWith.nix b/lib/tests/modules/lazy-attrsWith.nix
new file mode 100644
index 00000000000000..0fdd5534dac649
--- /dev/null
+++ b/lib/tests/modules/lazy-attrsWith.nix
@@ -0,0 +1,45 @@
+# Check that AttrsWith { lazy = true; } is lazy
+{ lib, ... }:
+let
+  inherit (lib) types mkOption;
+in
+{
+  imports = [
+    #  Module A
+    (
+      { ... }:
+      {
+        options.mergedLazy = mkOption {
+          # Same as lazyAttrsOf
+          type = types.attrsWith {
+            lazy = true;
+            elemType = types.int;
+          };
+        };
+      }
+    )
+    # Module B
+    (
+      { ... }:
+      {
+        options.mergedLazy = lib.mkOption {
+          # Same as lazyAttrsOf
+          type = types.attrsWith {
+            lazy = true;
+            elemType = types.int;
+          };
+        };
+      }
+    )
+    # Result
+    (
+      { config, ... }:
+      {
+        # Can only evaluate if lazy
+        config.mergedLazy.bar = config.mergedLazy.baz + 1;
+        config.mergedLazy.baz = 10;
+        options.result = mkOption { default = config.mergedLazy.bar; };
+      }
+    )
+  ];
+}
diff --git a/lib/types.nix b/lib/types.nix
index 86b717afd1227e..41a695165afcb7 100644
--- a/lib/types.nix
+++ b/lib/types.nix
@@ -94,11 +94,11 @@ rec {
             && (f.payload == null && f'.payload == null)
        then f.type
     # composed types
-    else if (f.wrapped != null && f'.wrapped != null) && (wrapped != null)
-       then f.type wrapped
     # value types
     else if (f.payload != null && f'.payload != null) && (payload != null)
        then f.type payload
+    else if (f.wrapped != null && f'.wrapped != null) && (wrapped != null)
+       then f.type wrapped
     else null;
 
   # Default type functor
@@ -574,48 +574,75 @@ rec {
         substSubModules = m: nonEmptyListOf (elemType.substSubModules m);
       };
 
-    attrsOf = elemType: mkOptionType rec {
-      name = "attrsOf";
-      description = "attribute set of ${optionDescriptionPhrase (class: class == "noun" || class == "composite") elemType}";
-      descriptionClass = "composite";
-      check = isAttrs;
-      merge = loc: defs:
-        mapAttrs (n: v: v.value) (filterAttrs (n: v: v ? value) (zipAttrsWith (name: defs:
-            (mergeDefinitions (loc ++ [name]) elemType defs).optionalValue
-          )
-          # Push down position info.
-          (map (def: mapAttrs (n: v: { inherit (def) file; value = v; }) def.value) defs)));
-      emptyValue = { value = {}; };
-      getSubOptions = prefix: elemType.getSubOptions (prefix ++ ["<name>"]);
-      getSubModules = elemType.getSubModules;
-      substSubModules = m: attrsOf (elemType.substSubModules m);
-      functor = (defaultFunctor name) // { wrapped = elemType; };
-      nestedTypes.elemType = elemType;
-    };
+    attrsOf = elemType: attrsWith { inherit elemType; };
 
     # A version of attrsOf that's lazy in its values at the expense of
     # conditional definitions not working properly. E.g. defining a value with
     # `foo.attr = mkIf false 10`, then `foo ? attr == true`, whereas with
     # attrsOf it would correctly be `false`. Accessing `foo.attr` would throw an
     # error that it's not defined. Use only if conditional definitions don't make sense.
-    lazyAttrsOf = elemType: mkOptionType rec {
-      name = "lazyAttrsOf";
-      description = "lazy attribute set of ${optionDescriptionPhrase (class: class == "noun" || class == "composite") elemType}";
+    lazyAttrsOf = elemType: attrsWith { inherit elemType; lazy = true; };
+
+    # base type for lazyAttrsOf and attrsOf
+    attrsWith = {
+      elemType,
+      lazy ? false,
+    }:
+    let
+      typeName = if lazy then "lazyAttrsOf" else "attrsOf";
+      # Push down position info.
+      pushPositions = map (def: mapAttrs (n: v: { inherit (def) file; value = v; }) def.value);
+    in
+    mkOptionType {
+      name = typeName;
+      description = (if lazy then "lazy attribute set" else "attribute set") + " of ${optionDescriptionPhrase (class: class == "noun" || class == "composite") elemType}";
       descriptionClass = "composite";
       check = isAttrs;
-      merge = loc: defs:
-        zipAttrsWith (name: defs:
-          let merged = mergeDefinitions (loc ++ [name]) elemType defs;
-          # mergedValue will trigger an appropriate error when accessed
-          in merged.optionalValue.value or elemType.emptyValue.value or merged.mergedValue
-        )
-        # Push down position info.
-        (map (def: mapAttrs (n: v: { inherit (def) file; value = v; }) def.value) defs);
+      merge = if lazy then (
+        # Lazy merge Function
+        loc: defs:
+          zipAttrsWith (name: defs:
+            let merged = mergeDefinitions (loc ++ [name]) elemType defs;
+            # mergedValue will trigger an appropriate error when accessed
+            in merged.optionalValue.value or elemType.emptyValue.value or merged.mergedValue
+          )
+          # Push down position info.
+          (pushPositions defs)
+      ) else (
+        # Non-lazy merge Function
+        loc: defs:
+          mapAttrs (n: v: v.value) (filterAttrs (n: v: v ? value) (zipAttrsWith (name: defs:
+              (mergeDefinitions (loc ++ [name]) elemType (defs)).optionalValue
+            )
+          # Push down position info.
+          (pushPositions defs)))
+      );
       emptyValue = { value = {}; };
       getSubOptions = prefix: elemType.getSubOptions (prefix ++ ["<name>"]);
       getSubModules = elemType.getSubModules;
-      substSubModules = m: lazyAttrsOf (elemType.substSubModules m);
-      functor = (defaultFunctor name) // { wrapped = elemType; };
+      substSubModules = m: attrsWith { elemType = elemType.substSubModules m; inherit lazy; };
+      functor = defaultFunctor "attrsWith" // {
+        wrapped = elemType;
+        payload = {
+          # Important!: Add new function attributes here in case of future changes
+          inherit elemType lazy;
+        };
+        binOp = lhs: rhs:
+          let
+            elemType = lhs.elemType.typeMerge rhs.elemType.functor;
+            lazy =
+              if lhs.lazy == rhs.lazy then
+                lhs.lazy
+              else
+                null;
+          in
+          if elemType == null || lazy == null then
+            null
+          else
+            {
+              inherit elemType lazy;
+            };
+      };
       nestedTypes.elemType = elemType;
     };