Refactor Link_time_code_gen (#2095)

- share more of the work between the various odes - add a helper function for generating a module Signed-off-by: Jeremie Dimino <[email protected]>
ocaml · Apr 30, 2019 · a37a9be · a37a9be
1 parent 49d9384
commit a37a9be
Show file tree

Hide file tree

Showing 7 changed files with 135 additions and 102 deletions.
diff --git a/src/arg_spec.ml b/src/arg_spec.ml
@@ -86,11 +86,13 @@ let quote_args =
   in
   fun quote args -> As (loop quote args)
 
+let fail e = Fail { fail = fun _ -> raise e }
+
 let of_result = function
   | Ok x -> x
-  | Error e -> Fail {fail = fun _ -> raise e}
+  | Error e -> fail e
 
 let of_result_map res ~f =
   match res with
   | Ok    x -> f x
-  | Error e -> Fail {fail = fun _ -> raise e}
+  | Error e -> fail e
diff --git a/src/arg_spec.mli b/src/arg_spec.mli
@@ -60,3 +60,4 @@ val quote_args : string -> string list -> _ t
 
 val of_result : ('a, 'b) t Or_exn.t -> ('a, 'b) t
 val of_result_map : 'a Or_exn.t -> f:('a -> ('b, 'c) t) -> ('b, 'c) t
+val fail : exn -> ('a, 'b) t
diff --git a/src/exe.ml b/src/exe.ml
@@ -150,9 +150,6 @@ let link_exe
       Build.dyn_paths (Build.arr (fun (modules, _) ->
         artifacts modules ~ext:ctx.ext_obj))
   in
-  let arg_spec_for_requires =
-    Lazy.force (Mode.Dict.get arg_spec_for_requires mode)
-  in
   (* The rule *)
   SC.add_rule sctx ~loc ~dir
     (Build.fanout3
@@ -169,7 +166,7 @@ let link_exe
        ; A "-o"; Target exe
        ; As linkage.flags
        ; Dyn (fun (_, _, link_flags) -> As link_flags)
-       ; Arg_spec.of_result_map arg_spec_for_requires ~f:Fn.id
+       ; arg_spec_for_requires
        ; Dyn (fun (cm_files, _, _) -> Deps cm_files)
        ]);
   if linkage.ext = ".bc" then
@@ -197,19 +194,29 @@ let build_and_link_many
   (* CR-someday jdimino: this should probably say [~dynlink:false] *)
   Module_compilation.build_modules cctx ~js_of_ocaml ~dep_graphs;
 
+  let arg_spec_for_requires =
+    let f =
+      Staged.unstage (Link_time_code_gen.libraries_link cctx)
+    in
+    List.map linkages ~f:(fun x -> x.Linkage.mode)
+    |> Mode.Dict.Set.of_list
+    |> Mode.Dict.mapi ~f:(fun mode x ->
+      if x then
+        f mode
+      else
+        Arg_spec.fail Exit)
+  in
   List.iter programs ~f:(fun { Program.name; main_module_name ; loc } ->
     let top_sorted_modules =
       let main = Option.value_exn
                    (Module.Name.Map.find (CC.modules cctx) main_module_name) in
       Dep_graph.top_closed_implementations dep_graphs.impl
         [main]
     in
-    let arg_spec_for_requires =
-      Mode.Dict.of_func (fun ~mode ->
-        lazy (Result.map (CC.requires_link cctx)
-                ~f:(Link_time_code_gen.libraries_link ~loc ~name ~mode cctx)))
-    in
-    List.iter linkages ~f:(fun linkage ->
+    List.iter linkages ~f:(fun (linkage : Linkage.t) ->
+      let arg_spec_for_requires =
+        Mode.Dict.get arg_spec_for_requires linkage.mode
+      in
       link_exe cctx
         ~loc
         ~name

diff --git a/src/link_time_code_gen.ml b/src/link_time_code_gen.ml
@@ -3,80 +3,106 @@ open Import
 module CC = Compilation_context
 module SC = Super_context
 
-let of_libs = List.map ~f:(fun l -> Lib.Lib_and_module.Lib l)
-
-let rec cut_after_libs ~pkg_name before = function
-  | [] -> None
-  | a::l when Lib_name.compare (Lib.name a) pkg_name = Eq ->
-    Some (List.rev (a::before),l)
-  | a::l -> cut_after_libs (a::before) ~pkg_name l
-
-let findlib_dynload = Lib_name.of_string_exn ~loc:None "findlib.dynload"
-
-let libraries_link ~name ~loc ~mode cctx libs =
+let generate_and_compile_module cctx ~name:basename ~code ~requires =
   let sctx       = CC.super_context cctx in
-  let ctx        = SC.context       sctx in
   let obj_dir    = CC.obj_dir       cctx in
   let dir        = CC.dir           cctx in
-  let stdlib_dir = ctx.stdlib_dir in
-  match cut_after_libs [] ~pkg_name:findlib_dynload libs with
-  | Some (before, after) ->
-    (* If findlib.dynload is linked, we stores in the binary the packages linked
-       by linking just after findlib.dynload a module containing the info *)
-    let libs =
-      List.filter
-        ~f:(fun lib -> not (Lib_info.Status.is_private (Lib.status lib)))
-        libs
-    in
-    let preds = Variant.Set.add Findlib.Package.preds (Mode.variant mode) in
-    let s =
-      Format.asprintf "%a@\nFindlib.record_package_predicates %a;;@."
-        (Fmt.list ~pp_sep:Fmt.nl (fun fmt lib ->
-           Format.fprintf fmt "Findlib.record_package Findlib.Record_core %a;;"
-             Lib_name.pp_quoted (Lib.name lib)))
-        libs
-        (Fmt.ocaml_list Variant.pp) (Variant.Set.to_list preds)
-    in
-    let basename = Format.asprintf "%s_findlib_initl_%a" name Mode.pp mode in
-    let ml  = Path.relative (Obj_dir.obj_dir obj_dir) (basename ^ ".ml") in
-    SC.add_rule ~dir sctx (Build.write_file ml s);
-    let impl = Module.File.make OCaml ml in
-    let name = Module.Name.of_string basename in
-    let module_ =
-      Module.make ~impl name ~visibility:Public ~obj_dir ~kind:Impl in
-    let requires =
-      Lib.DB.find_many ~loc (SC.public_libs sctx)
-        [Lib_name.of_string_exn ~loc:(Some loc) "findlib"]
-    in
-    let opaque =
-      Ocaml_version.supports_opaque_for_mli
-        (Super_context.context sctx).version
-    in
-    let cctx =
-      Compilation_context.create
-        ~super_context:sctx
-        ~expander:(Compilation_context.expander cctx)
-        ~scope:(Compilation_context.scope cctx)
-        ~dir_kind:(Compilation_context.dir_kind cctx)
-        ~obj_dir:(Compilation_context.obj_dir cctx)
-        ~modules:(Module.Name.Map.singleton name module_)
-        ~requires_compile:requires
-        ~requires_link:(lazy requires)
-        ~flags:Ocaml_flags.empty
-        ~opaque
-        ()
-    in
-    Module_compilation.build_module
-      ~dep_graphs:(Dep_graph.Ml_kind.dummy module_)
-      cctx
-      module_;
-    let lm =
-      of_libs before
-      @ [Lib.Lib_and_module.Module module_]
-      @ of_libs after
+  let ml = Path.relative (Obj_dir.obj_dir obj_dir) (basename ^ ".ml") in
+  SC.add_rule ~dir sctx (Build.write_file ml code);
+  let impl = Module.File.make OCaml ml in
+  let name = Module.Name.of_string basename in
+  let module_ =
+    Module.make ~impl name ~visibility:Public ~obj_dir ~kind:Impl
+  in
+  let opaque =
+    Ocaml_version.supports_opaque_for_mli
+      (Super_context.context sctx).version
+  in
+  let cctx =
+    Compilation_context.create
+      ~super_context:sctx
+      ~expander:(Compilation_context.expander cctx)
+      ~scope:(Compilation_context.scope cctx)
+      ~dir_kind:(Compilation_context.dir_kind cctx)
+      ~obj_dir:(Compilation_context.obj_dir cctx)
+      ~modules:(Module.Name.Map.singleton name module_)
+      ~requires_compile:requires
+      ~requires_link:(lazy requires)
+      ~flags:Ocaml_flags.empty
+      ~opaque
+      ()
+  in
+  Module_compilation.build_module
+    ~dep_graphs:(Dep_graph.Ml_kind.dummy module_)
+    cctx
+    module_;
+  module_
+
+let is_findlib_dynload lib =
+  match Lib_name.to_string (Lib.name lib) with
+  | "findlib.dynload" -> true
+  | _ -> false
+
+let libraries_link cctx  =
+  match CC.requires_link cctx with
+  | Error exn ->
+    let arg_spec = Arg_spec.fail exn in
+    Staged.stage (fun _mode -> arg_spec)
+  | Ok libs ->
+    let sctx       = CC.super_context cctx in
+    let ctx        = SC.context       sctx in
+    let stdlib_dir = ctx.stdlib_dir in
+    let has_findlib_dynload =
+      List.exists libs ~f:is_findlib_dynload
     in
-    Arg_spec.S [ A "-linkall"
-               ; Lib.Lib_and_module.link_flags lm ~mode ~stdlib_dir
-               ]
-  | None ->
-    Lib.L.link_flags libs ~mode ~stdlib_dir
+    if not has_findlib_dynload then
+      Staged.stage (fun mode -> Lib.L.link_flags libs ~mode ~stdlib_dir)
+    else begin
+      (* If findlib.dynload is linked, we stores in the binary the
+         packages linked by linking just after findlib.dynload a
+         module containing the info *)
+      let public_libs =
+        List.filter
+          ~f:(fun lib -> not (Lib_info.Status.is_private (Lib.status lib)))
+          libs
+      in
+      let requires =
+        (* This shouldn't fail since findlib.dynload depends on
+           findlib. That's why it's ok to use a dummy location. *)
+        Lib.DB.find_many ~loc:Loc.none (SC.public_libs sctx)
+          [Lib_name.of_string_exn ~loc:None "findlib"]
+      in
+      Staged.stage (fun mode ->
+        let preds = Variant.Set.add Findlib.Package.preds (Mode.variant mode) in
+        let code =
+          Format.asprintf "%a@\nFindlib.record_package_predicates %a;;@."
+            (Fmt.list ~pp_sep:Fmt.nl (fun fmt lib ->
+               Format.fprintf fmt "Findlib.record_package Findlib.Record_core %a;;"
+                 Lib_name.pp_quoted (Lib.name lib)))
+            public_libs
+            (Fmt.ocaml_list Variant.pp)
+            (Variant.Set.to_list preds)
+        in
+        let module_ =
+          generate_and_compile_module
+            cctx
+            ~name:(Format.sprintf "findlib_initl_%s"
+                     (Mode.to_string mode))
+            ~code
+            ~requires
+        in
+        let rec insert = function
+          | [] -> assert false
+          | lib :: libs ->
+            if is_findlib_dynload lib then
+              Lib.Lib_and_module.Lib lib
+              :: Module module_
+              :: List.map libs ~f:(fun lib -> Lib.Lib_and_module.Lib lib)
+            else
+              Lib lib :: insert libs
+        in
+        Arg_spec.S
+          [ A "-linkall"
+          ; Lib.Lib_and_module.link_flags (insert libs) ~mode ~stdlib_dir
+          ])
+    end
diff --git a/src/link_time_code_gen.mli b/src/link_time_code_gen.mli
@@ -2,11 +2,5 @@
 
 open Stdune
 
-val libraries_link
-  :  name:string
-  -> loc:Loc.t
-  -> mode:Mode.t
-  -> Compilation_context.t
-  -> Lib.L.t
-  -> _ Arg_spec.t
 (** Insert link time generated code for findlib_dynload in the list *)
+val libraries_link : Compilation_context.t -> (Mode.t -> _ Arg_spec.t) Staged.t
diff --git a/src/mode.ml b/src/mode.ml
@@ -12,20 +12,15 @@ let decode =
     ; "native" , Native
     ]
 
-let encode =
-  let open Dune_lang.Encoder in
-  function
-  | Byte -> string "byte"
-  | Native -> string "native"
-
-let pp fmt = function
-  | Byte -> Format.pp_print_string fmt "byte"
-  | Native -> Format.pp_print_string fmt "native"
-
 let choose byte native = function
   | Byte   -> byte
   | Native -> native
 
+let to_string = choose "byte" "native"
+
+let encode t = Dune_lang.Encoder.string (to_string t)
+let pp fmt t = Format.pp_print_string fmt (to_string t)
+
 let compiled_unit_ext = choose (Cm_kind.ext Cmo) (Cm_kind.ext Cmx)
 let compiled_lib_ext = choose ".cma" ".cmxa"
 let plugin_ext = choose ".cma" ".cmxs"
@@ -73,6 +68,11 @@ module Dict = struct
     ; native = f t.native
     }
 
+  let mapi t ~f =
+    { byte = f Byte t.byte
+    ; native = f Native t.native
+    }
+
   let make_both x =
     { byte   = x
     ; native = x

diff --git a/src/mode.mli b/src/mode.mli
@@ -16,6 +16,8 @@ val of_cm_kind : Cm_kind.t -> t
 
 val variant : t -> Variant.t
 
+val to_string : t -> string
+
 val pp : t Fmt.t
 
 module Dict : sig
@@ -45,6 +47,7 @@ module Dict : sig
   val map2 : 'a t -> 'b t -> f:('a -> 'b -> 'c) -> 'c t
 
   val map : 'a t -> f:('a -> 'b) -> 'b t
+  val mapi : 'a t -> f:(mode -> 'a -> 'b) -> 'b t
 
   val make_both : 'a -> 'a t
Original file line number	Diff line number	Diff line change
Expand Up		@@ -60,3 +60,4 @@ val quote_args : string -> string list -> _ t

		val of_result : ('a, 'b) t Or_exn.t -> ('a, 'b) t
		val of_result_map : 'a Or_exn.t -> f:('a -> ('b, 'c) t) -> ('b, 'c) t
		val fail : exn -> ('a, 'b) t