Add support for staged ppx rewriters such as ones using the typer (#1080)

Fix #193

Signed-off-by: Jeremie Dimino <jeremie@dimino.org>

CHANGES.md

@@ -21,6 +21,9 @@ next
 - Add support for multi directory libraries by writing
   `(include_subdirs qualified)` (#1034, @diml)
 
+- Add `(staged_pps ...)` to support staged ppx rewriters such as ones
+  using the OCaml typer like `ppx_import` (#1080, fix #193, @diml)
+
 1.0.1 (19/07/2018)
 ------------------
 

doc/dune-files.rst

@@ -1107,9 +1107,31 @@ Dune accepts three kinds of preprocessing:
 - ``(action <action>)`` to preprocess files using the given action
 - ``(pps <ppx-rewriters-and-flags>)`` to preprocess files using the given list
   of ppx rewriters
+- ``(staged_pps <ppx-rewriters-and-flags>)`` is similar to ``(pps
+    ...)`` but behave slightly differently and is needed for certain
+    ppx rewriters (see below for details)
 
-Note that in any cases, files are preprocessed only once. Dune doesn't use
-the ``-pp`` or ``-ppx`` of the various OCaml tools.
+Dune normally assumes that the compilation pipeline is sequenced as
+follow:
+
+- code generation (including preprocessing)
+- dependency analysis
+- compilation
+
+Dune uses this fact to optimize the pipeline and in particular share
+the result of code generation and preprocessing between the dependency
+analysis and compilation phases. However, some specific code
+generators or preprocessors require feedback from the compilation
+phase. As a result they must be applied in stages as follows:
+
+- first stage of code geneneration
+- dependency analysis
+- second step of code generation in parallel with compilation
+
+This is the case for ppx rewriters using the OCaml typer for
+instance. When using such ppx rewriters, you must use ``staged_pps``
+instead of ``pps`` in order to force Dune to use the second pipeline,
+which is slower but necessary in this case.
 
 Preprocessing with actions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~

src/jbuild.ml

@@ -356,7 +356,12 @@ module Dep_conf = struct
 end
 
 module Preprocess = struct
-  type pps = { loc : Loc.t; pps : (Loc.t * Pp.t) list; flags : string list }
+  type pps =
+    { loc : Loc.t
+    ; pps : (Loc.t * Pp.t) list
+    ; flags : string list
+    ; staged : bool
+    }
   type t =
     | No_preprocessing
     | Action of Loc.t * Action.Unexpanded.t
@@ -371,7 +376,12 @@ module Preprocess = struct
       ; "pps",
         (let%map loc = loc
          and pps, flags = Pps_and_flags.t in
-         Pps { loc; pps; flags })
+         Pps { loc; pps; flags; staged = false })
+      ; "staged_pps",
+        (let%map () = Syntax.since Stanza.syntax (1, 1)
+         and loc = loc
+         and pps, flags = Pps_and_flags.t in
+         Pps { loc; pps; flags; staged = true })
       ]
 
   let pps = function

src/jbuild.mli

@@ -15,6 +15,7 @@ module Preprocess : sig
     { loc   : Loc.t
     ; pps   : (Loc.t * Pp.t) list
     ; flags : string list
+    ; staged : bool
     }
 
   type t =

src/merlin.ml

@@ -18,14 +18,17 @@ module Preprocess = struct
     | Other, Other -> Other
     | Pps _, Other -> a
     | Other, Pps _ -> b
-    | Pps { loc = _; pps = pps1; flags = flags1 },
-      Pps { loc = _; pps = pps2; flags = flags2 } ->
+    | Pps { loc = _; pps = pps1; flags = flags1; staged = s1 },
+      Pps { loc = _; pps = pps2; flags = flags2; staged = s2 } ->
       match
-        match List.compare flags1 flags2 ~compare:String.compare with
+        match Bool.compare s1 s2 with
+        | Gt| Lt as ne -> ne
         | Eq ->
-          List.compare pps1 pps2 ~compare:(fun (_, a) (_, b) ->
-            Jbuild.Pp.compare a b)
-        | ne -> ne
+          match List.compare flags1 flags2 ~compare:String.compare with
+          | Gt | Lt as ne -> ne
+          | Eq ->
+            List.compare pps1 pps2 ~compare:(fun (_, a) (_, b) ->
+              Jbuild.Pp.compare a b)
       with
       | Eq -> a
       | _  -> Other
@@ -98,7 +101,7 @@ let add_source_dir t dir =
 
 let ppx_flags sctx ~dir:_ ~scope ~dir_kind { preprocess; libname; _ } =
   match preprocess with
-  | Pps { loc = _; pps; flags } -> begin
+  | Pps { loc = _; pps; flags; staged = _ } -> begin
     match Preprocessing.get_ppx_driver sctx ~scope ~dir_kind pps with
     | Ok exe ->
       (Path.to_absolute_filename exe

src/module.ml

@@ -39,6 +39,7 @@ type t =
   ; impl     : File.t option
   ; intf     : File.t option
   ; obj_name : string
+  ; pp       : (unit, string list) Build.t option
   }
 
 let name t = t.name
@@ -68,6 +69,7 @@ let make ?impl ?intf ?obj_name name =
   ; impl
   ; intf
   ; obj_name
+  ; pp = None
   }
 
 let real_unit_name t = Name.of_string (Filename.basename t.obj_name)
@@ -126,3 +128,5 @@ let dir t =
     | None -> Option.value_exn t.impl
   in
   Path.parent_exn file.path
+
+let set_pp t pp = { t with pp }

src/module.mli

@@ -44,6 +44,7 @@ type t = private
   ; intf      : File.t option
   ; obj_name  : string (** Object name. It is different from [name]
                            for wrapped modules. *)
+  ; pp        : (unit, string list) Build.t option (** Preprocessing flags *)
   }
 
 val make
@@ -84,3 +85,5 @@ val has_impl : t -> bool
 val with_wrapper : t -> libname:string -> t
 
 val map_files : t -> f:(Ml_kind.t -> File.t -> File.t) -> t
+
+val set_pp : t -> (unit, string list) Build.t option -> t

src/module_compilation.ml

@@ -101,12 +101,21 @@ let build_cm cctx ?sandbox ?(dynlink=true) ~dep_graphs ~cm_kind (m : Module.t) =
         end else
           (ctx.build_dir, As [])
       in
+      let flags =
+        let flags = Ocaml_flags.get_for_cm (CC.flags cctx) ~cm_kind in
+        match m.pp with
+        | None -> flags
+        | Some pp ->
+          Build.fanout flags pp >>^ fun (flags, pp_flags) ->
+          flags @ pp_flags
+      in
       SC.add_rule sctx ?sandbox
         (Build.paths extra_deps >>>
          other_cm_files >>>
-         Ocaml_flags.get_for_cm (CC.flags cctx) ~cm_kind >>>
+         flags
+         >>>
          Build.run ~dir ~context:ctx (Ok compiler)
-           [ Dyn (fun ocaml_flags -> As ocaml_flags)
+           [ Dyn (fun flags -> As flags)
            ; no_keep_locs
            ; cmt_args
            ; A "-I"; Path obj_dir

src/ocamldep.ml

@@ -137,9 +137,15 @@ let deps_of cctx ~ml_kind unit =
       let all_deps_file = all_deps_path file in
       let ocamldep_output = file_in_obj_dir file ~suffix:".d" in
       SC.add_rule sctx
-        ( Build.run ~context (Ok context.ocamldep)
-            [A "-modules"; Ml_kind.flag ml_kind; Dep file]
-            ~stdout_to:ocamldep_output
+        (let flags = Option.value unit.pp ~default:(Build.return []) in
+         flags >>>
+         Build.run ~context (Ok context.ocamldep)
+           [ A "-modules"
+           ; Dyn (fun flags -> As flags)
+           ; Ml_kind.flag ml_kind
+           ; Dep file
+           ]
+           ~stdout_to:ocamldep_output
         );
       let build_paths dependencies =
         let dependency_file_path m =

src/ordered_set_lang.ml

@@ -231,7 +231,13 @@ let standard =
   ; context = Univ_map.empty
   }
 
-let field ?(default=standard) name = Sexp.Of_sexp.field name t ~default
+let field ?(default=standard) ?check name =
+  let t =
+    match check with
+    | None -> t
+    | Some x -> Sexp.Of_sexp.(>>>) x t
+  in
+  Sexp.Of_sexp.field name t ~default
 
 module Unexpanded = struct
   type ast = (String_with_vars.t, Ast.unexpanded) Ast.t
@@ -265,7 +271,13 @@ module Unexpanded = struct
 
   let standard = standard
 
-  let field ?(default=standard) name = Stanza.Of_sexp.field name t ~default
+  let field ?(default=standard) ?check name =
+    let t =
+      match check with
+      | None -> t
+      | Some x -> Sexp.Of_sexp.(>>>) x t
+    in
+    Sexp.Of_sexp.field name t ~default
 
   let files t ~f =
     let rec loop acc (ast : ast) =

src/ordered_set_lang.mli

@@ -66,7 +66,11 @@ end
 val standard : t
 val is_standard : t -> bool
 
-val field : ?default:t -> string -> t Sexp.Of_sexp.fields_parser
+val field
+  :  ?default:t
+  -> ?check:unit Sexp.Of_sexp.t
+  -> string
+  -> t Sexp.Of_sexp.fields_parser
 
 module Unexpanded : sig
   type expanded = t
@@ -75,7 +79,11 @@ module Unexpanded : sig
   include Sexp.Sexpable with type t := t
   val standard : t
 
-  val field : ?default:t -> string -> t Sexp.Of_sexp.fields_parser
+  val field
+    :  ?default:t
+    -> ?check:unit Sexp.Of_sexp.t
+    -> string
+    -> t Sexp.Of_sexp.fields_parser
 
   val has_special_forms : t -> bool
 

src/preprocessing.ml

@@ -21,11 +21,12 @@ module Driver = struct
     module Info = struct
       let name = Sub_system_name.make "ppx.driver"
       type t =
-        { loc        : Loc.t
-        ; flags      : Ordered_set_lang.Unexpanded.t
-        ; lint_flags : Ordered_set_lang.Unexpanded.t
-        ; main       : string
-        ; replaces   : (Loc.t * string) list
+        { loc          : Loc.t
+        ; flags        : Ordered_set_lang.Unexpanded.t
+        ; as_ppx_flags : Ordered_set_lang.Unexpanded.t
+        ; lint_flags   : Ordered_set_lang.Unexpanded.t
+        ; main         : string
+        ; replaces     : (Loc.t * string) list
         }
 
       type Jbuild.Sub_system_info.t += T of t
@@ -46,12 +47,16 @@ module Driver = struct
         record
           (let%map loc = loc
            and flags = Ordered_set_lang.Unexpanded.field "flags"
+           and as_ppx_flags =
+             Ordered_set_lang.Unexpanded.field "flags"
+               ~check:(Syntax.since syntax (1, 1))
            and lint_flags = Ordered_set_lang.Unexpanded.field "lint_flags"
            and main = field "main" string
            and replaces = field "replaces" (list (located string)) ~default:[]
            in
            { loc
            ; flags
+           ; as_ppx_flags
            ; lint_flags
            ; main
            ; replaces
@@ -482,7 +487,10 @@ let lint_module sctx ~dir ~dep_kind ~lint ~lib_name ~scope ~dir_kind =
                         ~targets:(Static [])
                         ~targets_dir:dir
                         ~scope)))
-        | Pps { loc; pps; flags } ->
+        | Pps { loc; pps; flags; staged } ->
+          if staged then
+            Loc.fail loc
+              "Staged ppx rewriters cannot be used as linters.";
           let args : _ Arg_spec.t =
             S [ As flags
               ; As (cookie_library_name lib_name)
@@ -569,49 +577,82 @@ let make sctx ~dir ~dep_kind ~lint ~preprocess
            |> setup_reason_rules sctx in
          if lint then lint_module ~ast ~source:m;
          ast)
-    | Pps { loc; pps; flags } ->
-      let args : _ Arg_spec.t =
-        S [ As flags
-          ; As (cookie_library_name lib_name)
-          ]
-      in
-      let corrected_suffix = ".ppx-corrected" in
-      let driver_and_flags =
-        let open Result.O in
-        get_ppx_driver sctx ~loc ~scope ~dir_kind pps >>| fun (exe, driver) ->
-        (exe,
-         let bindings =
-           Pform.Map.singleton "corrected-suffix"
-             (Values [String corrected_suffix])
-         in
-         Build.memoize "ppx flags"
-           (SC.expand_and_eval_set sctx driver.info.flags
-              ~scope
-              ~dir
-              ~bindings
-              ~standard:(Build.return [])))
-      in
-      (fun m ~lint ->
-         let ast = setup_reason_rules sctx m in
-         if lint then lint_module ~ast ~source:m;
-         pped_module ast ~f:(fun kind src dst ->
-           SC.add_rule sctx
-             (promote_correction ~suffix:corrected_suffix
-                (Option.value_exn (Module.file m kind))
-                (preprocessor_deps >>^ ignore
-                 >>>
-                 Build.of_result_map driver_and_flags
-                   ~targets:[dst]
-                   ~f:(fun (exe, flags) ->
-                     flags
-                     >>>
-                     Build.run ~context:(SC.context sctx)
-                       (Ok exe)
-                       [ args
-                       ; A "-o"; Target dst
-                       ; Ml_kind.ppx_driver_flag kind; Dep src
-                       ; Dyn (fun x -> As x)
-                       ]))))))
+    | Pps { loc; pps; flags; staged } ->
+      if not staged then begin
+        let args : _ Arg_spec.t =
+          S [ As flags
+            ; As (cookie_library_name lib_name)
+            ]
+        in
+        let corrected_suffix = ".ppx-corrected" in
+        let driver_and_flags =
+          let open Result.O in
+          get_ppx_driver sctx ~loc ~scope ~dir_kind pps >>| fun (exe, driver) ->
+          (exe,
+           let bindings =
+             Pform.Map.singleton "corrected-suffix"
+               (Values [String corrected_suffix])
+           in
+           Build.memoize "ppx flags"
+             (SC.expand_and_eval_set sctx driver.info.flags
+                ~scope
+                ~dir
+                ~bindings
+                ~standard:(Build.return ["--as-ppx"])))
+        in
+        (fun m ~lint ->
+           let ast = setup_reason_rules sctx m in
+           if lint then lint_module ~ast ~source:m;
+           pped_module ast ~f:(fun kind src dst ->
+             SC.add_rule sctx
+               (promote_correction ~suffix:corrected_suffix
+                  (Option.value_exn (Module.file m kind))
+                  (preprocessor_deps >>^ ignore
+                   >>>
+                   Build.of_result_map driver_and_flags
+                     ~targets:[dst]
+                     ~f:(fun (exe, flags) ->
+                       flags
+                       >>>
+                       Build.run ~context:(SC.context sctx)
+                         (Ok exe)
+                         [ args
+                         ; A "-o"; Target dst
+                         ; Ml_kind.ppx_driver_flag kind; Dep src
+                         ; Dyn (fun x -> As x)
+                         ])))))
+      end else begin
+        let pp_flags = Build.of_result (
+          let open Result.O in
+          get_ppx_driver sctx ~loc ~scope ~dir_kind pps >>| fun (exe, driver) ->
+          Build.memoize "ppx command"
+            (Build.path exe
+             >>>
+             preprocessor_deps >>^ ignore
+             >>>
+             SC.expand_and_eval_set sctx driver.info.as_ppx_flags
+               ~scope
+               ~dir
+               ~standard:(Build.return [])
+             >>^ fun flags ->
+             let command =
+               List.map
+                 (List.concat
+                    [ [Path.reach exe ~from:(SC.context sctx).build_dir]
+                    ; flags
+                    ; cookie_library_name lib_name
+                    ])
+                 ~f:quote_for_shell
+               |> String.concat ~sep:" "
+             in
+             ["-ppx"; command]))
+        in
+        let pp = Some pp_flags in
+        (fun m ~lint ->
+           let ast = setup_reason_rules sctx m in
+           if lint then lint_module ~ast ~source:m;
+           Module.set_pp m pp)
+      end)
 
 let pp_modules t ?(lint=true) modules =
   Module.Name.Map.map modules ~f:(fun (m : Module.t) ->

test/blackbox-tests/test-cases/dune-ppx-driver-system/dune

@@ -66,3 +66,22 @@
  (name test_ppx_args)
  (modules test_ppx_args)
  (preprocess (pps -arg1 driver_print_args -arg2 -- -foo bar)))
+
+(library
+ (name driver_print_tool)
+ (modules ())
+ (libraries compiler-libs.common)
+ (ppx.driver (main "\| (fun () ->
+                   "\|    Ast_mapper.run_main (fun argv ->
+                   "\|      Printf.eprintf "tool name: %s\nargs:%s\n"
+                   "\|        (Ast_mapper.tool_name ())
+                   "\|        (String.concat " " argv);
+                   "\|      Ast_mapper.default_mapper))
+                   )))
+
+(rule (with-stdout-to test_ppx_staged.ml (echo "")))
+
+(library
+ (name test_ppx_staged)
+ (modules test_ppx_staged)
+ (preprocess (staged_pps -arg1 driver_print_tool -arg2 -- -foo bar)))

test/blackbox-tests/test-cases/dune-ppx-driver-system/dune-project

@@ -1 +1 @@
-(lang dune 1.0)
+(lang dune 1.1)

test/blackbox-tests/test-cases/dune-ppx-driver-system/run.t

@@ -53,6 +53,17 @@ Test the argument syntax
   test_ppx_args.pp.ml
   --impl
   test_ppx_args.ml
+  --as-ppx
   Error: Rule failed to generate the following targets:
   - test_ppx_args.pp.ml
   [1]
+
+Test that going throught the -ppx option of the compiler works
+
+  $ dune build test_ppx_staged.cma
+      ocamldep .test_ppx_staged.objs/test_ppx_staged.ml.d
+  tool name: ocamldep
+  args:--cookie library-name="test_ppx_staged"
+        ocamlc .test_ppx_staged.objs/test_ppx_staged.{cmi,cmo,cmt}
+  tool name: ocamlc
+  args:--cookie library-name="test_ppx_staged"