WIP: Narrow which in pkg context

src/dune_rules/artifacts.ml

@@ -33,6 +33,8 @@ type t =
        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 : local_bins Memo.Lazy.t
+  ; (* Package names of lock dir packages that are dependencies of the "current dir"'s stanza *)
+    lockdir_deps : Package.Name.Set.t option
   }
 
 let force { local_bins; _ } =
@@ -65,7 +67,7 @@ let analyze_binary t ~dir name =
       match lookup_name with
       | None -> Memo.return `None
       | Some lookup_name ->
-        Context.which t.context lookup_name
+        Context.which t.context ~narrow_to_packages:t.lockdir_deps lookup_name
         >>| (function
          | None -> `None
          | Some path -> `Resolved path)
@@ -151,6 +153,8 @@ let add_binaries t ~dir l =
   { t with local_bins }
 ;;
 
+let set_lockdir_deps t ~lockdir_deps = { t with lockdir_deps }
+
 let create =
   fun (context : Context.t)
     ~(local_bins : origin Appendable_list.t Filename.Map.t Memo.Lazy.t) ->
@@ -162,5 +166,5 @@ let create =
         , Origin (Appendable_list.to_list sources) ))
       |> Filename.Map.of_list_exn)
   in
-  { context; local_bins }
+  { context; local_bins; lockdir_deps = None }
 ;;

src/dune_rules/artifacts.mli

@@ -47,3 +47,5 @@ val create
   :  Context.t
   -> local_bins:origin Appendable_list.t Filename.Map.t Memo.Lazy.t
   -> t
+
+val set_lockdir_deps : t -> lockdir_deps:Package.Name.Set.t option -> t

src/dune_rules/cc_flags.ml

@@ -134,7 +134,7 @@ let resolve_c_compiler context ~dir program =
   | Relative_to_current_dir -> Memo.return (Ok (Path.relative (Path.build dir) program))
   | In_path ->
     let program = Filename.of_string_exn program in
-    Context.which context program
+    Context.which ~narrow_to_packages:None context program
     >>| (function
      | Some path -> Ok path
      | None ->

src/dune_rules/context.ml

@@ -90,7 +90,8 @@ and t =
   ; default_ocamlpath : Path.t list Memo.Lazy.t
   ; build_context : Build_context.t
   ; builder : builder
-  ; which : Filename.t -> Path.t option Memo.t
+  ; which :
+      narrow_to_packages:Package.Name.Set.t option -> Filename.t -> Path.t option Memo.t
   }
 
 let default_target_exec ~target_exec toolchain =
@@ -207,7 +208,7 @@ let cms_cmt_dependency t = t.builder.cms_cmt_dependency
 let equal x y = Context_name.equal x.builder.name y.builder.name
 let hash t = Context_name.hash t.builder.name
 let build_context t = t.build_context
-let which t fname = t.which fname
+let which t ~narrow_to_packages fname = t.which ~narrow_to_packages fname
 let name t = t.builder.name
 let path t = t.builder.path
 let installed_env t = t.builder.env
@@ -425,6 +426,8 @@ let create (builder : Builder.t) ~(kind : Kind.t) =
           (fun () ->
              let+ current_env = builder.env
              and+ pkg_env = Pkg_rules.exported_env builder.name in
+             (* drop the PATH variable from the env *)
+             let pkg_env = Env.remove pkg_env ~var:Env_path.var in
              Env_path.extend_env_concat_path current_env pkg_env)
         |> Memo.Lazy.force
       in
@@ -433,18 +436,25 @@ let create (builder : Builder.t) ~(kind : Kind.t) =
   let which_outside_lockdir = Which.which ~path:builder.path in
   let which =
     match kind with
-    | Default | Opam _ -> which_outside_lockdir
+    | Default | Opam _ ->
+      fun ~narrow_to_packages ->
+        (* Narrow to deps is only relevant in the Lock-kind contexts *)
+        ignore narrow_to_packages;
+        which_outside_lockdir
     | Lock _ ->
-      let which = Staged.unstage @@ Pkg_rules.which builder.name in
-      fun prog ->
+      let which ~narrow_to_packages =
+        Staged.unstage
+        @@ Pkg_rules.which_narrowed_to_packages ~packages:narrow_to_packages builder.name
+      in
+      fun ~narrow_to_packages prog ->
         Memo.push_stack_frame
           ~human_readable_description:(fun () ->
             Pp.textf
               "looking up binary %S in context %S"
               (Filename.to_string prog)
               (Context_name.to_string builder.name))
           (fun () ->
-             which prog
+             which ~narrow_to_packages prog
              >>= function
              | Some p -> Memo.return (Some p)
              | None -> Which.which ~path:builder.path prog)
@@ -470,7 +480,11 @@ let create (builder : Builder.t) ~(kind : Kind.t) =
          let findlib_toolchain =
            Option.map builder.findlib_toolchain ~f:Context_name.to_string
          in
-         Findlib_config.discover_from_env ~env ~which ~ocamlpath ~findlib_toolchain)
+         Findlib_config.discover_from_env
+           ~env
+           ~which:(which ~narrow_to_packages:None)
+           ~ocamlpath
+           ~findlib_toolchain)
   in
   let ocaml_and_build_env_kind =
     Memo.Lazy.create

src/dune_rules/context.mli

@@ -57,7 +57,12 @@ val hash : t -> int
 val to_dyn : t -> Dyn.t
 val to_dyn_concise : t -> Dyn.t
 val name : t -> Context_name.t
-val which : t -> Filename.t -> Path.t option Memo.t
+
+val which
+  :  t
+  -> narrow_to_packages:Package.Name.Set.t option
+  -> Filename.t
+  -> Path.t option Memo.t
 
 (** [Some path/to/foo.exe] if this contexts is for feedback-directed
     optimization of target path/to/foo.exe *)

src/dune_rules/env_node.ml

@@ -2,11 +2,13 @@ open Import
 
 type t =
   { local_binaries : File_binding.Expanded.t list Memo.Lazy.t
+  ; base_env : Env.t Memo.Lazy.t
   ; external_env : Env.t Memo.Lazy.t
   ; artifacts : Artifacts.t Memo.Lazy.t
   }
 
 let local_binaries t = Memo.Lazy.force t.local_binaries
+let base_env t = Memo.Lazy.force t.base_env
 let external_env t = Memo.Lazy.force t.external_env
 let artifacts t = Memo.Lazy.force t.artifacts
 
@@ -27,6 +29,8 @@ let make
       ~expander
       ~default_env
       ~default_artifacts
+      ~lockdir_deps
+      ~lockdir_bin_env
   =
   let open Memo.O in
   let config = Dune_env.find config_stanza ~profile in
@@ -38,8 +42,8 @@ let make
       >>= extend)
   in
   let config_binaries = Option.value config.binaries ~default:[] in
-  let external_env =
-    inherited ~field:external_env ~root:default_env (fun env ->
+  let base_env =
+    inherited ~field:base_env ~root:default_env (fun env ->
       let env =
         let env = Env.extend_env env config.env_vars in
         match config_binaries with
@@ -50,15 +54,23 @@ let make
       in
       Memo.return env)
   in
+  let external_env =
+    Memo.lazy_ (fun () ->
+      let* env = Memo.Lazy.force base_env in
+      let+ bin_env = lockdir_bin_env in
+      Env_path.extend_env_concat_path env bin_env)
+  in
   let artifacts =
     inherited ~field:artifacts ~root:default_artifacts (fun binaries ->
+      let* lockdir_deps = lockdir_deps in
       Memo.parallel_map
         config_binaries
         ~f:(File_binding_expand.expand ~dir ~f:(expand_str_lazy expander))
-      >>| Artifacts.add_binaries binaries ~dir)
+      >>| Artifacts.add_binaries binaries ~dir
+      >>| Artifacts.set_lockdir_deps ~lockdir_deps)
   in
   let local_binaries =
     Memo.lazy_ (fun () -> Memo.Lazy.force artifacts >>= Artifacts.local_binaries)
   in
-  { external_env; artifacts; local_binaries }
+  { base_env; external_env; artifacts; local_binaries }
 ;;

src/dune_rules/env_node.mli

@@ -12,6 +12,8 @@ val make
   -> expander:Expander.t Memo.t
   -> default_env:Env.t Memo.t
   -> default_artifacts:Artifacts.t Memo.t
+  -> lockdir_deps:Package.Name.Set.t option Memo.t
+  -> lockdir_bin_env:Env.t Memo.t
   -> t
 
 val external_env : t -> Env.t Memo.t

src/dune_rules/expander.ml

@@ -759,6 +759,37 @@ let expand_pkg_macro ~loc { context; _ } macro_invocation =
   [ Value.Path path ]
 ;;
 
+let package_depends_by_src_dir t =
+  let open Memo.O in
+  let context_name = Context.name t.context in
+  let* lock_active = Pkg_rules.lock_dir_active context_name in
+  if not lock_active
+  then Memo.return None
+  else (
+    let src_dir = Path.Build.drop_build_context_exn t.dir in
+    match Dune_project.exclusive_package t.project ~dir:src_dir with
+    | None ->
+      (* No owning package implies no narrowing. Use the previous default of
+         all lock packages *)
+      Memo.return None
+    | Some pkg_id ->
+      let package_deps =
+        match
+          Package.Name.Map.find (Dune_project.packages t.project) (Package.Id.name pkg_id)
+        with
+        | None ->
+          (* Package ID couldn't be mapped to an actual package in the
+             project. Use the previous default of all lock packages *)
+          None
+        | Some pkg ->
+          Package.depends pkg
+          |> List.map ~f:(fun (pd : Package_dependency.t) -> pd.name)
+          |> Package.Name.Set.of_list
+          |> Option.some
+      in
+      Memo.return package_deps)
+;;
+
 let expand_pform_macro
       (context : Context.t)
       ~dir

src/dune_rules/expander.mli

@@ -8,6 +8,12 @@ val dir : t -> Path.Build.t
 val context : t -> Context_name.t
 val project : t -> Dune_project.t
 
+(** The declared package dependencies of the package owning this expander's
+    directory (via its [(dir ...)] field), used to narrow lock-directory binary
+    lookup. [None] when there is no lock directory or no owning package, meaning
+    no narrowing. *)
+val package_depends_by_src_dir : t -> Package.Name.Set.t option Memo.t
+
 val make_root
   :  project:Dune_project.t
   -> scope:Scope.t Memo.t

src/dune_rules/pkg_rules.ml

@@ -2566,6 +2566,24 @@ let all_deps universe =
 
 let all_project_deps context = all_deps (Dependencies context)
 
+let project_deps_closure ?(packages : Package.Name.Set.t option = None) context =
+  match packages with
+  | None -> all_project_deps context
+  | Some packages ->
+    let+ all_project_deps = all_project_deps context in
+    let known =
+      all_project_deps
+      |> List.map ~f:(fun (pkg : Pkg.t) -> pkg.info.name)
+      |> Package.Name.Set.of_list
+    in
+    let filtered_packages = Package.Name.Set.inter known packages in
+    let filtered_pkgs =
+      List.filter all_project_deps ~f:(fun (pkg : Pkg.t) ->
+        Package.Name.Set.mem filtered_packages pkg.info.name)
+    in
+    Pkg.top_closure filtered_pkgs
+;;
+
 let which context =
   let artifacts_and_deps =
     Memo.lazy_
@@ -2584,6 +2602,31 @@ let which context =
     Filename.Map.find artifacts program)
 ;;
 
+let which_narrowed_to_packages ~(packages : Package.Name.Set.t option) context =
+  match packages with
+  | None -> which context
+  | Some packages ->
+    let artifacts_and_deps =
+      Memo.lazy_
+        ~human_readable_description:(fun () ->
+          Pp.textf
+            "Loading binaries for packages %s in the lock directory for %S"
+            (String.concat
+               ~sep:", "
+               (List.map (Package.Name.Set.to_list packages) ~f:Package.Name.to_string))
+            (Context_name.to_string context))
+        (fun () ->
+           let+ { binaries; dep_info = _ } =
+             project_deps_closure ~packages:(Some packages) context
+             >>= Action_expander.Artifacts_and_deps.of_closure
+           in
+           binaries)
+    in
+    Staged.stage (fun program ->
+      let+ artifacts = Memo.Lazy.force artifacts_and_deps in
+      Filename.Map.find artifacts program)
+;;
+
 let ocamlpath universe =
   let+ all_project_deps = all_deps universe in
   let env = Pkg.build_env_of_deps all_project_deps in
@@ -2621,6 +2664,33 @@ let exported_env context =
   Env.extend Env.empty ~vars
 ;;
 
+let bin_env_for_packages ?(packages : Package.Name.Set.t option = None) context =
+  Memo.push_stack_frame ~human_readable_description:(fun () ->
+    Pp.textf
+      "lock directory PATH for context %S and packages %s"
+      (Context_name.to_string context)
+      (match packages with
+       | None -> "all"
+       | Some packages ->
+         String.concat
+           ~sep:", "
+           (List.map (Package.Name.Set.to_list packages) ~f:Package.Name.to_string)))
+  @@ fun () ->
+  lock_dir_active context
+  >>= function
+  | false -> Memo.return Env.empty
+  | true ->
+    let+ deps = project_deps_closure ~packages context in
+    let env = Pkg.build_env_of_deps deps in
+    (match Env.Map.find env Env_path.var with
+     | None -> Env.empty
+     | Some values ->
+       Env.add
+         Env.empty
+         ~var:Env_path.var
+         ~value:(Value_list_env.string_of_env_values values))
+;;
+
 let find_package ctx pkg =
   lock_dir_active ctx
   >>= function

src/dune_rules/pkg_rules.mli

@@ -21,7 +21,21 @@ val setup_rules
 val lock_dir_path : Context_name.t -> Path.t option Memo.t
 val lock_dir_active : Context_name.t -> bool Memo.t
 val ocaml_toolchain : Context_name.t -> Ocaml_toolchain.t Action_builder.t option Memo.t
-val which : Context_name.t -> (Filename.t -> Path.t option Memo.t) Staged.t
+
+val which_narrowed_to_packages
+  :  packages:Package.Name.Set.t option
+  -> Context_name.t
+  -> (Filename.t -> Path.t option Memo.t) Staged.t
+
+(** [bin_env_for_packages ~packages context] is an env holding only the PATH
+    (bin-layout) entries for the dependency closure of [packages] in the lock
+    directory. [None] means the whole lock directory. Empty when the context
+    has no lock directory. *)
+val bin_env_for_packages
+  :  ?packages:Package.Name.Set.t option
+  -> Context_name.t
+  -> Env.t Memo.t
+
 val exported_env : Context_name.t -> Env.t Memo.t
 val project_ocamlpath : Context_name.t -> Path.t list Memo.t
 val dev_tool_ocamlpath : Dune_pkg.Dev_tool.t -> Path.t list Memo.t

src/dune_rules/rocq/rocq_path.ml

@@ -188,7 +188,7 @@ let of_rocq_install rocq =
 
 let of_rocq_install context =
   let open Memo.O in
-  let* rocq = Context.which context Filename.rocq in
+  let* rocq = Context.which ~narrow_to_packages:None context Filename.rocq in
   match rocq with
   | None -> Memo.return []
   | Some rocq -> of_rocq_install rocq