perf: run parse_compilation_units once (#7187)

fix: speed up compilation with many modules

Only generate the map used to look up dependencies once per library/executable. Rather than once per module.

Signed-off-by: Javier Chávarri <javier.chavarri@gmail.com>
Signed-off-by: Rudi Grinberg <me@rgrinberg.com>
Co-authored-by: Rudi Grinberg <me@rgrinberg.com>

.github/workflows/bench.yml

@@ -5,7 +5,7 @@ on:
   push:
     branches:
       - main
-      - bench/add-warm-run
+      - perf/run-parse_compilation_units-once
 
 permissions:
   # deployments permission to deploy GitHub pages website

CHANGES.md

@@ -1,6 +1,9 @@
 Unreleased
 ----------
 
+- Speed up rule generation for libraries and executables with many modules
+  (#7187, @jchavarri)
+
 - Do not re-render UI on every frame if the UI doesn't change (#7186, fix
   #7184, @rgrinberg)
 

src/dune_rules/dep_graph.ml

@@ -3,29 +3,29 @@ open Action_builder.O
 
 type t =
   { dir : Path.Build.t
-  ; per_module : Module.t list Action_builder.t Module.Obj_map.t
+  ; per_module : Module.t list Action_builder.t Module_name.Unique.Map.t
   }
 
 let make ~dir ~per_module = { dir; per_module }
 
 let deps_of t (m : Module.t) =
-  match Module.Obj_map.find t.per_module m with
+  match Module_name.Unique.Map.find t.per_module (Module.obj_name m) with
   | Some x -> x
   | None ->
     Code_error.raise "Ocamldep.Dep_graph.deps_of"
       [ ("dir", Path.Build.to_dyn t.dir)
       ; ( "modules"
         , Dyn.(list Module_name.Unique.to_dyn)
-            (Module.Obj_map.keys t.per_module |> List.map ~f:Module.obj_name) )
+            (Module_name.Unique.Map.keys t.per_module) )
       ; ("m", Module.to_dyn m)
       ]
 
 module Top_closure = Top_closure.Make (Module_name.Unique.Set) (Action_builder)
 
 let top_closed t modules =
   let+ res =
-    Top_closure.top_closure modules ~key:Module.obj_name
-      ~deps:(Module.Obj_map.find_exn t.per_module)
+    Top_closure.top_closure modules ~key:Module.obj_name ~deps:(fun m ->
+        Module_name.Unique.Map.find_exn t.per_module (Module.obj_name m))
   in
   match res with
   | Ok modules -> modules
@@ -58,7 +58,9 @@ let top_closed_implementations t modules =
 
 let dummy (m : Module.t) =
   { dir = Path.Build.root
-  ; per_module = Module.Obj_map.singleton m (Action_builder.return [])
+  ; per_module =
+      Module_name.Unique.Map.singleton (Module.obj_name m)
+        (Action_builder.return [])
   }
 
 module Ml_kind = struct

src/dune_rules/dep_graph.mli

@@ -6,7 +6,7 @@ type t
 
 val make :
      dir:Path.Build.t
-  -> per_module:Module.t list Action_builder.t Module.Obj_map.t
+  -> per_module:Module.t list Action_builder.t Module_name.Unique.Map.t
   -> t
 
 val deps_of : t -> Module.t -> Module.t list Action_builder.t

src/dune_rules/dep_rules.ml

@@ -3,11 +3,15 @@ open Memo.O
 open Ocamldep.Modules_data
 
 let transitive_deps_contents modules =
-  List.map modules ~f:(fun m -> Module_name.to_string (Module.name m))
+  List.map modules ~f:(fun m ->
+      (* TODO use object names *)
+      Modules.Sourced_module.to_module m |> Module.name |> Module_name.to_string)
   |> String.concat ~sep:"\n"
 
 let ooi_deps { vimpl; sctx; dir; obj_dir; modules = _; stdlib = _; sandbox = _ }
-    ~dune_version ~vlib_obj_map ~(ml_kind : Ml_kind.t) (m : Module.t) =
+    ~dune_version ~vlib_obj_map ~(ml_kind : Ml_kind.t)
+    (sourced_module : Modules.Sourced_module.t) =
+  let m = Modules.Sourced_module.to_module sourced_module in
   let cm_kind =
     match ml_kind with
     | Intf -> Cm_kind.Cmi
@@ -61,7 +65,9 @@ let deps_of_module ({ modules; _ } as md) ~ml_kind m =
       let+ deps = deps in
       aliases @ deps)
 
-let deps_of_vlib_module ({ obj_dir; vimpl; dir; sctx; _ } as md) ~ml_kind m =
+let deps_of_vlib_module ({ obj_dir; vimpl; dir; sctx; _ } as md) ~ml_kind
+    sourced_module =
+  let m = Modules.Sourced_module.to_module sourced_module in
   let vimpl = Option.value_exn vimpl in
   let vlib = Vimpl.vlib vimpl in
   match Lib.Local.of_lib vlib with
@@ -71,7 +77,11 @@ let deps_of_vlib_module ({ obj_dir; vimpl; dir; sctx; _ } as md) ~ml_kind m =
       let impl = Vimpl.impl vimpl in
       Dune_project.dune_version impl.project
     in
-    ooi_deps md ~dune_version ~vlib_obj_map ~ml_kind m
+    let open Memo.O in
+    let+ deps =
+      ooi_deps md ~dune_version ~vlib_obj_map ~ml_kind sourced_module
+    in
+    Action_builder.map deps ~f:(List.map ~f:Modules.Sourced_module.to_module)
   | Some lib ->
     let modules = Vimpl.vlib_modules vimpl in
     let info = Lib.Local.info lib in
@@ -96,12 +106,13 @@ let rec deps_of md ~ml_kind (m : Modules.Sourced_module.t) =
   in
   if is_alias then Memo.return (Action_builder.return [])
   else
-    let skip_if_source_absent f m =
-      if Module.has m ~ml_kind then f m
+    let skip_if_source_absent f sourced_module =
+      let m = Modules.Sourced_module.to_module m in
+      if Module.has m ~ml_kind then f sourced_module
       else Memo.return (Action_builder.return [])
     in
     match m with
-    | Imported_from_vlib m ->
+    | Imported_from_vlib _ ->
       skip_if_source_absent (deps_of_vlib_module md ~ml_kind) m
     | Normal m -> skip_if_source_absent (deps_of_module md ~ml_kind) m
     | Impl_of_virtual_module impl_or_vlib -> (
@@ -144,6 +155,8 @@ let rules md =
   | None ->
     dict_of_func_concurrently (fun ~ml_kind ->
         let+ per_module =
-          Modules.obj_map_build modules ~f:(deps_of md ~ml_kind)
+          Modules.obj_map modules
+          |> Module_name.Unique.Map_traversals.parallel_map
+               ~f:(fun _obj_name m -> deps_of md ~ml_kind m)
         in
         Dep_graph.make ~dir:md.dir ~per_module)

src/dune_rules/module.ml

@@ -326,8 +326,6 @@ module Obj_map = struct
   end)
 end
 
-module Obj_map_traversals = Memo.Make_map_traversals (Obj_map)
-
 let encode ({ source; obj_name; pp = _; visibility; kind; install_as = _ } as t)
     ~src_dir =
   let open Dune_lang.Encoder in

src/dune_rules/module.mli

@@ -112,13 +112,6 @@ module Obj_map : sig
   val find_exn : 'a t -> module_ -> 'a
 end
 
-module Obj_map_traversals : sig
-  val parallel_iter : 'a Obj_map.t -> f:(t -> 'a -> unit Memo.t) -> unit Memo.t
-
-  val parallel_map :
-    'a Obj_map.t -> f:(t -> 'a -> 'b Memo.t) -> 'b Obj_map.t Memo.t
-end
-
 val sources : t -> Path.t list
 
 val visibility : t -> Visibility.t

src/dune_rules/module_name.ml

@@ -117,6 +117,7 @@ module Unique = struct
 
   module Map = Map
   module Set = Set
+  module Map_traversals = Map_traversals
 end
 
 module Path = struct

src/dune_rules/module_name.mli

@@ -66,6 +66,12 @@ module Unique : sig
   include Dune_lang.Conv.S with type t := t
 
   include Comparable_intf.S with type key := t
+
+  module Map_traversals : sig
+    val parallel_iter : 'a Map.t -> f:(t -> 'a -> unit Memo.t) -> unit Memo.t
+
+    val parallel_map : 'a Map.t -> f:(t -> 'a -> 'b Memo.t) -> 'b Map.t Memo.t
+  end
 end
 
 module Path : sig