refactor(readTree): Pass all readTree parameters as function args r/2824

Instead of having a mix of depot-passed args (for the filter) and args
to the readTree function itself, make everything a single attribute
set of arguments passed to the function.

This also makes it a bit easier to extend this in the future.

Change-Id: I633c1fc96026d137b451bb604ef92be32571a0f5
Reviewed-on: https://cl.tvl.fyi/c/depot/+/3498
Tested-by: BuildkiteCI
Reviewed-by: grfn <grfn@gws.fyi>

4 files changed, 53 insertions, 24 deletions

diff --git a/default.nix b/default.nix
index 19c00bb2c4d2..bb9bc703bf2c 100644
--- a/default.nix
+++ b/default.nix
@@ -56,8 +56,10 @@ let
       };
     };
 
-    readTree' = import ./nix/readTree {
-      argsFilter = depotArgsFilter;
+    readDepot = depotArgs: import ./nix/readTree {} {
+      args = depotArgs;
+      path = ./.;
+      filter = depotArgsFilter;
     };
 
   # To determine build targets, we walk through the depot tree and
@@ -91,7 +93,7 @@ let
            (node.meta.targets or []))
     else [];
 
-in fix(self: (readTree' {
+in fix(self: (readDepot {
   depot = self;
 
   # Pass third_party as 'pkgs' (for compatibility with external
@@ -107,7 +109,7 @@ in fix(self: (readTree' {
   # Note that it is intended for exceptional circumstance, such as
   # debugging by bisecting nixpkgs.
   externalArgs = args;
-} ./.) // {
+}) // {
   # Make the path to the depot available for things that might need it
   # (e.g. NixOS module inclusions)
   path = ./.;
diff --git a/nix/readTree/README.md b/nix/readTree/README.md
index 138abbe30583..b56bc944caba 100644
--- a/nix/readTree/README.md
+++ b/nix/readTree/README.md
@@ -70,8 +70,13 @@ the tree as empty nodes (`{}`).
 
 ## Import structure
 
-`readTree` is called with two parameters: The arguments to pass to all imports,
-and the initial path at which to start the traversal.
+`readTree` is called with an argument set containing a few parameters:
+
+* `path`: Initial path at which to start the traversal.
+* `args`: Arguments to pass to all imports.
+* `filter`: (optional) A function to filter the argument set on each
+  import based on the location in the tree. This can be used to, for
+  example, implement a "visibility" system inside of a tree.
 
 The package headers in this repository follow the form `{ pkgs, ... }:` where
 `pkgs` is a fixed-point of the entire package tree (see the `default.nix` at the
diff --git a/nix/readTree/default.nix b/nix/readTree/default.nix
index 633915f78b3d..b105738a4f80 100644
--- a/nix/readTree/default.nix
+++ b/nix/readTree/default.nix
@@ -5,11 +5,17 @@
 # Provides a function to automatically read a a filesystem structure
 # into a Nix attribute set.
 #
-# Optionally accepts an argument `argsFilter` on import, which is a
-# function that receives the current tree location (as a list of
-# strings) and the argument set and can arbitrarily modify it.
-{ argsFilter ? (x: _parts: x)
-, ... }:
+# Called with an attribute set taking the following arguments:
+#
+#   path: Path to a directory from which to start reading the tree.
+#
+#   args: Argument set to pass to each imported file.
+#
+#   filter: Function to filter `args` based on the tree location. This should
+#           be a function of the form `args -> location -> args`, where the
+#           location is a list of strings representing the path components of
+#           the current readTree target. Optional.
+{ ... }:
 
 let
   inherit (builtins)
@@ -53,7 +59,7 @@ let
 
   # The marker is added to every set that was imported directly by
   # readTree.
-  importWithMark = args: path: parts:
+  importWithMark = args: path: parts: filter:
     let
       importedFile = import path;
       pathType = builtins.typeOf importedFile;
@@ -61,7 +67,7 @@ let
         assert assertMsg
           (pathType == "lambda")
           "readTree: trying to import ${toString path}, but it’s a ${pathType}, you need to make it a function like { depot, pkgs, ... }";
-        importedFile (argsFilter (argsWithPath args parts) parts);
+        importedFile (filter (argsWithPath args parts) parts);
     in if (isAttrs imported)
       then imported // (marker parts)
       else imported;
@@ -70,14 +76,14 @@ let
     let res = match "(.*)\\.nix" file;
     in if res == null then null else head res;
 
-  readTree = { args, initPath, rootDir, parts }:
+  readTree = { args, initPath, rootDir, parts, argsFilter }:
     let
       dir = readDirVisible initPath;
       joinChild = c: initPath + ("/" + c);
 
       self = if rootDir
         then { __readTree = []; }
-        else importWithMark args initPath parts;
+        else importWithMark args initPath parts argsFilter;
 
       # Import subdirectories of the current one, unless the special
       # `.skip-subtree` file exists which makes readTree ignore the
@@ -90,6 +96,7 @@ let
       children = if hasAttr ".skip-subtree" dir then [] else map (c: {
         name = c;
         value = readTree {
+          inherit argsFilter;
           args = args;
           initPath = (joinChild c);
           rootDir = false;
@@ -101,16 +108,22 @@ let
       nixFiles = filter (f: f != null) (map nixFileName (attrNames dir));
       nixChildren = map (c: let p = joinChild (c + ".nix"); in {
         name = c;
-        value = importWithMark args p (parts ++ [ c ]);
+        value = importWithMark args p (parts ++ [ c ]) argsFilter;
       }) nixFiles;
     in if dir ? "default.nix"
       then (if isAttrs self then self // (listToAttrs children) else self)
       else (listToAttrs (nixChildren ++ children) // (marker parts));
 
 in {
-  __functor = _: args: initPath: readTree {
-    inherit args initPath;
-    rootDir = true;
-    parts = [];
-  };
+  __functor = _:
+    { path
+    , args
+    , filter ? (x: _parts: x) }:
+      readTree {
+        inherit args;
+        argsFilter = filter;
+        initPath = path;
+        rootDir = true;
+        parts = [];
+      };
 }
diff --git a/nix/readTree/tests/default.nix b/nix/readTree/tests/default.nix
index f3cab2844785..06d4a12dd08f 100644
--- a/nix/readTree/tests/default.nix
+++ b/nix/readTree/tests/default.nix
@@ -8,7 +8,10 @@ let
     assertThrows
     ;
 
-  tree-ex = depot.nix.readTree {} ./test-example;
+  tree-ex = depot.nix.readTree {
+    path = ./test-example;
+    args = {};
+  };
 
   example = it "corresponds to the README example" [
     (assertEq "third_party attrset"
@@ -32,7 +35,10 @@ let
       "roquefort")
   ];
 
-  tree-tl = depot.nix.readTree {} ./test-tree-traversal;
+  tree-tl = depot.nix.readTree {
+    path = ./test-tree-traversal;
+    args = {};
+  };
 
   traversal-logic = it "corresponds to the traversal logic in the README" [
     (assertEq "skip subtree default.nix is read"
@@ -82,7 +88,10 @@ let
   # these each call readTree themselves because the throws have to happen inside assertThrows
   wrong = it "cannot read these files and will complain" [
     (assertThrows "this file is not a function"
-      (depot.nix.readTree {} ./test-wrong-not-a-function).not-a-function)
+      (depot.nix.readTree {
+        path = ./test-wrong-not-a-function;
+        args = {};
+      }).not-a-function)
     # can’t test for that, assertThrows can’t catch this error
     # (assertThrows "this file is a function but doesn’t have dots"
     #   (depot.nix.readTree {} ./test-wrong-no-dots).no-dots-in-function)