lib.fileset.fileFilter: Don't run predicate unnecessarily

Before:

    nix-repl> fileset.trace (fileset.fileFilter (file: builtins.trace file.name false) ./default.nix)
    trace: README.md
    trace: benchmark.sh
    trace: default.nix
    trace: internal.nix
    trace: mock-splitRoot.nix
    trace: tests.sh

After:

    nix-repl> fileset.trace (fileset.fileFilter (file: builtins.trace file.name false) ./default.nix)
    trace: default.nix

2 files changed, 43 insertions, 16 deletions

diff --git a/lib/fileset/internal.nix b/lib/fileset/internal.nix
index a4a230a4798..622e9506679 100644
--- a/lib/fileset/internal.nix
+++ b/lib/fileset/internal.nix
@@ -732,29 +732,36 @@ rec {
   # Type: ({ name, type, ... } -> Bool) -> FileSet -> FileSet
   _fileFilter = predicate: fileset:
     let
-      # Check the predicate for a path and a filesetTree, returning a new filesetTree
-      # Type: Path -> filesetTree -> filesetTree
-      recurse = path: tree:
+      # Check the predicate for a single file
+      # Type: String -> String -> filesetTree
+      fromFile = name: type:
+        if
+          predicate {
+            inherit name type;
+            # To ensure forwards compatibility with more arguments being added in the future,
+            # adding an attribute which can't be deconstructed :)
+            "lib.fileset.fileFilter: The predicate function passed as the first argument must be able to handle extra attributes for future compatibility. If you're using `{ name, file }:`, use `{ name, file, ... }:` instead." = null;
+          }
+        then
+          type
+        else
+          null;
+
+      # Check the predicate for all files in a directory
+      # Type: Path -> filesetTree
+      fromDir = path: tree:
         mapAttrs (name: subtree:
           if isAttrs subtree || subtree == "directory" then
-            recurse (path + "/${name}") subtree
-          else if
-            predicate {
-              inherit name;
-              type = subtree;
-              # To ensure forwards compatibility with more arguments being added in the future,
-              # adding an attribute which can't be deconstructed :)
-              "lib.fileset.fileFilter: The predicate function passed as the first argument must be able to handle extra attributes for future compatibility. If you're using `{ name, file }:`, use `{ name, file, ... }:` instead." = null;
-            }
-          then
-            subtree
-          else
+            fromDir (path + "/${name}") subtree
+          else if subtree == null then
             null
+          else
+            fromFile name subtree
         ) (_directoryEntries path tree);
     in
     if fileset._internalIsEmptyWithoutBase then
       _emptyWithoutBase
     else
       _create fileset._internalBase
-        (recurse fileset._internalBase fileset._internalTree);
+        (fromDir fileset._internalBase fileset._internalTree);
 }
diff --git a/lib/fileset/tests.sh b/lib/fileset/tests.sh
index 4247c71efcf..4ddfe3a82a4 100755
--- a/lib/fileset/tests.sh
+++ b/lib/fileset/tests.sh
@@ -857,6 +857,26 @@ checkFileset 'union ./c/a (fileFilter (file: assert file.name != "a"; true) ./.)
 # but here we need to use ./c
 checkFileset 'union (fileFilter (file: assert file.name != "a"; true) ./.) ./c'
 
+# Also lazy, the filter isn't called on a filtered out path
+tree=(
+    [a]=1
+    [b]=0
+    [c]=0
+)
+checkFileset 'fileFilter (file: assert file.name != "c"; file.name == "a") (difference ./. ./c)'
+
+# Make sure single files are filtered correctly
+tree=(
+    [a]=1
+    [b]=0
+)
+checkFileset 'fileFilter (file: assert file.name == "a"; true) ./a'
+tree=(
+    [a]=0
+    [b]=0
+)
+checkFileset 'fileFilter (file: assert file.name == "a"; false) ./a'
+
 ## Tracing
 
 # The second trace argument is returned