perf: memoise per-(dep_m, ml_kind, cm_kind, is_consumer) raw-refs builder

src/dune_rules/compilation_context.ml

@@ -33,6 +33,75 @@ module Includes = struct
   let empty = Lib_mode.Cm_kind.Map.make_all Command.Args.empty
 end
 
+(* Variant key: [Direct] for the consumer-side row (the module being compiled),
+   [Transitive] for trans-dep rows. Each carries only the fields that identify
+   the row's closure behaviour — [ml_kind] for [Direct], [cm_kind] for
+   [Transitive] — so cache entries collapse across irrelevant dimensions. *)
+module Raw_refs = struct
+  module Key = struct
+    type t =
+      | Direct of
+          { obj_name : Module_name.Unique.t
+          ; ml_kind : Ml_kind.t
+          }
+      | Transitive of
+          { obj_name : Module_name.Unique.t
+          ; cm_kind : Lib_mode.Cm_kind.t
+          }
+
+    let cm_kind_tag : Lib_mode.Cm_kind.t -> int = function
+      | Ocaml Cmi -> 0
+      | Ocaml Cmo -> 1
+      | Ocaml Cmx -> 2
+      | Melange Cmi -> 3
+      | Melange Cmj -> 4
+    ;;
+
+    let ml_kind_tag : Ml_kind.t -> int = function
+      | Intf -> 0
+      | Impl -> 1
+    ;;
+
+    let equal a b =
+      match a, b with
+      | Direct a, Direct b ->
+        Module_name.Unique.equal a.obj_name b.obj_name
+        && ml_kind_tag a.ml_kind = ml_kind_tag b.ml_kind
+      | Transitive a, Transitive b ->
+        Module_name.Unique.equal a.obj_name b.obj_name
+        && cm_kind_tag a.cm_kind = cm_kind_tag b.cm_kind
+      | Direct _, Transitive _ | Transitive _, Direct _ -> false
+    ;;
+
+    let hash = function
+      | Direct { obj_name; ml_kind } -> Poly.hash (0, obj_name, ml_kind_tag ml_kind)
+      | Transitive { obj_name; cm_kind } -> Poly.hash (1, obj_name, cm_kind_tag cm_kind)
+    ;;
+
+    let repr =
+      let open Repr in
+      let obj_name_repr = view string ~to_:Module_name.Unique.to_string in
+      let ml_kind_repr = view string ~to_:Ml_kind.to_string in
+      let cm_kind_repr = abstract Lib_mode.Cm_kind.to_dyn in
+      variant
+        "Raw_refs.Key"
+        [ case "Direct" (pair obj_name_repr ml_kind_repr) ~proj:(function
+            | Direct { obj_name; ml_kind } -> Some (obj_name, ml_kind)
+            | Transitive _ -> None)
+        ; case "Transitive" (pair obj_name_repr cm_kind_repr) ~proj:(function
+            | Transitive { obj_name; cm_kind } -> Some (obj_name, cm_kind)
+            | Direct _ -> None)
+        ]
+    ;;
+
+    let to_dyn = Repr.to_dyn repr
+  end
+
+  type t = (Key.t, Module_name.Set.t Action_builder.t) Table.t
+
+  let create () : t = Table.create (module Key) 64
+end
+
 (* Per-cctx cache of [Lib_flags.L.include_flags] keyed on [(lib_mode, sorted
    kept_libs)]; two compile rules in this cctx sharing those values share one
    [Args.t]. Output also depends on [t.requires_compile] / [t.requires_hidden],
@@ -123,6 +192,7 @@ type t =
   ; ocaml : Ocaml_toolchain.t
   ; for_ : Compilation_mode.t
   ; filtered_includes : Filtered_includes.t
+  ; raw_refs : Raw_refs.t
   }
 
 let loc t = t.loc
@@ -387,11 +457,21 @@ let create
   ; instances
   ; for_
   ; filtered_includes = Filtered_includes.create ()
+  ; raw_refs = Raw_refs.create ()
   }
 ;;
 
 let for_ t = t.for_
 
+let cached_raw_refs t ~key ~compute =
+  match Table.find t.raw_refs key with
+  | Some builder -> builder
+  | None ->
+    let builder = compute () in
+    Table.set t.raw_refs key builder;
+    builder
+;;
+
 let filtered_include_flags t ~cm_kind ~kept_libs =
   let lib_mode = Lib_mode.of_cm_kind cm_kind in
   let cache_key =

src/dune_rules/compilation_context.mli

@@ -64,6 +64,32 @@ val requires_compile : t -> Lib.t list Resolve.Memo.t
 val parameters : t -> Module_name.t list Resolve.Memo.t
 val includes : t -> Command.Args.without_targets Command.Args.t Lib_mode.Cm_kind.Map.t
 
+module Raw_refs : sig
+  module Key : sig
+    type t =
+      | Direct of
+          { obj_name : Module_name.Unique.t
+          ; ml_kind : Ml_kind.t
+          }
+      | Transitive of
+          { obj_name : Module_name.Unique.t
+          ; cm_kind : Lib_mode.Cm_kind.t
+          }
+  end
+end
+
+(** Memoise the raw-refs [Action_builder.t] computed for each
+    [Raw_refs.Key.t] within this cctx. [compute ()] is invoked
+    only on cache miss; subsequent callers with the same key get
+    the cached builder back. The cache short-circuits before
+    allocating, so siblings sharing [trans_deps] don't redo
+    construction. *)
+val cached_raw_refs
+  :  t
+  -> key:Raw_refs.Key.t
+  -> compute:(unit -> Module_name.Set.t Action_builder.t)
+  -> Module_name.Set.t Action_builder.t
+
 (** Include flags ([-I]/[-H]) for compiling a module against [kept_libs]. The
     cctx's [requires_compile] and [requires_hidden] are each restricted to
     libraries in [kept_libs]; the kept direct entries become [-I], the kept

src/dune_rules/module_compilation.ml

@@ -122,21 +122,33 @@ let lib_deps_for_module ~cctx ~obj_dir ~for_ ~dep_graph ~opaque ~cm_kind ~ml_kin
           | Ocaml (Cmi | Cmo) | Melange _ -> false
       in
       let read_dep_m_raw dep_m ~is_consumer =
-        let* impl_deps =
-          if need_impl_deps_of dep_m ~is_consumer
-          then
+        let key : Compilation_context.Raw_refs.Key.t =
+          let obj_name = Module.obj_name dep_m in
+          if is_consumer
+          then Direct { obj_name; ml_kind }
+          else Transitive { obj_name; cm_kind }
+        in
+        Compilation_context.cached_raw_refs cctx ~key ~compute:(fun () ->
+          let* impl_deps =
+            if need_impl_deps_of dep_m ~is_consumer
+            then
+              Ocamldep.read_immediate_deps_raw_of
+                ~sandbox
+                ~sctx
+                ~obj_dir
+                ~ml_kind:Impl
+                dep_m
+            else Action_builder.return Module_name.Set.empty
+          in
+          let+ intf_deps =
             Ocamldep.read_immediate_deps_raw_of
               ~sandbox
               ~sctx
               ~obj_dir
-              ~ml_kind:Impl
+              ~ml_kind:Intf
               dep_m
-          else Action_builder.return Module_name.Set.empty
-        in
-        let+ intf_deps =
-          Ocamldep.read_immediate_deps_raw_of ~sandbox ~sctx ~obj_dir ~ml_kind:Intf dep_m
-        in
-        Module_name.Set.union impl_deps intf_deps
+          in
+          Module_name.Set.union impl_deps intf_deps)
       in
       let* m_raw = read_dep_m_raw m ~is_consumer:true in
       let* trans_raw =