Refactor a bit the functions computing directories contents

Signed-off-by: Jeremie Dimino <jeremie@dimino.org>

src/gen_rules.ml

@@ -237,16 +237,28 @@ module Gen(P : Install_rules.Params) = struct
       file_dst)
 
   (* +-----------------------------------------------------------------+
-     | "text" file listing                                             |
+     | Directory contents                                              |
      +-----------------------------------------------------------------+ *)
 
-  (* Compute the list of "text" files (.ml, .c, ...). This is the list
-     of source files + user generated ones. As a side-effect, setup
-     user rules and copy_files rules. *)
-  let text_files =
-    let cache = Hashtbl.create 32 in
-    fun ~dir ->
-      Hashtbl.find_or_add cache dir ~f:(fun dir ->
+  module Dir_contents : sig
+    type t =
+      { (* Set of "text" files (.ml, .c, ...). This is
+           the set of source files + user generated ones. *)
+        text_files : String.Set.t
+      ; modules : Module.t Module.Name.Map.t Lazy.t
+      ; mlds : string String.Map.t Lazy.t
+      }
+
+    val get : dir:Path.t -> t
+  end = struct
+    type t =
+      { text_files : String.Set.t
+      ; modules : Module.t Module.Name.Map.t Lazy.t
+      ; mlds : string String.Map.t Lazy.t
+      }
+
+    (* As a side-effect, setup user rules and copy_files rules. *)
+    let load_text_files ~dir =
       match Path.Map.find stanzas_per_dir dir with
       | None -> String.Set.empty
       | Some { stanzas; src_dir; scope; _ } ->
@@ -273,13 +285,9 @@ module Gen(P : Install_rules.Params) = struct
           |> String.Set.of_list
         in
         String.Set.union generated_files
-            (SC.source_files sctx ~src_path:src_dir))
+          (SC.source_files sctx ~src_path:src_dir)
 
-  (* +-----------------------------------------------------------------+
-     | Modules listing                                                 |
-     +-----------------------------------------------------------------+ *)
-
-  let guess_modules ~dir ~files =
+    let extract_modules ~dir ~files =
       let make_module syntax base fn =
         (Module.Name.of_string base,
          Module.File.make syntax (Path.relative dir fn))
@@ -314,35 +322,30 @@ module Gen(P : Install_rules.Params) = struct
       Module.Name.Map.merge impls intfs ~f:(fun name impl intf ->
         Some (Module.make name ?impl ?intf))
 
-  let guess_mlds ~files =
-    String.Set.to_list files
-    |> List.filter_map ~f:(fun fn ->
+    let extract_mlds ~files =
+      String.Set.fold files ~init:String.Map.empty ~f:(fun fn acc ->
         match String.lsplit2 fn ~on:'.' with
-      | Some (s, "mld") -> Some (s, fn)
-      | _ -> None)
-    |> String.Map.of_list_exn
+        | Some (s, "mld") -> String.Map.add acc s fn
+        | _ -> acc)
 
-  let mlds_by_dir =
+    let get =
       let cache = Hashtbl.create 32 in
       fun ~dir ->
         Hashtbl.find_or_add cache dir ~f:(fun dir ->
-        let files = text_files ~dir in
-        guess_mlds ~files)
+          let files = load_text_files ~dir in
+          { text_files = files
+          ; modules = lazy (extract_modules ~dir ~files)
+          ; mlds = lazy (extract_mlds ~files)
+          })
+  end
 
   let mlds_of_dir (doc : Documentation.t) ~dir =
     parse_mlds ~dir
-      ~all_mlds:(mlds_by_dir ~dir)
+      ~all_mlds:(Lazy.force (Dir_contents.get ~dir).mlds)
       ~mlds_written_by_user:doc.mld_files
     |> String.Map.values
     |> List.map ~f:(Path.relative dir)
 
-  let modules_by_dir =
-    let cache = Hashtbl.create 32 in
-    fun ~dir ->
-      Hashtbl.find_or_add cache dir ~f:(fun dir ->
-        let files = text_files ~dir in
-        guess_modules ~dir ~files)
-
   type modules_by_lib =
     { modules          : Module.t Module.Name.Map.t
     ; alias_module     : Module.t option
@@ -353,7 +356,7 @@ module Gen(P : Install_rules.Params) = struct
     let cache = Hashtbl.create 32 in
     fun (lib : Library.t) ~dir ->
       Hashtbl.find_or_add cache (dir, lib.name) ~f:(fun _ ->
-        let all_modules = modules_by_dir ~dir in
+        let all_modules = Lazy.force (Dir_contents.get ~dir).modules in
         let modules =
           parse_modules ~all_modules ~buildable:lib.buildable
         in
@@ -1026,8 +1029,12 @@ module Gen(P : Install_rules.Params) = struct
 
   let gen_rules { SC.Dir_with_jbuild. src_dir; ctx_dir; stanzas; scope; kind } =
     (* This interprets "rule" and "copy_files" stanzas. *)
-    let files = text_files ~dir:ctx_dir in
-    let all_modules = modules_by_dir ~dir:ctx_dir in
+    let { Dir_contents.
+          text_files = files
+        ; modules = lazy all_modules
+        ; _
+        } = Dir_contents.get ~dir:ctx_dir
+    in
     let modules_partitioner = Modules_partitioner.create ~dir_kind:kind in
     let merlins =
       List.filter_map stanzas ~f:(fun stanza ->