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>

3 files changed, 47 insertions, 20 deletions

diff --git a/nix/readTree/README.md b/nix/readTree/README.md
index 138abbe305..b56bc944ca 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 633915f78b..b105738a4f 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 f3cab28447..06d4a12dd0 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)