feature(pkg): workspace package pins (#10072)

Allow a named (pin ..) stanza in workspace files. This allows
customizing the sources per lock directory.

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

bin/pkg/lock.ml

@@ -29,6 +29,15 @@ let solve_lock_dir
   =
   let open Fiber.O in
   let lock_dir = Workspace.find_lock_dir workspace lock_dir_path in
+  let project_sources =
+    match lock_dir with
+    | None -> project_sources
+    | Some lock_dir ->
+      let workspace =
+        Pin_stanza.DB.Workspace.extract workspace.sources ~names:lock_dir.pins
+      in
+      Pin_stanza.DB.combine_exn workspace project_sources
+  in
   let solver_env =
     solver_env
       ~solver_env_from_context:
@@ -126,7 +135,7 @@ let solve
 let project_sources =
   let open Memo.O in
   Dune_rules.Dune_load.projects ()
-  >>| List.fold_left ~init:(Pin_stanza.DB.empty Workspace) ~f:(fun acc project ->
+  >>| List.fold_left ~init:Pin_stanza.DB.empty ~f:(fun acc project ->
     Pin_stanza.DB.combine_exn acc (Dune_project.sources project))
 ;;
 

src/dune_pkg/pin_stanza.ml

@@ -16,6 +16,15 @@ module Package = struct
     and+ version = field_o "version" Package_version.decode in
     { name; version; loc }
   ;;
+
+  let to_local_package t ~url ~origin =
+    { Local_package.url
+    ; version = Option.value ~default:Package_version.dev t.version
+    ; loc = t.loc
+    ; origin
+    ; name = t.name
+    }
+  ;;
 end
 
 type t =
@@ -25,15 +34,24 @@ type t =
 
 let url t = t.url
 
-let decode =
+let common_fields =
   let open Dune_lang.Decoder in
-  fields
-  @@
   let+ url = field "url" OpamUrl.decode_loc
   and+ packages = multi_field "package" Package.decode in
   { url; packages }
 ;;
 
+let decode = Dune_lang.Decoder.fields common_fields
+
+let decode_with_name =
+  let open Dune_lang.Decoder in
+  fields
+  @@
+  let+ t = common_fields
+  and+ name = field "name" (located string) in
+  name, t
+;;
+
 module DB = struct
   type context =
     | Workspace
@@ -45,7 +63,7 @@ module DB = struct
     ; context : context
     }
 
-  let empty context = { all = []; map = Package_name.Map.empty; context }
+  let empty = { all = []; map = Package_name.Map.empty; context = Workspace }
   let to_dyn = Dyn.opaque
   let hash = Poly.hash
   let equal = Poly.equal
@@ -66,34 +84,32 @@ module DB = struct
     }
   ;;
 
-  let decode context =
+  let package_map_of_list list ~pin =
+    match Package_name.Map.of_list list with
+    | Ok map -> map
+    | Error (name, p, _) ->
+      let pin : Local_package.pin = pin p in
+      User_error.raise
+        ~loc:pin.loc
+        [ Pp.textf "package %S is already defined" (Package_name.to_string name) ]
+  ;;
+
+  let gen_decode context =
     let open Dune_lang.Decoder in
     let+ all = multi_field "pin" decode in
     let map =
-      match
-        List.concat_map all ~f:(fun source ->
-          List.map source.packages ~f:(fun (package : Package.t) ->
-            let name = package.name in
-            let package =
-              { Local_package.url = source.url
-              ; version = Option.value ~default:Package_version.dev package.version
-              ; loc = package.loc
-              ; origin = `Dune
-              ; name
-              }
-            in
-            name, (package, context)))
-        |> Package_name.Map.of_list
-      with
-      | Ok map -> map
-      | Error (name, (pin, _), _) ->
-        User_error.raise
-          ~loc:pin.loc
-          [ Pp.textf "package %S is already defined" (Package_name.to_string name) ]
+      List.concat_map all ~f:(fun source ->
+        List.map source.packages ~f:(fun (package : Package.t) ->
+          let name = package.name in
+          let package = Package.to_local_package package ~url:source.url ~origin:`Dune in
+          name, (package, context)))
+      |> package_map_of_list ~pin:fst
     in
     { all; map; context }
   ;;
 
+  let decode ~dir = gen_decode (Project { dir })
+
   let super_context ((_, ctx) as x) ((_, ctx') as x') =
     match ctx, ctx' with
     | Workspace, Workspace -> Code_error.raise "more than one workspace context" []
@@ -132,6 +148,47 @@ module DB = struct
   ;;
 
   let encode _ = (* CR-rgrinberg: needed for dune init *) []
+
+  module Workspace = struct
+    type nonrec t = Local_package.pin Package_name.Map.t String.Map.t
+
+    let empty = String.Map.empty
+
+    let decode =
+      let open Dune_lang.Decoder in
+      let+ pins = Dune_lang.Decoder.multi_field "pin" decode_with_name in
+      match
+        String.Map.of_list_map pins ~f:(fun ((loc, name), pin) ->
+          let packages =
+            List.map pin.packages ~f:(fun (package : Package.t) ->
+              let package = Package.to_local_package package ~url:pin.url ~origin:`Dune in
+              package.name, package)
+            |> package_map_of_list ~pin:Fun.id
+          in
+          name, (loc, packages))
+      with
+      | Ok s -> String.Map.map ~f:snd s
+      | Error (name, ((loc, _), _), _) ->
+        User_error.raise ~loc [ Pp.textf "a pin named %S already defined" name ]
+    ;;
+
+    let extract (t : t) ~names =
+      let map =
+        List.concat_map names ~f:(fun (loc, name) ->
+          match String.Map.find t name with
+          | None -> User_error.raise ~loc [ Pp.textf "pin %S doesn't exist" name ]
+          | Some packages ->
+            Package_name.Map.to_list_map packages ~f:(fun name package ->
+              name, (package, Workspace)))
+        |> package_map_of_list ~pin:fst
+      in
+      { all = []; map; context = Workspace }
+    ;;
+
+    let equal = String.Map.equal ~equal:Poly.equal
+    let to_dyn = Dyn.opaque
+    let hash = Poly.hash
+  end
 end
 
 module Scan_project = struct

src/dune_pkg/pin_stanza.mli

@@ -9,18 +9,26 @@ module Package : sig end
 module DB : sig
   type t
 
-  type context =
-    | Workspace
-    | Project of { dir : Path.Source.t }
-
-  val empty : context -> t
+  val empty : t
   val to_dyn : t -> Dyn.t
   val equal : t -> t -> bool
   val hash : t -> int
-  val decode : context -> t Dune_lang.Decoder.fields_parser
+  val decode : dir:Path.Source.t -> t Dune_lang.Decoder.fields_parser
   val encode : t -> Dune_lang.t list
   val combine_exn : t -> t -> t
   val add_opam_pins : t -> Dune_lang.Package.t Package_name.Map.t -> t
+
+  module Workspace : sig
+    type db := t
+    type t
+
+    val decode : t Dune_lang.Decoder.fields_parser
+    val empty : t
+    val extract : t -> names:(Loc.t * string) list -> db
+    val equal : t -> t -> bool
+    val to_dyn : t -> Dyn.t
+    val hash : t -> int
+  end
 end
 
 module Scan_project : sig

src/dune_rules/dune_project.ml

@@ -427,7 +427,7 @@ let infer ~dir info packages =
   let opam_file_location = opam_file_location_default ~lang in
   { name
   ; allow_approximate_merlin = None
-  ; sources = Dune_pkg.Pin_stanza.DB.empty (Project { dir })
+  ; sources = Dune_pkg.Pin_stanza.DB.empty
   ; packages
   ; root
   ; info
@@ -724,7 +724,7 @@ let parse ~dir ~(lang : Lang.Instance.t) ~file =
      and+ version = field_o "version" Package_version.decode
      and+ info = Package_info.decode ()
      and+ packages = multi_field "package" (Package.decode ~dir)
-     and+ sources = Dune_pkg.Pin_stanza.DB.decode (Project { dir })
+     and+ sources = Dune_pkg.Pin_stanza.DB.decode ~dir
      and+ explicit_extensions =
        multi_field
          "using"

src/dune_rules/workspace.ml

@@ -12,7 +12,7 @@ module Lock_dir = struct
     ; unset_solver_vars : Package_variable_name.Set.t option
     ; repositories : (Loc.t * Dune_pkg.Pkg_workspace.Repository.Name.t) list
     ; constraints : Dune_lang.Package_dependency.t list
-    ; sources : string list
+    ; pins : (Loc.t * string) list
     }
 
   let to_dyn
@@ -22,7 +22,7 @@ module Lock_dir = struct
     ; unset_solver_vars
     ; repositories
     ; constraints
-    ; sources
+    ; pins
     }
     =
     Dyn.record
@@ -36,7 +36,7 @@ module Lock_dir = struct
             Dune_pkg.Pkg_workspace.Repository.Name.to_dyn
             (List.map repositories ~f:snd) )
       ; "constraints", Dyn.list Dune_lang.Package_dependency.to_dyn constraints
-      ; "sources", (Dyn.list Dyn.string) sources
+      ; "pins", (Dyn.list Dyn.string) (List.map pins ~f:snd)
       ]
   ;;
 
@@ -47,7 +47,7 @@ module Lock_dir = struct
     ; unset_solver_vars
     ; repositories
     ; constraints
-    ; sources
+    ; pins
     }
     =
     Poly.hash
@@ -57,7 +57,7 @@ module Lock_dir = struct
       , unset_solver_vars
       , repositories
       , constraints
-      , sources )
+      , pins )
   ;;
 
   let equal
@@ -67,7 +67,7 @@ module Lock_dir = struct
     ; unset_solver_vars
     ; repositories
     ; constraints
-    ; sources
+    ; pins
     }
     t
     =
@@ -83,7 +83,7 @@ module Lock_dir = struct
          repositories
          t.repositories
     && List.equal Dune_lang.Package_dependency.equal constraints t.constraints
-    && List.equal String.equal sources t.sources
+    && List.equal (Tuple.T2.equal Loc.equal String.equal) pins t.pins
   ;;
 
   let decode ~dir =
@@ -108,7 +108,7 @@ module Lock_dir = struct
       and+ repositories = Dune_lang.Ordered_set_lang.field "repositories"
       and+ constraints =
         field ~default:[] "constraints" (repeat Dune_lang.Package_dependency.decode)
-      and+ sources = field ~default:[] "sources" (repeat string) in
+      and+ pins = field ~default:[] "pins" (repeat (located string)) in
       Option.iter solver_env ~f:(fun solver_env ->
         Option.iter
           unset_solver_vars
@@ -133,7 +133,7 @@ module Lock_dir = struct
       ; version_preference
       ; repositories = repositories_of_ordered_set repositories
       ; constraints
-      ; sources
+      ; pins
       }
     in
     fields decode
@@ -596,7 +596,7 @@ type t =
   ; repos : Dune_pkg.Pkg_workspace.Repository.t list
   ; lock_dirs : Lock_dir.t list
   ; dir : Path.Source.t
-  ; sources : Dune_pkg.Pin_stanza.DB.t
+  ; sources : Dune_pkg.Pin_stanza.DB.Workspace.t
   }
 
 let to_dyn { merlin_context; contexts; env; config; repos; lock_dirs; sources; dir } =
@@ -609,7 +609,7 @@ let to_dyn { merlin_context; contexts; env; config; repos; lock_dirs; sources; d
     ; "repos", list Repository.to_dyn repos
     ; "solver", (list Lock_dir.to_dyn) lock_dirs
     ; "dir", Path.Source.to_dyn dir
-    ; "sources", Dune_pkg.Pin_stanza.DB.to_dyn sources
+    ; "sources", Dune_pkg.Pin_stanza.DB.Workspace.to_dyn sources
     ]
 ;;
 
@@ -621,7 +621,7 @@ let equal { merlin_context; contexts; env; config; repos; lock_dirs; dir; source
   && List.equal Repository.equal repos w.repos
   && List.equal Lock_dir.equal lock_dirs w.lock_dirs
   && Path.Source.equal dir w.dir
-  && Dune_pkg.Pin_stanza.DB.equal sources w.sources
+  && Dune_pkg.Pin_stanza.DB.Workspace.equal sources w.sources
 ;;
 
 let hash { merlin_context; contexts; env; config; repos; lock_dirs; dir; sources } =
@@ -633,7 +633,7 @@ let hash { merlin_context; contexts; env; config; repos; lock_dirs; dir; sources
     , List.hash Repository.hash repos
     , List.hash Lock_dir.hash lock_dirs
     , Path.Source.hash dir
-    , Dune_pkg.Pin_stanza.DB.hash sources )
+    , Dune_pkg.Pin_stanza.DB.Workspace.hash sources )
 ;;
 
 let find_lock_dir t path =
@@ -801,7 +801,7 @@ let step1 clflags =
          ~default:(lazy []))
   and+ config_from_workspace_file = Dune_config.decode_fields_of_workspace_file
   and+ lock_dirs = multi_field "lock_dir" (Lock_dir.decode ~dir)
-  and+ sources = Dune_pkg.Pin_stanza.DB.decode Workspace in
+  and+ sources = Dune_pkg.Pin_stanza.DB.Workspace.decode in
   let+ contexts = multi_field "context" (lazy_ Context.decode) in
   let config =
     create_final_config
@@ -907,7 +907,7 @@ let default clflags =
   ; repos = default_repositories
   ; lock_dirs = []
   ; dir = Path.Source.root
-  ; sources = Dune_pkg.Pin_stanza.DB.empty Workspace
+  ; sources = Dune_pkg.Pin_stanza.DB.Workspace.empty
   }
 ;;
 

src/dune_rules/workspace.mli

@@ -10,7 +10,7 @@ module Lock_dir : sig
     ; unset_solver_vars : Dune_lang.Package_variable_name.Set.t option
     ; repositories : (Loc.t * Dune_pkg.Pkg_workspace.Repository.Name.t) list
     ; constraints : Dune_lang.Package_dependency.t list
-    ; sources : string list
+    ; pins : (Loc.t * string) list
     }
 
   val equal : t -> t -> bool
@@ -107,7 +107,7 @@ type t = private
   ; repos : Dune_pkg.Pkg_workspace.Repository.t list
   ; lock_dirs : Lock_dir.t list
   ; dir : Path.Source.t
-  ; sources : Dune_pkg.Pin_stanza.DB.t
+  ; sources : Dune_pkg.Pin_stanza.DB.Workspace.t
   }
 
 val equal : t -> t -> bool