From ff24d5d4330db69a4dc4e6fdb97ca3ce9f8881a2 Mon Sep 17 00:00:00 2001
From: Archit Gupta <archit@accelbread.com>
Date: Mon, 6 Nov 2023 23:13:47 -0800
Subject: [PATCH] Allow bundlers to be functions that take pkgs

In addition to a regular bundler of the form `x: x`, this allows setting
bundler options to a function of the form `pkgs: x: x` which is passed
the package set for the system (for example: `{ hello, ... }: x: hello`
to always return hello). This allows, in particular, an autoloaded
bundler in its own file to access the package set.

This is non-trivial as we must tell `x: x` and `pkgs: x: x` apart.

Fortunately, given some derivation `drv`, `pkgs // drv` is a valid
derivation and a set with attr names matching pkgs. When applying this
to the function, if it returns a derivation, it was of the `x: x` form,
and if it returns a function, it was of the `pkgs: x: x` form.

In order to prevent IFD when evaluating the flake if the bundler is of
form `x: x` and uses IFD, we determine the form and apply pkgs when
applying the bundler instead of during flake evaluation. This is done by
wrapping the bundler.

Of note, we cannot rely on `builtins.functionArgs`, since
`pkgs: { hello, ... }: (x: hello) pkgs` is the same as the inner
function but with args hidden.
---
 builtinModules/bundlers.nix | 15 ++++++++-------
 default.nix                 |  8 ++++++++
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/builtinModules/bundlers.nix b/builtinModules/bundlers.nix
index decc406..d723b6c 100644
--- a/builtinModules/bundlers.nix
+++ b/builtinModules/bundlers.nix
@@ -4,21 +4,22 @@
 
 { config, lib, flakelight, ... }:
 let
-  inherit (lib) mkMerge mkOption mkIf;
-  inherit (lib.types) functionTo lazyAttrsOf nullOr package uniq;
-  inherit (flakelight.types) optFunctionTo;
+  inherit (lib) isFunction mapAttrs mkMerge mkOption mkIf;
+  inherit (lib.types) lazyAttrsOf nullOr;
+  inherit (flakelight.types) function optFunctionTo;
 
-  bundler = uniq (functionTo package);
+  wrapBundler = pkgs: bundler: drv:
+    if isFunction (bundler (pkgs // drv)) then bundler pkgs drv else bundler drv;
 in
 {
   options = {
     bundler = mkOption {
-      type = nullOr bundler;
+      type = nullOr function;
       default = null;
     };
 
     bundlers = mkOption {
-      type = nullOr (optFunctionTo (lazyAttrsOf bundler));
+      type = nullOr (optFunctionTo (lazyAttrsOf function));
       default = { };
     };
   };
@@ -30,7 +31,7 @@ in
 
     (mkIf (config.bundlers != null) {
       perSystem = pkgs: {
-        bundlers = config.bundlers pkgs;
+        bundlers = mapAttrs (_: wrapBundler pkgs) (config.bundlers pkgs);
       };
     })
   ];
diff --git a/default.nix b/default.nix
index ba4ee8a..34fb9e9 100644
--- a/default.nix
+++ b/default.nix
@@ -61,6 +61,14 @@ let
         merge = mergeEqualOption;
       };
 
+      function = mkOptionType {
+        name = "function";
+        description = "function";
+        descriptionClass = "noun";
+        check = isFunction;
+        merge = mergeOneOption;
+      };
+
       module = mkOptionType {
         name = "module";
         description = "module";