lib.hm.dag: performance improvements

Avoid repeated attribute lookups and unnecessary intermediate allocations in
DAG helpers. Cache DAG names while normalizing before-edges to improve topoSort
performance on large DAGs.

Co-authored-by: Eman Resu <78693624+quatquatt@users.noreply.github.com>

modules/lib/dag.nix

@@ -12,20 +12,32 @@
 let
   inherit (lib)
     all
-    filterAttrs
+    attrNames
+    attrValues
+    elem
+    filter
     head
     hm
+    isAttrs
     mapAttrs
     length
     tail
     toposort
     ;
+
+  inherit (hm.dag)
+    empty
+    entriesBetween
+    entryAfter
+    entryBetween
+    isEntry
+    ;
 in
 {
   empty = { };
 
   isEntry = e: e ? data && e ? after && e ? before;
-  isDag = dag: builtins.isAttrs dag && all hm.dag.isEntry (builtins.attrValues dag);
+  isDag = dag: isAttrs dag && all isEntry (attrValues dag);
 
   # Takes an attribute set containing entries built by entryAnywhere,
   # entryAfter, and entryBefore to a topologically sorted list of
@@ -84,16 +96,19 @@ in
   #              }
   #    true
   topoSort =
+    let
+      before = a: b: elem a.name b.after;
+    in
     dag:
     let
-      dagBefore = dag: name: builtins.attrNames (filterAttrs (_n: v: builtins.elem name v.before) dag);
+      names = attrNames dag;
+      dagBefore = name: filter (n: elem name dag.${n}.before) names;
       normalizedDag = mapAttrs (n: v: {
         name = n;
         inherit (v) data;
-        after = v.after ++ dagBefore dag n;
+        after = v.after ++ dagBefore n;
       }) dag;
-      before = a: b: builtins.elem a.name b.after;
-      sorted = toposort before (builtins.attrValues normalizedDag);
+      sorted = toposort before (attrValues normalizedDag);
     in
     if sorted ? result then
       {
@@ -108,10 +123,10 @@ in
   entryBetween = before: after: data: { inherit data before after; };
 
   # Create a DAG entry with no particular dependency information.
-  entryAnywhere = hm.dag.entryBetween [ ] [ ];
+  entryAnywhere = entryBetween [ ] [ ];
 
-  entryAfter = hm.dag.entryBetween [ ];
-  entryBefore = before: hm.dag.entryBetween before [ ];
+  entryAfter = entryBetween [ ];
+  entryBefore = before: entryBetween before [ ];
 
   # Given a list of entries, this function places them in order within the DAG.
   # Each entry is labeled "${tag}-${entry index}" and other DAG entries can be
@@ -124,24 +139,25 @@ in
     let
       go =
         i: before: after: entries:
-        let
-          name = "${tag}-${toString i}";
-        in
         if entries == [ ] then
-          hm.dag.empty
-        else if length entries == 1 then
-          {
-            "${name}" = hm.dag.entryBetween before after (head entries);
-          }
+          empty
         else
-          {
-            "${name}" = hm.dag.entryAfter after (head entries);
-          }
-          // go (i + 1) before [ name ] (tail entries);
+          let
+            name = "${tag}-${toString i}";
+          in
+          if length entries == 1 then
+            {
+              "${name}" = entryBetween before after (head entries);
+            }
+          else
+            {
+              "${name}" = entryAfter after (head entries);
+            }
+            // go (i + 1) before [ name ] (tail entries);
     in
     go 0;
 
-  entriesAnywhere = tag: hm.dag.entriesBetween tag [ ] [ ];
-  entriesAfter = tag: hm.dag.entriesBetween tag [ ];
-  entriesBefore = tag: before: hm.dag.entriesBetween tag before [ ];
+  entriesAnywhere = tag: entriesBetween tag [ ] [ ];
+  entriesAfter = tag: entriesBetween tag [ ];
+  entriesBefore = tag: before: entriesBetween tag before [ ];
 }