Share some computations on pps when preprocessing modules

src/import.ml

@@ -141,6 +141,8 @@ module Hashtbl = struct
 end
 
 module Map = struct
+  module type OrderedType = MoreLabels.Map.OrderedType
+
   module type S = sig
     include MoreLabels.Map.S
 
@@ -155,7 +157,7 @@ module Map = struct
     val values : 'a t -> 'a list
   end
 
-  module Make(Key : MoreLabels.Map.OrderedType) : S with type key = Key.t = struct
+  module Make(Key : OrderedType) : S with type key = Key.t = struct
     include MoreLabels.Map.Make(Key)
 
     let add_multi t ~key ~data =

src/jbuild.ml

@@ -264,46 +264,39 @@ module Preprocess = struct
 end
 
 module Per_module = struct
-  type 'a t =
+  include Per_item.Make(String)
-    | For_all    of 'a
-    | Per_module of 'a String_map.t
 
-  let t a sexp =
+  let t ~default a sexp =
     match sexp with
     | List (_, Atom (_, "per_module") :: rest) -> begin
-        List.concat_map rest ~f:(fun sexp ->
+      List.map rest ~f:(fun sexp ->
-          let pp, names = pair a module_names sexp in
+        let pp, names = pair a module_names sexp in
-          List.map (String_set.elements names) ~f:(fun name -> (name, pp)))
+        (String_set.elements names, pp))
-        |> String_map.of_alist
+      |> of_mapping ~default
-        |> function
+      |> function
-        | Ok map -> Per_module map
+      | Ok t -> t
-        | Error (name, _, _) ->
+      | Error (name, _, _) ->
-          of_sexp_error sexp (sprintf "module %s present in two different sets" name)
+        of_sexp_error sexp (sprintf "module %s present in two different sets" name)
-      end
+    end
-    | sexp -> For_all (a sexp)
+    | sexp -> for_all (a sexp)
 end
 
 module Preprocess_map = struct
   type t = Preprocess.t Per_module.t
-  let t = Per_module.t Preprocess.t
+  let t = Per_module.t Preprocess.t ~default:Preprocess.No_preprocessing
 
-  let no_preprocessing = Per_module.For_all Preprocess.No_preprocessing
+  let no_preprocessing = Per_module.for_all Preprocess.No_preprocessing
 
-  let find module_name (t : t) =
+  let find module_name t = Per_module.get t module_name
-    match t with
-    | For_all pp -> pp
-    | Per_module map -> String_map.find_default module_name map ~default:No_preprocessing
 
-  let default : t = For_all No_preprocessing
+  let default = Per_module.for_all Preprocess.No_preprocessing
 
   module Pp_set = Set.Make(Pp)
 
-  let pps : t -> _ = function
+  let pps t =
-    | For_all pp -> Preprocess.pps pp
+    Per_module.fold t ~init:Pp_set.empty ~f:(fun pp acc ->
-    | Per_module map ->
+      Pp_set.union acc (Pp_set.of_list (Preprocess.pps pp)))
-      String_map.fold map ~init:Pp_set.empty ~f:(fun ~key:_ ~data:pp acc ->
+    |> Pp_set.elements
-        Pp_set.union acc (Pp_set.of_list (Preprocess.pps pp)))
-      |> Pp_set.elements
 end
 
 module Lint = struct
@@ -498,9 +491,10 @@ module Buildable = struct
       }
 
   let single_preprocess t =
-    match t.preprocess with
+    if Per_module.is_constant t.preprocess then
-    | For_all pp -> pp
+      Per_module.get t.preprocess ""
-    | Per_module _ -> No_preprocessing
+    else
+      Preprocess.No_preprocessing
 end
 
 module Public_lib = struct

src/jbuild.mli

@@ -45,13 +45,16 @@ module Preprocess : sig
     | Pps    of pps
 end
 
+module Per_module : Per_item.S with type key = string
+
 module Preprocess_map : sig
-  type t
+  type t = Preprocess.t Per_module.t
 
   val no_preprocessing : t
   val default : t
 
-  (** [find module_name] find the preprocessing specification for a given module *)
+  (** [find module_name] find the preprocessing specification for a
+      given module *)
   val find : string -> t -> Preprocess.t
 
   val pps : t -> Pp.t list

src/per_item.ml

@@ -0,0 +1,63 @@
+open Import
+
+module type S = sig
+  type key
+
+  type 'a t
+
+  val for_all : 'a -> 'a t
+  val of_mapping
+    :  (key list * 'a) list
+    -> default:'a
+    -> ('a t, key * 'a * 'a) result
+  val get : 'a t -> key -> 'a
+  val is_constant : _ t -> bool
+  val map : 'a t -> f:('a -> 'b) -> 'b t
+  val fold : 'a t -> init:'acc -> f:('a -> 'acc -> 'acc) -> 'acc
+end
+
+module Make(Key : Map.OrderedType) : S with type key = Key.t = struct
+  module Map = Map.Make(Key)
+
+  type key = Key.t
+
+  type 'a t =
+    { map    : int Map.t
+    ; values : 'a array
+    }
+
+  let for_all x =
+    { map    = Map.empty
+    ; values = [|x|]
+    }
+
+  let of_mapping l ~default =
+    let values =
+      Array.of_list (default :: List.map l ~f:snd)
+    in
+    List.mapi l ~f:(fun i (keys, _) ->
+      List.map keys ~f:(fun key -> (key, i + 1)))
+    |> List.concat
+    |> Map.of_alist
+    |> function
+    | Ok map ->
+      Ok { map; values }
+    | Error (key, x, y) ->
+      Error (key, values.(x), values.(y))
+
+  let get t key =
+    let index =
+      match Map.find key t.map with
+      | None   -> 0
+      | Some i -> i
+    in
+    t.values.(index)
+
+  let map t ~f =
+    { t with values = Array.map t.values ~f }
+
+  let fold t ~init ~f =
+    Array.fold_right t.values ~init ~f
+
+  let is_constant t = Array.length t.values = 1
+end

src/per_item.mli

@@ -0,0 +1,36 @@
+(** Module used to represent the [(per_xxx ...)] forms
+
+    The main different between this module and a plain [Map] is that
+    the [map] operation applies transformations only once per distinct
+    value.
+*)
+
+open Import
+
+module type S = sig
+  type key
+
+  type 'a t
+
+  (** Create a mapping where all keys map to the same value *)
+  val for_all : 'a -> 'a t
+
+  (** Create a mapping from a list of bindings *)
+  val of_mapping
+    :  (key list * 'a) list
+    -> default:'a
+    -> ('a t, key * 'a * 'a) result
+
+  (** Get the configuration for the given item *)
+  val get : 'a t -> key -> 'a
+
+  (** Returns [true] if the mapping returns the same value for all
+      keys. Note that the mapping might still be constant if
+      [is_constant] returns [false]. *)
+  val is_constant : _ t -> bool
+
+  val map : 'a t -> f:('a -> 'b) -> 'b t
+  val fold : 'a t -> init:'acc -> f:('a -> 'acc -> 'acc) -> 'acc
+end
+
+module Make(Key : Map.OrderedType) : S with type key = Key.t

src/super_context.ml

@@ -960,8 +960,7 @@ module PP = struct
                (Path.extend_basename fn ~suffix:".ppx-corrected"))
         ]
 
-  let lint_module sctx ~(source : Module.t) ~(ast : Module.t) ~dir
+  let lint_module sctx ~dir ~dep_kind ~lint ~lib_name ~scope = Staged.stage (
-        ~dep_kind ~lint ~lib_name ~scope =
     let alias = Alias.lint ~dir in
     let add_alias fn build =
       Alias.add_action sctx.build_system alias build
@@ -970,50 +969,58 @@ module PP = struct
                      ; Atom fn
                      ])
     in
-    match Preprocess_map.find source.name lint with
+    let lint =
-    | No_preprocessing -> ()
+      Per_module.map lint ~f:(function
-    | Action action ->
+        | Preprocess.No_preprocessing ->
-      let action = Action.U.Chdir (root_var, action) in
+          (fun ~source:_ ~ast:_ -> ())
-      Module.iter source ~f:(fun _ (src : Module.File.t) ->
+        | Action action ->
-        let src_path = Path.relative dir src.name in
+          (fun ~source ~ast:_ ->
-        add_alias src.name
+             let action = Action.U.Chdir (root_var, action) in
-          (Build.path src_path
+             Module.iter source ~f:(fun _ (src : Module.File.t) ->
-           >>^ (fun _ -> [src_path])
+               let src_path = Path.relative dir src.name in
-           >>> Action.run sctx
+               add_alias src.name
-                 action
+                 (Build.path src_path
-                 ~dir
+                  >>^ (fun _ -> [src_path])
-                 ~dep_kind
+                  >>> Action.run sctx
-                 ~targets:(Static [])
+                        action
-                 ~scope)
+                        ~dir
-      )
+                        ~dep_kind
-    | Pps { pps; flags } ->
+                        ~targets:(Static [])
-      let ppx_exe = get_ppx_driver sctx ~scope pps in
+                        ~scope)))
-      Module.iter ast ~f:(fun kind src ->
+        | Pps { pps; flags } ->
-        let src_path = Path.relative dir src.name in
+          let ppx_exe = get_ppx_driver sctx ~scope pps in
-        let args =
+          let uses_ppx_driver = uses_ppx_driver ~pps in
-          [ Arg_spec.As flags
+          let args : _ Arg_spec.t =
-          ; As (cookie_library_name lib_name)
+            S [ As flags
-          ; Ml_kind.ppx_driver_flag kind
+              ; As (cookie_library_name lib_name)
-          ; Dep src_path
+              (* This hack is needed until -null is standard:
-          ]
+                 https://github.com/ocaml-ppx/ocaml-migrate-parsetree/issues/35
-        in
+              *)
-        let uses_ppx_driver = uses_ppx_driver ~pps in
+              ; As (if uses_ppx_driver then
-        let args =
+                      [ "-null"; "-diff-cmd"; "-" ]
-          (* This hack is needed until -null is standard:
+                    else
-             https://github.com/ocaml-ppx/ocaml-migrate-parsetree/issues/35 *)
+                      [])
-          if uses_ppx_driver then
+              ]
-            args @ [ A "-null"; A "-diff-cmd"; A "-" ]
+          in
-          else
+          (fun ~source ~ast ->
-            args
+             Module.iter ast ~f:(fun kind src ->
-        in
+               let args =
-        add_alias src.name
+                 [ args
-          (promote_correction ~uses_ppx_driver
+                 ; Ml_kind.ppx_driver_flag kind
-             (Option.value_exn (Module.file ~dir source kind))
+                 ; Dep (Path.relative dir src.name)
-             (Build.run ~context:sctx.context (Ok ppx_exe) args))
+                 ]
-      )
+               in
+               add_alias src.name
+                 (promote_correction ~uses_ppx_driver
+                    (Option.value_exn (Module.file ~dir source kind))
+                    (Build.run ~context:sctx.context (Ok ppx_exe) args))
+             )))
+    in
+    fun ~(source : Module.t) ~ast ->
+      Per_module.get lint source.name ~source ~ast)
 
-  (* Generate rules to build the .pp files and return a new module map where all filenames
+  (* Generate rules to build the .pp files and return a new module map
-     point to the .pp files *)
+     where all filenames point to the .pp files *)
   let pp_and_lint_modules sctx ~dir ~dep_kind ~modules ~lint ~preprocess
         ~preprocessor_deps ~lib_name
         ~(scope : Lib_db.Scope.t With_required_by.t) =
@@ -1021,56 +1028,67 @@ module PP = struct
       Build.memoize "preprocessor deps"
         (Deps.interpret sctx ~scope:scope.data ~dir preprocessor_deps)
     in
-    let lint_module = lint_module sctx ~dir ~dep_kind ~lint ~lib_name ~scope in
+    let lint_module =
+      Staged.unstage (lint_module sctx ~dir ~dep_kind ~lint ~lib_name ~scope)
+    in
+    let preprocess =
+      Per_module.map preprocess ~f:(function
+        | Preprocess.No_preprocessing ->
+          (fun m ->
+             let ast = setup_reason_rules sctx ~dir m in
+             lint_module ~ast ~source:m;
+             ast)
+        | Action action ->
+          (fun m ->
+             let ast =
+               pped_module m ~dir ~f:(fun _kind src dst ->
+                 add_rule sctx
+                   (preprocessor_deps
+                    >>>
+                    Build.path src
+                    >>^ (fun _ -> [src])
+                    >>>
+                    Action.run sctx
+                      (Redirect
+                         (Stdout,
+                          target_var,
+                          Chdir (root_var,
+                                 action)))
+                      ~dir
+                      ~dep_kind
+                      ~targets:(Static [dst])
+                      ~scope))
+               |> setup_reason_rules sctx ~dir in
+             lint_module ~ast ~source:m;
+             ast)
+        | Pps { pps; flags } ->
+          let ppx_exe = get_ppx_driver sctx ~scope pps in
+          let uses_ppx_driver = uses_ppx_driver ~pps in
+          let args : _ Arg_spec.t =
+            S [ As flags
+              ; A "--dump-ast"
+              ; As (cookie_library_name lib_name)
+              ; As (if uses_ppx_driver then ["-diff-cmd"; "-"] else [])
+              ]
+          in
+          (fun m ->
+             let ast = setup_reason_rules sctx ~dir m in
+             lint_module ~ast ~source:m;
+             pped_module ast ~dir ~f:(fun kind src dst ->
+               add_rule sctx
+                 (promote_correction ~uses_ppx_driver
+                    (Option.value_exn (Module.file m ~dir kind))
+                    (preprocessor_deps
+                     >>>
+                     Build.run ~context:sctx.context
+                       (Ok ppx_exe)
+                       [ args
+                       ; A "-o"; Target dst
+                       ; Ml_kind.ppx_driver_flag kind; Dep src
+                       ])))))
+    in
     String_map.map modules ~f:(fun (m : Module.t) ->
-      match Preprocess_map.find m.name preprocess with
+      Per_module.get preprocess m.name m)
-      | No_preprocessing ->
-        let ast = setup_reason_rules sctx ~dir m in
-        lint_module ~ast ~source:m;
-        ast
-      | Action action ->
-        let ast =
-          pped_module m ~dir ~f:(fun _kind src dst ->
-            add_rule sctx
-              (preprocessor_deps
-               >>>
-               Build.path src
-               >>^ (fun _ -> [src])
-               >>>
-               Action.run sctx
-                 (Redirect
-                    (Stdout,
-                     target_var,
-                     Chdir (root_var,
-                            action)))
-                 ~dir
-                 ~dep_kind
-                 ~targets:(Static [dst])
-                 ~scope))
-          |> setup_reason_rules sctx ~dir in
-        lint_module ~ast ~source:m;
-        ast
-      | Pps { pps; flags } ->
-        let ppx_exe = get_ppx_driver sctx ~scope pps in
-        let ast = setup_reason_rules sctx ~dir m in
-        lint_module ~ast ~source:m;
-        let uses_ppx_driver = uses_ppx_driver ~pps in
-        pped_module ast ~dir ~f:(fun kind src dst ->
-          add_rule sctx
-            (promote_correction ~uses_ppx_driver
-               (Option.value_exn (Module.file m ~dir kind))
-               (preprocessor_deps
-                >>>
-                Build.run ~context:sctx.context
-                  (Ok ppx_exe)
-                  [ As flags
-                  ; A "--dump-ast"
-                  ; As (cookie_library_name lib_name)
-                  ; As (if uses_ppx_driver then ["-diff-cmd"; "-"] else [])
-                  ; A "-o"; Target dst
-                  ; Ml_kind.ppx_driver_flag kind; Dep src
-                  ])))
-    )
 end
 
 module Eval_strings = Ordered_set_lang.Make(struct