feat(back): fluidattacks#1101 globbed attrs

- Collect attr paths using python globs
for initialization performance
- Pass them as attrPaths to evaluator
- Redefine evaluator logic for building outputs from attrs
- Initialization for universe takes 1 second now

Signed-off-by: Daniel Salazar <podany270895@gmail.com>

src/cli/main/cli.py

@@ -5,6 +5,9 @@
 from functools import (
     partial,
 )
+from glob import (
+    glob,
+)
 from hashlib import (
     sha256,
 )
@@ -264,9 +267,16 @@ def _clone_src_cache_refresh(head: str, cache_key: str) -> None:
         shutil.copytree(head, cached)
 
 
+def _get_attr_paths(head: str) -> str:
+    rel_paths: list[str] = glob(f"**/main.nix", root_dir=head, recursive=True)
+    abs_paths: list[str] = [f"/{path}" for path in rel_paths]
+    return json.dumps({"attrs": abs_paths}, separators=(",", ":"))
+
+
 def _nix_build(
     *,
     attr: str,
+    attr_paths: str,
     cache: Optional[List[Dict[str, str]]],
     head: str,
     out: str = "",
@@ -298,6 +308,7 @@ def _nix_build(
         *_if(not NIX_STABLE, "build"),
         *_if(NIX_STABLE, "--argstr", "makesSrc", __MAKES_SRC__),
         *_if(NIX_STABLE, "--argstr", "projectSrc", head),
+        *["--argstr", "attrPaths", attr_paths],
         *_if(NIX_STABLE, "--attr", attr),
         *["--option", "cores", "0"],
         *_if(not NIX_STABLE, "--impure"),
@@ -400,7 +411,7 @@ class Config(NamedTuple):
     cache: List[Dict[str, str]]
 
 
-def _get_config(head: str) -> Config:
+def _get_config(head: str, attr_paths: str) -> Config:
     CON.out()
     CON.rule("Building project configuration")
     CON.out()
@@ -410,6 +421,7 @@ def _get_config(head: str) -> Config:
             attr="config.configAsJson"
             if NIX_STABLE
             else f'{head}#__makes__."config:configAsJson"',
+            attr_paths=attr_paths,
             cache=None,
             head=head,
             out=out,
@@ -566,7 +578,8 @@ def cli(args: List[str]) -> None:
         _help_and_exit_base()
 
     head: str = _get_head(src)
-    config: Config = _get_config(head)
+    attr_paths: str = _get_attr_paths(head)
+    config: Config = _get_config(head, attr_paths)
 
     args, attr = _cli_get_args_and_attr(args, config.attrs, src)
 
@@ -575,7 +588,7 @@ def cli(args: List[str]) -> None:
         MAKES_DIR,
         f"provenance{attr.replace('/', '-')}.json",
     )
-    code = _cli_build(attr, config, head, out, src)
+    code = _cli_build(attr, attr_paths, config, head, out, src)
 
     if code == 0:
         write_provenance(args, head, out, provenance, src)
@@ -601,8 +614,9 @@ def _cli_get_args_and_attr(
     return args, attr
 
 
-def _cli_build(
+def _cli_build(  # pylint: disable=too-many-arguments
     attr: str,
+    attr_paths: str,
     config: Config,
     head: str,
     out: str,
@@ -623,6 +637,7 @@ def _cli_build(
             attr=f'config.outputs."{attr}"'
             if NIX_STABLE
             else f'{head}#__makes__."config:outputs:{attr}"',
+            attr_paths=attr_paths,
             cache=config.cache,
             head=head,
             out=out,

src/evaluator/default.nix

@@ -10,6 +10,8 @@
 # You better avoid changing this function signature...
 # Ask a maintainer first.
 {
+  # JSON String containing complete list of main.nix files found within projectSrc
+  attrPaths,
   # flake inputs to inject, if any
   flakeInputs ? {},
   # Source code of makes, can be overriden by the user.
@@ -56,7 +58,7 @@
       "${makesSrcOverriden}/src/evaluator/modules/default.nix"
       makesNix
     ];
-    specialArgs = args;
+    specialArgs = args // {inherit attrPaths;};
   };
 in
   result

src/evaluator/modules/default.nix

@@ -1,8 +1,9 @@
 {
-  attrsMapToList,
+  attrsOptional,
+  attrPaths,
   config,
+  fromJson,
   lib,
-  attrsOptional,
   projectPath,
   projectSrc,
   toFileJson,
@@ -91,34 +92,37 @@
     configAsJson = toFileJson "config.json" config.config;
     outputs = let
       # Load an attr set distributed across many files and directories
-      attrsFromPath = path: position:
+      attrs = let
+        pathInConfig = dir:
+          builtins.any
+          (makesDir: lib.hasInfix makesDir dir)
+          config.extendingMakesDirs;
+        pathExists = dir: builtins.pathExists (projectSrc + dir);
+        attrName = dir: let
+          makesDirsNoRoot = lib.remove "/" config.extendingMakesDirs;
+          from = makesDirsNoRoot ++ ["/main.nix"];
+          to = builtins.genList (_: "") ((builtins.length makesDirsNoRoot) + 1);
+          replaced = builtins.replaceStrings from to dir;
+        in
+          if replaced != ""
+          then replaced
+          else "/";
+      in
         builtins.foldl'
         lib.mergeAttrs
         {}
-        (lib.lists.flatten
-          (attrsMapToList
-            (name: type:
-              if type == "directory"
-              then attrsFromPath "${path}/${name}" (position ++ [name])
-              else if name == "main.nix"
-              then {
-                "/${builtins.concatStringsSep "/" position}" =
-                  import "${path}/main.nix" args;
-              }
-              else {})
-            (builtins.readDir path)));
+        (
+          builtins.map
+          (
+            dir:
+              attrsOptional
+              (pathInConfig dir && pathExists dir)
+              {"${attrName dir}" = import (projectSrc + dir) args;}
+          )
+          (fromJson attrPaths).attrs
+        );
     in
-      (
-        builtins.foldl'
-        lib.mergeAttrs
-        {}
-        (builtins.map
-          (dir:
-            attrsOptional
-            (builtins.pathExists (projectSrc + dir))
-            (attrsFromPath (projectPath dir) []))
-          config.extendingMakesDirs)
-      )
+      attrs
       // {
         __all__ =
           toFileJson "all"