fix: allow reolving binaries to original paths

Previously, we'd always binaries from the install context
(_build/install/$context/bin). This would unnecessarily load the install
rules.

Now, we add an argument that allows us to resolve to the original paths.
Reducing the amount of rules that need to be loaded.

We use this argument in a few cases where we don't need to build the
path in _build/install

<!-- ps-id: cb6fc8f6-5f50-439b-b542-6b933c656af0 -->

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

otherlibs/stdune/src/map.ml

@@ -44,6 +44,12 @@ module Make (Key : Key) : S with type key = Key.t = struct
   let merge a b ~f = merge a b ~f
   let union a b ~f = union a b ~f
 
+  let union_all maps ~f =
+    match maps with
+    | [] -> empty
+    | init :: maps -> List.fold_left maps ~init ~f:(fun acc map -> union acc map ~f)
+  ;;
+
   let union_exn a b =
     union a b ~f:(fun key _ _ ->
       Code_error.raise

otherlibs/stdune/src/map_intf.ml

@@ -20,6 +20,7 @@ module type S = sig
   val add_multi : 'a list t -> key -> 'a -> 'a list t
   val merge : 'a t -> 'b t -> f:(key -> 'a option -> 'b option -> 'c option) -> 'c t
   val union : 'a t -> 'a t -> f:(key -> 'a -> 'a -> 'a option) -> 'a t
+  val union_all : 'a t list -> f:(key -> 'a -> 'a -> 'a option) -> 'a t
 
   (** Like [union] but raises a code error if a key appears in both maps. *)
   val union_exn : 'a t -> 'a t -> 'a t

src/dune_rules/artifacts.ml

@@ -4,49 +4,84 @@ open Memo.O
 let bin_dir_basename = ".bin"
 let local_bin p = Path.Build.relative p bin_dir_basename
 
+type origin =
+  { binding : File_binding.Unexpanded.t
+  ; dir : Path.Build.t
+  ; dst : Path.Local.t
+  }
+
+type where =
+  | Install_dir
+  | Original_path
+
+type path =
+  | Resolved of Path.Build.t
+  | Origin of origin
+
+type local_bins = path Filename.Map.t
+
 type t =
   { context : Context.t
   ; (* Mapping from executable names to their actual path in the workspace.
        The keys are the executable names without the .exe, even on Windows.
        Enumerating binaries from install stanzas may involve expanding globs,
        but the artifacts database is depended on by the logic which expands
        globs. The computation of this field is deferred to break the cycle. *)
-    local_bins : Path.Build.t Filename.Map.t Memo.Lazy.t
+    local_bins : local_bins Memo.Lazy.t
   }
 
 let force { local_bins; _ } =
-  let+ (_ : Path.Build.t Filename.Map.t) = Memo.Lazy.force local_bins in
+  let+ (_ : local_bins) = Memo.Lazy.force local_bins in
   ()
 ;;
 
+let expand = Fdecl.create Dyn.opaque
+
 let analyze_binary t name =
   match Filename.is_relative name with
-  | false -> Memo.return (Some (Path.of_filename_relative_to_initial_cwd name))
+  | false -> Memo.return (`Resolved (Path.of_filename_relative_to_initial_cwd name))
   | true ->
     let* local_bins = Memo.Lazy.force t.local_bins in
     (match Filename.Map.find local_bins name with
-     | Some path -> Memo.return (Some (Path.build path))
-     | None -> Context.which t.context name)
+     | Some (Resolved p) -> Memo.return (`Resolved (Path.build p))
+     | Some (Origin o) -> Memo.return (`Origin o)
+     | None ->
+       Context.which t.context name
+       >>| (function
+        | None -> `None
+        | Some path -> `Resolved path))
 ;;
 
-let binary t ?hint ~loc name =
+let binary t ?hint ?(where = Install_dir) ~loc name =
   analyze_binary t name
-  >>| function
-  | Some path -> Ok path
-  | None ->
+  >>= function
+  | `Resolved path -> Memo.return @@ Ok path
+  | `None ->
     let context = Context.name t.context in
-    Error (Action.Prog.Not_found.create ~program:name ?hint ~context ~loc ())
+    Memo.return
+    @@ Error (Action.Prog.Not_found.create ~program:name ?hint ~context ~loc ())
+  | `Origin { dir; binding; dst } ->
+    (match where with
+     | Install_dir ->
+       let install_dir = Install.Context.bin_dir ~context:(Context.name t.context) in
+       Memo.return @@ Ok (Path.build @@ Path.Build.append_local install_dir dst)
+     | Original_path ->
+       let+ expanded =
+         File_binding.Unexpanded.expand
+           binding
+           ~dir
+           ~f:(Fdecl.get expand ~context:t.context ~dir)
+       in
+       let src = File_binding.Expanded.src expanded in
+       Ok (Path.build src))
 ;;
 
 let binary_available t name =
-  analyze_binary t name
-  >>= function
-  | None -> Memo.return false
-  | Some path ->
-    (match path with
-     | External e -> Fs_memo.file_exists @@ External e
-     | In_source_tree e -> Fs_memo.file_exists @@ In_source_dir e
-     | In_build_dir _ -> Memo.return true)
+  match Filename.is_relative name with
+  | false -> Memo.return false
+  | true ->
+    let+ local_bins = Memo.Lazy.force t.local_bins in
+    Filename.Map.mem local_bins name
 ;;
 
 let add_binaries t ~dir l =
@@ -55,7 +90,7 @@ let add_binaries t ~dir l =
       let+ local_bins = Memo.Lazy.force t.local_bins in
       List.fold_left l ~init:local_bins ~f:(fun acc fb ->
         let path = File_binding.Expanded.dst_path fb ~dir:(local_bin dir) in
-        Filename.Map.set acc (Path.Build.basename path) path))
+        Filename.Map.set acc (Path.Build.basename path) (Resolved path)))
   in
   { t with local_bins }
 ;;
@@ -70,10 +105,9 @@ let create =
     let local_bins =
       Memo.lazy_ (fun () ->
         let+ local_bins = Memo.Lazy.force local_bins in
-        Path.Build.Set.fold local_bins ~init:Filename.Map.empty ~f:(fun path acc ->
-          let name = Path.Build.basename path in
+        Filename.Map.foldi local_bins ~init:Filename.Map.empty ~f:(fun name origin acc ->
           let key = drop_suffix name in
-          Filename.Map.set acc key path))
+          Filename.Map.set acc key (Origin origin)))
     in
     { context; local_bins }
 ;;

src/dune_rules/artifacts.mli

@@ -2,6 +2,16 @@ open Import
 
 type t
 
+type origin =
+  { binding : File_binding.Unexpanded.t
+  ; dir : Path.Build.t
+  ; dst : Path.Local.t
+  }
+
+type where =
+  | Install_dir
+  | Original_path
+
 (** Force the computation of the internal list of binaries. This is exposed as
     some error checking is only performed during this computation and some
     errors will go unreported unless this computation takes place. *)
@@ -15,8 +25,17 @@ val local_bin : Path.Build.t -> Path.Build.t
 
 (** A named artifact that is looked up in the PATH if not found in the tree If
     the name is an absolute path, it is used as it. *)
-val binary : t -> ?hint:string -> loc:Loc.t option -> string -> Action.Prog.t Memo.t
+val binary
+  :  t
+  -> ?hint:string
+  -> ?where:where
+  -> loc:Loc.t option
+  -> Filename.t
+  -> Action.Prog.t Memo.t
 
 val binary_available : t -> string -> bool Memo.t
 val add_binaries : t -> dir:Path.Build.t -> File_binding.Expanded.t list -> t
-val create : Context.t -> local_bins:Path.Build.Set.t Memo.Lazy.t -> t
+val create : Context.t -> local_bins:origin Filename.Map.t Memo.Lazy.t -> t
+
+val expand
+  : (context:Context.t -> dir:Path.Build.t -> String_with_vars.t -> string Memo.t) Fdecl.t

src/dune_rules/artifacts_db.ml

@@ -35,8 +35,8 @@ let available_exes ~dir (exes : Dune_file.Executables.t) =
 ;;
 
 let get_installed_binaries ~(context : Context.t) stanzas =
+  let merge _ _ x = Some x in
   let open Memo.O in
-  let install_dir = Install.Context.bin_dir ~context:(Context.name context) in
   let expand ~dir sw = Expander.With_reduced_var_set.expand ~context ~dir sw in
   let expand_str ~dir sw = Expander.With_reduced_var_set.expand_str ~context ~dir sw in
   let expand_str_partial ~dir sw =
@@ -56,12 +56,14 @@ let get_installed_binaries ~(context : Context.t) stanzas =
             ~expand:(expand_str ~dir)
             ~expand_partial:(expand_str_partial ~dir)
         in
-        let p = Path.Local.of_string (Install.Entry.Dst.to_string p) in
-        if Path.Local.is_root (Path.Local.parent_exn p)
-        then Some (Path.Build.append_local install_dir p)
+        let dst = Path.Local.of_string (Install.Entry.Dst.to_string p) in
+        if Path.Local.is_root (Path.Local.parent_exn dst)
+        then (
+          let origin = { Artifacts.binding = fb; dir; dst } in
+          Some (Path.Local.basename dst, origin))
         else None)
       >>| List.filter_opt
-      >>| Path.Build.Set.of_list
+      >>| Filename.Map.of_list_reduce ~f:(fun _ y -> y)
     in
     Memo.List.map d.stanzas ~f:(fun stanza ->
       match Stanza.repr stanza with
@@ -79,12 +81,10 @@ let get_installed_binaries ~(context : Context.t) stanzas =
              | false -> Memo.return true
              | true -> available_exes ~dir exes)
         in
-        if available
-        then binaries_from_install files
-        else Memo.return Path.Build.Set.empty
-      | _ -> Memo.return Path.Build.Set.empty)
-    >>| Path.Build.Set.union_all)
-  >>| Path.Build.Set.union_all
+        if available then binaries_from_install files else Memo.return Filename.Map.empty
+      | _ -> Memo.return Filename.Map.empty)
+    >>| Filename.Map.union_all ~f:merge)
+  >>| Filename.Map.union_all ~f:merge
 ;;
 
 let all =

src/dune_rules/coq/coq_rules.ml

@@ -95,6 +95,7 @@ let coqc ~loc ~dir ~sctx =
   Super_context.resolve_program_memo
     sctx
     "coqc"
+    ~where:Original_path
     ~dir
     ~loc:(Some loc)
     ~hint:"opam install coq"

src/dune_rules/expander.ml

@@ -931,3 +931,8 @@ let expand_lock ~base expander (Locks.Lock sw) =
 let expand_locks ~base expander locks =
   Memo.List.map locks ~f:(expand_lock ~base expander) |> Action_builder.of_memo
 ;;
+
+let () =
+  Fdecl.set Artifacts.expand (fun ~context ~dir sw ->
+    With_reduced_var_set.expand_str ~context ~dir sw)
+;;

src/dune_rules/melange/melange_binary.ml

@@ -1,7 +1,13 @@
 open Import
 
 let melc sctx ~loc ~dir =
-  Super_context.resolve_program_memo sctx ~loc ~dir ~hint:"opam install melange" "melc"
+  Super_context.resolve_program_memo
+    sctx
+    ~loc
+    ~dir
+    ~where:Original_path
+    ~hint:"opam install melange"
+    "melc"
 ;;
 
 let where =

src/dune_rules/super_context.ml

@@ -275,13 +275,13 @@ let add_alias_action t alias ~dir ~loc action =
 
 let env_node = Env_tree.get_node
 
-let resolve_program_memo t ~dir ?hint ~loc bin =
+let resolve_program_memo t ~dir ?where ?hint ~loc bin =
   let* artifacts = Env_tree.artifacts_host t ~dir in
-  Artifacts.binary ?hint ~loc artifacts bin
+  Artifacts.binary ?hint ?where ~loc artifacts bin
 ;;
 
-let resolve_program t ~dir ?hint ~loc bin =
-  Action_builder.of_memo @@ resolve_program_memo t ~dir ?hint ~loc bin
+let resolve_program t ~dir ?where ?hint ~loc bin =
+  Action_builder.of_memo @@ resolve_program_memo t ~dir ?where ?hint ~loc bin
 ;;
 
 let add_packages_env context ~base stanzas packages =

src/dune_rules/super_context.mli

@@ -69,6 +69,7 @@ val add_alias_action
 val resolve_program
   :  t
   -> dir:Path.Build.t
+  -> ?where:Artifacts.where
   -> ?hint:string
   -> loc:Loc.t option
   -> string
@@ -78,6 +79,7 @@ val resolve_program
 val resolve_program_memo
   :  t
   -> dir:Path.Build.t
+  -> ?where:Artifacts.where
   -> ?hint:string
   -> loc:Loc.t option
   -> string