From e7b543e7cb511e20552d439d575e713689a0f66e Mon Sep 17 00:00:00 2001
From: Jeremie Dimino <jdimino@janestreet.com>
Date: Wed, 15 Mar 2017 09:10:22 +0000
Subject: [PATCH] Don't scan the tree a second time for files_recursively_in

---
 src/build.ml      | 14 +++-----------
 src/build.mli     |  2 +-
 src/file_tree.ml  |  9 +++++++++
 src/file_tree.mli |  2 ++
 src/gen_rules.ml  |  2 +-
 5 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/src/build.ml b/src/build.ml
index 28d3ea96..1c91c268 100644
--- a/src/build.ml
+++ b/src/build.ml
@@ -109,21 +109,13 @@ let fail ?targets x =
   | None -> Fail x
   | Some l -> Targets l >>> Fail x
 
-let files_recursively_in ~dir =
-  let ctx_dir, src_dir =
+let files_recursively_in ~dir ~file_tree =
+  let prefix_with, dir =
     match Path.extract_build_context_dir dir with
     | None -> (Path.root, dir)
     | Some (ctx_dir, src_dir) -> (ctx_dir, src_dir)
   in
-  let rec loop dir acc =
-    List.fold_left (Path.readdir dir) ~init:acc ~f:(fun acc fn ->
-      let path = Path.relative dir fn in
-      if Path.is_directory path then
-        loop path acc
-      else
-        Pset.add (Path.append ctx_dir path) acc)
-  in
-  path_set (loop src_dir Pset.empty)
+  path_set (File_tree.files_recursively_in file_tree dir ~prefix_with)
 
 let store_vfile spec = Store_vfile spec
 
diff --git a/src/build.mli b/src/build.mli
index 396851be..cc454e8d 100644
--- a/src/build.mli
+++ b/src/build.mli
@@ -37,7 +37,7 @@ val path  : Path.t      -> ('a, 'a) t
 val paths : Path.t list -> ('a, 'a) t
 val path_set : Path.Set.t -> ('a, 'a) t
 val paths_glob : dir:Path.t -> Re.re -> ('a, 'a) t
-val files_recursively_in : dir:Path.t -> ('a, 'a) t
+val files_recursively_in : dir:Path.t -> file_tree:File_tree.t -> ('a, 'a) t
 val vpath : 'a Vspec.t  -> (unit, 'a) t
 
 val dyn_paths : ('a, Path.t list) t -> ('a, 'a) t
diff --git a/src/file_tree.ml b/src/file_tree.ml
index ca03e69a..296da89f 100644
--- a/src/file_tree.ml
+++ b/src/file_tree.ml
@@ -74,3 +74,12 @@ let file_exists t path fn =
 let exists t path =
   Path.Map.mem path t.dirs ||
   file_exists t (Path.parent path) (Path.basename path)
+
+let files_recursively_in t ?(prefix_with=Path.root) path =
+  match find_dir t path with
+  | None -> Path.Set.empty
+  | Some dir ->
+    Dir.fold dir ~init:Path.Set.empty ~f:(fun dir acc ->
+      let path = Path.append prefix_with (Dir.path dir) in
+      String_set.fold (Dir.files dir) ~init:acc ~f:(fun fn acc ->
+        Path.Set.add (Path.relative path fn) acc))
diff --git a/src/file_tree.mli b/src/file_tree.mli
index 96b61f55..586a94d1 100644
--- a/src/file_tree.mli
+++ b/src/file_tree.mli
@@ -21,3 +21,5 @@ val find_dir : t -> Path.t -> Dir.t option
 
 val exists : t -> Path.t -> bool
 val file_exists : t -> Path.t -> string -> bool
+
+val files_recursively_in : t -> ?prefix_with:Path.t -> Path.t -> Path.Set.t
diff --git a/src/gen_rules.ml b/src/gen_rules.ml
index 32c288e7..c7c3c169 100644
--- a/src/gen_rules.ml
+++ b/src/gen_rules.ml
@@ -425,7 +425,7 @@ module Gen(P : Params) = struct
         end
       | Files_recursively_in s ->
         let path = Path.relative dir (expand_vars ~dir s) in
-        Build.files_recursively_in ~dir:path
+        Build.files_recursively_in ~dir:path ~file_tree:P.file_tree
 
     let dep_of_list ~dir ts =
       let rec loop acc = function