From 610cc3b0c40f1022829cd78384abdbdd5eecdda5 Mon Sep 17 00:00:00 2001
From: Archit Gupta <archit@accelbread.com>
Date: Sat, 24 Feb 2024 21:22:00 -0800
Subject: [PATCH] Refactor checks to use types

---
 builtinModules/checks.nix | 41 +++++++++++++++++++++++++++++----------
 default.nix               | 26 ++++++++++++++++++++++---
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/builtinModules/checks.nix b/builtinModules/checks.nix
index 61761ca..4c2afd5 100644
--- a/builtinModules/checks.nix
+++ b/builtinModules/checks.nix
@@ -4,28 +4,49 @@
 
 { config, src, lib, flakelight, genSystems, ... }:
 let
-  inherit (lib) isDerivation isFunction mkOption mkIf mapAttrs;
-  inherit (lib.types) lazyAttrsOf raw;
-  inherit (flakelight.types) nullable optFunctionTo;
+  inherit (lib) isFunction last mapAttrs mergeDefinitions mkIf mkOption
+    mkOptionType;
+  inherit (lib.types) lazyAttrsOf optionDescriptionPhrase;
+  inherit (flakelight.types) coercedTo' drv nullable optFunctionTo stringLike;
 
-  mkCheck = pkgs: name: cmd:
-    let cmd' = if isFunction cmd then cmd pkgs else cmd; in
-    if isDerivation cmd' then cmd' else
+  mkCheck = name: pkgs: cmd:
     pkgs.runCommand "check-${name}" { } ''
       cp --no-preserve=mode -r ${src} src
       cd src
-      ${cmd'}
+      ${cmd}
       touch $out
     '';
+
+  checkType = mkOptionType {
+    name = "checkType";
+    description =
+      let
+        targetDesc = optionDescriptionPhrase
+          (class: class == "noun" || class == "composite")
+          (coercedTo' stringLike (abort "") drv);
+      in
+      "${targetDesc} or function that evaluates to it";
+    descriptionClass = "composite";
+    check = x: isFunction x || drv.check x || stringLike.check x;
+    merge = loc: defs: pkgs:
+      let
+        targetType = coercedTo' stringLike (mkCheck (last loc) pkgs) drv;
+      in
+      (mergeDefinitions loc targetType (map
+        (fn: {
+          inherit (fn) file;
+          value = if isFunction fn.value then fn.value pkgs else fn.value;
+        })
+        defs)).mergedValue;
+  };
 in
 {
   options.checks = mkOption {
-    type = nullable (optFunctionTo (lazyAttrsOf raw));
+    type = nullable (optFunctionTo (lazyAttrsOf checkType));
     default = null;
   };
 
   config.outputs = mkIf (config.checks != null) {
-    checks = genSystems (pkgs:
-      mapAttrs (mkCheck pkgs) (config.checks pkgs));
+    checks = genSystems (pkgs: mapAttrs (_: v: v pkgs) (config.checks pkgs));
   };
 }
diff --git a/default.nix b/default.nix
index a877776..fd1eeef 100644
--- a/default.nix
+++ b/default.nix
@@ -7,9 +7,9 @@ let
   inherit (inputs) nixpkgs;
   inherit (builtins) isAttrs isPath readDir;
   inherit (nixpkgs.lib) all attrNames composeManyExtensions evalModules filter
-    findFirst fix genAttrs getValues hasSuffix isFunction isList isStringLike
-    mapAttrs mapAttrsToList mkDefault mkOptionType pathExists pipe removePrefix
-    removeSuffix singleton warn;
+    findFirst fix genAttrs getValues hasSuffix isDerivation isFunction isList
+    isStringLike mapAttrs mapAttrsToList mkDefault mkOptionType pathExists pipe
+    removePrefix removeSuffix singleton warn;
   inherit (nixpkgs.lib.types) coercedTo defaultFunctor functionTo listOf
     optionDescriptionPhrase;
   inherit (nixpkgs.lib.options) mergeEqualOption mergeOneOption;
@@ -74,6 +74,14 @@ let
       merge = mergeOneOption;
     };
 
+    drv = mkOptionType {
+      name = "drv";
+      description = "derivation";
+      descriptionClass = "noun";
+      check = isDerivation;
+      merge = mergeOneOption;
+    };
+
     stringLike = mkOptionType {
       name = "stringLike";
       description = "string-convertible value";
@@ -99,6 +107,18 @@ let
 
     optListOf = elemType: coercedTo elemType singleton (listOf elemType);
 
+    coercedTo' = coercedType: coerceFunc: finalType:
+      (coercedTo coercedType coerceFunc finalType) // {
+        merge = loc: defs:
+          let
+            coerceVal = val:
+              if finalType.check val then val
+              else coerceFunc val;
+          in
+          finalType.merge loc
+            (map (def: def // { value = coerceVal def.value; }) defs);
+      };
+
     optFunctionTo =
       let
         nonFunction = mkOptionType {