fix: performance regression from #8447

Get rid of the slow dune file comparison in #9738

Signed-off-by: Rudi Grinberg <[email protected]>

<!-- ps-id: 6499de7c-e597-44e0-af7f-33bf7bf235ab -->

Author: rgrinberg

bin/describe/describe_pp.ml

@@ -56,7 +56,7 @@ let get_pped_file super_context file =
          >>| Source_tree.Dir.path
          >>| Path.source
        in
-       let* dune_file = Dune_rules.Dune_load.in_dir (dir |> in_build_dir) in
+       let* dune_file = Dune_rules.Only_packages.stanzas_in_dir (dir |> in_build_dir) in
        let staged_pps =
          Option.bind dune_file ~f:(fun dune_file ->
            dune_file.stanzas

src/dune_rules/dune_load.ml

@@ -215,42 +215,26 @@ module Dune_files = struct
 
   type t = one list
 
-  let interpret =
-    let impl (dir, project, dune_file) =
-      let file = Source_tree.Dune_file.path dune_file in
-      let static = Source_tree.Dune_file.get_static_sexp dune_file in
-      match Source_tree.Dune_file.kind dune_file with
-      | Ocaml_script ->
-        Memo.return
-          (Script
-             { script =
-                 { dir
-                 ; project
-                 ; file =
-                     (* we can't introduce ocaml syntax with [(sudir ..)] *)
-                     Option.value_exn file
-                 }
-             ; from_parent = static
-             })
-      | Plain ->
-        let open Memo.O in
-        let+ stanzas = Dune_file.parse static ~dir ~file ~project in
-        Literal stanzas
-    in
-    let module Input = struct
-      type t = Path.Source.t * Dune_project.t * Source_tree.Dune_file.t
-
-      let equal =
-        Tuple.T3.equal Path.Source.equal Dune_project.equal Source_tree.Dune_file.equal
-      ;;
-
-      let hash = Tuple.T3.hash Path.Source.hash Dune_project.hash Poly.hash
-      let to_dyn = Dyn.opaque
-    end
-    in
-    let memo = Memo.create "Dune_files.interpret" ~input:(module Input) impl in
-    fun ~dir ~project ~(dune_file : Source_tree.Dune_file.t) ->
-      Memo.exec memo (dir, project, dune_file)
+  let interpret ~dir project dune_file =
+    let file = Source_tree.Dune_file.path dune_file in
+    let static = Source_tree.Dune_file.get_static_sexp dune_file in
+    match Source_tree.Dune_file.kind dune_file with
+    | Ocaml_script ->
+      Memo.return
+        (Script
+           { script =
+               { dir
+               ; project
+               ; file =
+                   (* we can't introduce ocaml syntax with [(sudir ..)] *)
+                   Option.value_exn file
+               }
+           ; from_parent = static
+           })
+    | Plain ->
+      let open Memo.O in
+      let+ stanzas = Dune_file.parse static ~dir ~file ~project in
+      Literal stanzas
   ;;
 
   let eval dune_files ~context =
@@ -328,7 +312,7 @@ let load () =
   let+ dune_files =
     Appendable_list.to_list dune_files
     |> Memo.parallel_map ~f:(fun (dir, project, dune_file) ->
-      Dune_files.interpret ~dir ~project ~dune_file)
+      Dune_files.interpret ~dir project dune_file)
   in
   { dune_files
   ; packages
@@ -339,34 +323,6 @@ let load () =
   }
 ;;
 
-let in_dir dir =
-  let source_dir = Path.Build.drop_build_context_exn dir in
-  Source_tree.find_dir source_dir
-  >>= function
-  | None -> Memo.return None
-  | Some d ->
-    (match Source_tree.Dir.dune_file d with
-     | None -> Memo.return None
-     | Some dune_file ->
-       let project = Source_tree.Dir.project d in
-       Dune_files.interpret ~dir:source_dir ~project ~dune_file
-       >>= (function
-        | Literal dune_file -> Memo.return (Some dune_file)
-        | Script script ->
-          let context =
-            match Install.Context.of_path dir with
-            | Some c -> c
-            | None ->
-              User_error.raise
-                [ Pp.textf
-                    "no context in directory %s"
-                    (Path.Build.to_string_maybe_quoted dir)
-                ]
-          in
-          let+ dune_file = Script.eval_one ~context script in
-          Some dune_file))
-;;
-
 let load =
   let memo = Memo.lazy_ ~name:"dune_load" load in
   fun () -> Memo.Lazy.force memo

src/dune_rules/dune_load.mli

@@ -14,7 +14,6 @@ end
 
 type t
 
-val in_dir : Path.Build.t -> Dune_file.t option Memo.t
 val dune_files : t -> Dune_files.t
 val packages : t -> Package.t Package.Name.Map.t
 val projects : t -> Dune_project.t list

src/dune_rules/only_packages.ml

@@ -100,20 +100,33 @@ let filter_out_stanzas_from_hidden_packages ~visible_pkgs =
       | _ -> None))
 ;;
 
+type filtered_stanzas =
+  { all : Dune_file.t list
+  ; map : Dune_file.t Path.Source.Map.t
+  }
+
 let filtered_stanzas =
   let db =
     Per_context.create_by_name ~name:"filtered_stanzas"
     @@ fun context ->
     let* only_packages = Memo.Lazy.force conf
     and+ dune_files = Dune_load.load () >>| Dune_load.dune_files in
     let+ stanzas = Dune_load.Dune_files.eval ~context dune_files in
-    match only_packages with
-    | None -> stanzas
-    | Some visible_pkgs ->
-      List.map stanzas ~f:(fun (dir_conf : Dune_file.t) ->
-        { dir_conf with
-          stanzas = filter_out_stanzas_from_hidden_packages ~visible_pkgs dir_conf.stanzas
-        })
+    let all =
+      match only_packages with
+      | None -> stanzas
+      | Some visible_pkgs ->
+        List.map stanzas ~f:(fun (dune_file : Dune_file.t) ->
+          { dune_file with
+            stanzas =
+              filter_out_stanzas_from_hidden_packages ~visible_pkgs dune_file.stanzas
+          })
+    in
+    let map =
+      Path.Source.Map.of_list_map_exn all ~f:(fun (dune_file : Dune_file.t) ->
+        dune_file.dir, dune_file)
+    in
+    { all; map }
   in
   fun ctx -> Staged.unstage db ctx
 ;;
@@ -127,17 +140,16 @@ let get () =
 let stanzas_in_dir dir =
   if Path.Build.is_root dir
   then Memo.return None
-  else
-    Dune_load.in_dir dir
-    >>= function
+  else (
+    match Install.Context.of_path dir with
     | None -> Memo.return None
-    | Some dune_file ->
-      let+ stanzas =
-        Memo.Lazy.force conf
-        >>| function
-        | None -> dune_file.stanzas
-        | Some visible_pkgs ->
-          filter_out_stanzas_from_hidden_packages ~visible_pkgs dune_file.stanzas
-      in
-      Some { dune_file with stanzas }
+    | Some ctx ->
+      let+ filtered_stanzas = filtered_stanzas ctx in
+      let dir = Path.Build.drop_build_context_exn dir in
+      Path.Source.Map.find filtered_stanzas.map dir)
+;;
+
+let filtered_stanzas ctx =
+  let+ filtered_stanzas = filtered_stanzas ctx in
+  filtered_stanzas.all
 ;;