lib.extendMkDerivation: init

Add functions extendMkDerivation to lib.customisation.

Co-authored-by: Robert Hensing <roberth@users.noreply.github.com>
Co-authored-by: Valentin Gagarin <valentin.gagarin@tweag.io>
Co-authored-by: Lin Jian <me@linj.tech>
Co-authored-by: Philip Taron <philip.taron@gmail.com>

3 files changed, 195 insertions, 2 deletions

diff --git a/lib/customisation.nix b/lib/customisation.nix
index be6585d5fa6e..9c50e6a26c3d 100644
--- a/lib/customisation.nix
+++ b/lib/customisation.nix
@@ -31,6 +31,8 @@ let
     flatten
     deepSeq
     extends
+    toFunction
+    id
     ;
   inherit (lib.strings) levenshtein levenshteinAtMost;
 
@@ -730,4 +732,126 @@ rec {
     in
     self;
 
+  /**
+    Define a `mkDerivation`-like function based on another `mkDerivation`-like function.
+
+    [`stdenv.mkDerivation`](#part-stdenv) gives access to
+    its final set of derivation attributes when it is passed a function,
+    or when it is passed an overlay-style function in `overrideAttrs`.
+
+    Instead of composing new `stdenv.mkDerivation`-like build helpers
+    using normal function composition,
+    `extendMkDerivation` makes sure that the returned build helper
+    supports such first class recursion like `mkDerivation` does.
+
+    `extendMkDerivation` takes an extra attribute set to configure its behaviour.
+    One can optionally specify
+    `transformDrv` to specify a function to apply to the result derivation,
+    or `inheritFunctionArgs` to decide whether to inherit the `__functionArgs`
+    from the base build helper.
+
+    # Inputs
+
+    `extendMkDerivation`-specific configurations
+    : `constructDrv`: Base build helper, the `mkDerivation`-like build helper to extend.
+    : `excludeDrvArgNames`: Argument names not to pass from the input fixed-point arguments to `constructDrv`. Note: It doesn't apply to the updating arguments returned by `extendDrvArgs`.
+    : `extendDrvArgs` : An extension (overlay) of the argument set, like the one taken by [overrideAttrs](#sec-pkg-overrideAttrs) but applied before passing to `constructDrv`.
+    : `inheritFunctionArgs`: Whether to inherit `__functionArgs` from the base build helper (default to `true`).
+    : `transformDrv`: Function to apply to the result derivation (default to `lib.id`).
+
+    # Type
+
+    ```
+    extendMkDerivation ::
+      {
+        constructDrv :: ((FixedPointArgs | AttrSet) -> a)
+        excludeDrvArgNames :: [ String ],
+        extendDrvArgs :: (AttrSet -> AttrSet -> AttrSet)
+        inheritFunctionArgs :: Bool,
+        transformDrv :: a -> a,
+      }
+      -> (FixedPointArgs | AttrSet) -> a
+
+    FixedPointArgs = AttrSet -> AttrSet
+    a = Derivation when defining a build helper
+    ```
+
+    # Examples
+
+    :::{.example}
+    ## `lib.customisation.extendMkDerivation` usage example
+    ```nix-repl
+    mkLocalDerivation = lib.extendMkDerivation {
+      constructDrv = pkgs.stdenv.mkDerivation;
+      excludeDrvArgNames = [ "specialArg" ];
+      extendDrvArgs =
+        finalAttrs: args@{ preferLocalBuild ? true, allowSubstitute ? false, specialArg ? (_: false), ... }:
+        { inherit preferLocalBuild allowSubstitute; passthru = { inherit specialArg; } // args.passthru or { }; };
+    }
+
+    mkLocalDerivation.__functionArgs
+    => { allowSubstitute = true; preferLocalBuild = true; specialArg = true; }
+
+    mkLocalDerivation { inherit (pkgs.hello) pname version src; specialArg = _: false; }
+    => «derivation /nix/store/xirl67m60ahg6jmzicx43a81g635g8z8-hello-2.12.1.drv»
+
+    mkLocalDerivation (finalAttrs: { inherit (pkgs.hello) pname version src; specialArg = _: false; })
+    => «derivation /nix/store/xirl67m60ahg6jmzicx43a81g635g8z8-hello-2.12.1.drv»
+
+    (mkLocalDerivation (finalAttrs: { inherit (pkgs.hello) pname version src; passthru = { foo = "a"; bar = "${finalAttrs.passthru.foo}b"; }; })).bar
+    => "ab"
+    ```
+    :::
+
+    :::{.note}
+    If `transformDrv` is specified,
+    it should take care of existing attributes that perform overriding
+    (e.g., [`overrideAttrs`](#sec-pkg-overrideAttrs))
+    to ensure that the overriding functionality of the result derivation
+    work as expected.
+    Modifications that breaks the overriding include
+    direct [attribute set update](https://nixos.org/manual/nix/stable/language/operators#update)
+    and [`lib.extendDerivation`](#function-library-lib.customisation.extendDerivation).
+    :::
+  */
+  extendMkDerivation =
+    let
+      extendsWithExclusion =
+        excludedNames: g: f: final:
+        let
+          previous = f final;
+        in
+        removeAttrs previous excludedNames // g final previous;
+    in
+    {
+      constructDrv,
+      excludeDrvArgNames ? [ ],
+      extendDrvArgs,
+      inheritFunctionArgs ? true,
+      transformDrv ? id,
+    }:
+    setFunctionArgs
+      # Adds the fixed-point style support
+      (
+        fpargs:
+        transformDrv (
+          constructDrv (extendsWithExclusion excludeDrvArgNames extendDrvArgs (toFunction fpargs))
+        )
+      )
+      # Add __functionArgs
+      (
+        # Inherit the __functionArgs from the base build helper
+        optionalAttrs inheritFunctionArgs (removeAttrs (functionArgs constructDrv) excludeDrvArgNames)
+        # Recover the __functionArgs from the derived build helper
+        // functionArgs (extendDrvArgs { })
+      )
+    // {
+      inherit
+        # Expose to the result build helper.
+        constructDrv
+        excludeDrvArgNames
+        extendDrvArgs
+        transformDrv
+        ;
+    };
 }
diff --git a/lib/default.nix b/lib/default.nix
index 849d6ce8f9d6..f931524002f2 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -121,7 +121,8 @@ let
       noDepEntry fullDepEntry packEntry stringAfter;
     inherit (self.customisation) overrideDerivation makeOverridable
       callPackageWith callPackagesWith extendDerivation hydraJob
-      makeScope makeScopeWithSplicing makeScopeWithSplicing';
+      makeScope makeScopeWithSplicing makeScopeWithSplicing'
+      extendMkDerivation;
     inherit (self.derivations) lazyDerivation optionalDrvAttr warnOnInstantiate;
     inherit (self.meta) addMetaAttrs dontDistribute setName updateName
       appendToName mapDerivationAttrset setPrio lowPrio lowPrioSet hiPrio
diff --git a/pkgs/test/overriding.nix b/pkgs/test/overriding.nix
index 9ccc84e03356..27f57ab1ab9d 100644
--- a/pkgs/test/overriding.nix
+++ b/pkgs/test/overriding.nix
@@ -5,7 +5,7 @@
 }:
 
 let
-  tests = tests-stdenv // tests-go // tests-python;
+  tests = tests-stdenv // test-extendMkDerivation // tests-go // tests-python;
 
   tests-stdenv =
     let
@@ -64,6 +64,73 @@ let
       };
     };
 
+  test-extendMkDerivation =
+    let
+      mkLocalDerivation = lib.extendMkDerivation {
+        constructDrv = pkgs.stdenv.mkDerivation;
+        excludeDrvArgNames = [
+          "specialArg"
+        ];
+        extendDrvArgs =
+          finalAttrs:
+          {
+            preferLocalBuild ? true,
+            allowSubstitute ? false,
+            specialArg ? (_: false),
+            ...
+          }@args:
+          {
+            inherit preferLocalBuild allowSubstitute;
+            passthru = args.passthru or { } // {
+              greeting = if specialArg "Hi!" then "Hi!" else "Hello!";
+            };
+          };
+      };
+
+      helloLocalPlainAttrs = {
+        inherit (pkgs.hello) pname version src;
+        specialArg = throw "specialArg is broken.";
+      };
+
+      helloLocalPlain = mkLocalDerivation helloLocalPlainAttrs;
+
+      helloLocal = mkLocalDerivation (
+        finalAttrs:
+        helloLocalPlainAttrs
+        // {
+          passthru = pkgs.hello.passthru or { } // {
+            foo = "a";
+            bar = "${finalAttrs.passthru.foo}b";
+          };
+        }
+      );
+
+      hiLocal = mkLocalDerivation (
+        helloLocalPlainAttrs
+        // {
+          specialArg = s: lib.stringLength s == 3;
+        }
+      );
+    in
+    {
+      extendMkDerivation-helloLocal-imp-arguments = {
+        expr = helloLocal.preferLocalBuild;
+        expected = true;
+      };
+      extendMkDerivation-helloLocal-plain-equivalence = {
+        expr = helloLocal.drvPath == helloLocalPlain.drvPath;
+        expected = true;
+      };
+      extendMkDerivation-helloLocal-finalAttrs = {
+        expr = helloLocal.bar == "ab";
+        expected = true;
+      };
+      extendMkDerivation-helloLocal-specialArg = {
+        expr = hiLocal.greeting == "Hi!";
+        expected = true;
+      };
+    };
+
   tests-go =
     let
       pet_0_3_4 = pkgs.buildGoModule rec {
@@ -194,6 +261,7 @@ let
         expected = true;
       };
     };
+
 in
 
 stdenvNoCC.mkDerivation {