Preserve functor concrete syntax in the Parsetree (ocaml-ppx#2345)

Author: gpetiot

lib/Ast.ml

@@ -1350,7 +1350,7 @@ end = struct
          |Pexp_field (e, _)
          |Pexp_lazy e
          |Pexp_letexception (_, e)
-         |Pexp_letmodule (_, _, e)
+         |Pexp_letmodule (_, _, _, e)
          |Pexp_newtype (_, e)
          |Pexp_open (_, e)
          |Pexp_letopen (_, e)
@@ -1932,7 +1932,7 @@ end = struct
         | Pexp_let (_, e)
          |Pexp_letop {body= e; _}
          |Pexp_letexception (_, e)
-         |Pexp_letmodule (_, _, e) -> (
+         |Pexp_letmodule (_, _, _, e) -> (
           match cls with Match | Then | ThenElse -> continue e | _ -> false )
         | Pexp_match _ when match cls with Then -> true | _ -> false ->
             false
@@ -2000,7 +2000,7 @@ end = struct
       | Pexp_let (_, e)
        |Pexp_letop {body= e; _}
        |Pexp_letexception (_, e)
-       |Pexp_letmodule (_, _, e) ->
+       |Pexp_letmodule (_, _, _, e) ->
           continue e
       | Pexp_ifthenelse (eN, None) -> continue (List.last_exn eN).if_body
       | Pexp_extension (ext, PStr [{pstr_desc= Pstr_eval (e, _); _}])

lib/Extended_ast.ml

@@ -215,7 +215,42 @@ module Parse = struct
           {p with pexp_desc= Pexp_pack (name, Some pt)}
       | e -> Ast_mapper.default_mapper.expr m e
     in
-    Ast_mapper.{default_mapper with expr; pat; binding_op}
+    let module_expr (m : Ast_mapper.mapper) me =
+      let me =
+        match me with
+        (* [functor () -> functor () -> E] => [functor () () -> E] *)
+        | { pmod_desc=
+              Pmod_functor
+                ( args
+                , { pmod_desc= Pmod_functor (args', me')
+                  ; pmod_attributes= []
+                  ; _ } )
+          ; pmod_attributes= []
+          ; _ } ->
+            {me with pmod_desc= Pmod_functor (args @ args', me')}
+        | x -> x
+      in
+      Ast_mapper.default_mapper.module_expr m me
+    in
+    let module_type (m : Ast_mapper.mapper) mt =
+      let mt =
+        match mt with
+        (* [functor () -> functor () -> E] => [functor () () -> E] *)
+        | { pmty_desc=
+              Pmty_functor
+                ( args
+                , { pmty_desc= Pmty_functor (args', mt')
+                  ; pmty_attributes= []
+                  ; _ } )
+          ; pmty_attributes= []
+          ; _ } ->
+            {mt with pmty_desc= Pmty_functor (args @ args', mt')}
+        | x -> x
+      in
+      Ast_mapper.default_mapper.module_type m mt
+    in
+    Ast_mapper.
+      {default_mapper with expr; pat; binding_op; module_expr; module_type}
 
   let ast (type a) (fg : a t) ~preserve_beginend ~input_name str : a =
     map fg (normalize_mapper ~preserve_beginend)

lib/Fmt_ast.ml

@@ -158,14 +158,6 @@ let box_fun_sig_args c =
   | `Fit_or_vertical -> hvbox
   | `Wrap | `Smart -> hovbox
 
-let sugar_pmod_functor c ~for_functor_kw pmod =
-  let source_is_long = Source.is_long_pmod_functor c.source in
-  Sugar.functor_ c.cmts ~for_functor_kw ~source_is_long pmod
-
-let sugar_pmty_functor c ~for_functor_kw pmty =
-  let source_is_long = Source.is_long_pmty_functor c.source in
-  Sugar.functor_type c.cmts ~for_functor_kw ~source_is_long pmty
-
 let closing_paren ?force ?(offset = 0) c =
   match c.conf.fmt_opts.indicate_multiline_delimiters.v with
   | `No -> str ")"
@@ -2240,11 +2232,9 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
             $ fmt "@;<1000 0>"
             $ fmt_expression c (sub_exp ~ctx exp) )
         $ fmt_atrs )
-  | Pexp_letmodule (name, pmod, exp) ->
+  | Pexp_letmodule (name, args, pmod, exp) ->
       let keyword = "let module" in
-      let xargs, xbody =
-        sugar_pmod_functor c ~for_functor_kw:false (sub_mod ~ctx pmod)
-      in
+      let xbody = sub_mod ~ctx pmod in
       let xbody, xmty =
         match xbody.ast with
         | { pmod_desc= Pmod_constraint (body_me, body_mt)
@@ -2265,7 +2255,7 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
             (parens || not (List.is_empty pexp_attributes))
             c.conf
             ( hvbox 2
-                (fmt_module c ctx keyword ~eqty:":" name xargs (Some xbody)
+                (fmt_module c ctx keyword ~eqty:":" name args (Some xbody)
                    xmty [] ~epi:(str "in") ~can_sparse ?ext ~rec_flag:false )
             $ fmt "@;<1000 0>"
             $ fmt_expression c (sub_exp ~ctx exp) )
@@ -3447,17 +3437,15 @@ and fmt_module_type c ?(rec_ = false) ({ast= mty; _} as xmty) =
           Some
             ( str "end" $ after
             $ fmt_attributes_and_docstrings c pmty_attributes ) }
-  | Pmty_functor _ ->
-      let for_functor_kw = true in
-      let xargs, mt2 = sugar_pmty_functor c ~for_functor_kw xmty in
-      let blk = fmt_module_type c mt2 in
+  | Pmty_functor (args, mt) ->
+      let blk = fmt_module_type c (sub_mty ~ctx mt) in
       { blk with
         pro=
           Some
             ( str "functor"
             $ fmt_attributes c ~pre:Blank pmty_attributes
             $ fmt "@;<1 2>"
-            $ list xargs "@;<1 2>" (fmt_functor_param c ctx)
+            $ list args "@;<1 2>" (fmt_functor_param c ctx)
             $ fmt "@;<1 2>->"
             $ opt blk.pro (fun pro -> str " " $ pro) )
       ; epi= Some (fmt_opt blk.epi $ Cmts.fmt_after c pmty_loc)
@@ -3767,22 +3755,19 @@ and fmt_module c ctx ?rec_ ?ext ?epi ?(can_sparse = false) keyword
 and fmt_module_declaration ?ext c ~rec_flag ~first {ast= pmd; _} =
   protect c (Md pmd)
   @@
-  let {pmd_name; pmd_type; pmd_attributes; pmd_loc} = pmd in
+  let {pmd_name; pmd_args; pmd_type; pmd_attributes; pmd_loc} = pmd in
   update_config_maybe_disabled c pmd_loc pmd_attributes
   @@ fun c ->
   let ctx = Md pmd in
   let ext = if first then ext else None in
   let keyword = if first then "module" else "and" in
-  let xargs, xmty =
-    if rec_flag then ([], sub_mty ~ctx pmd_type)
-    else sugar_pmty_functor c ~for_functor_kw:false (sub_mty ~ctx pmd_type)
-  in
+  let xmty = sub_mty ~ctx pmd_type in
   let eqty =
     match xmty.ast.pmty_desc with Pmty_alias _ -> None | _ -> Some ":"
   in
   Cmts.fmt c pmd_loc
-    (fmt_module ~rec_:rec_flag ?ext c ctx keyword pmd_name xargs None ?eqty
-       (Some xmty) ~rec_flag:(rec_flag && first) pmd_attributes )
+    (fmt_module ~rec_:rec_flag ?ext c ctx keyword pmd_name pmd_args None
+       ?eqty (Some xmty) ~rec_flag:(rec_flag && first) pmd_attributes )
 
 and fmt_module_substitution ?ext c ctx pms =
   let {pms_name; pms_manifest; pms_attributes; pms_loc} = pms in
@@ -3995,8 +3980,7 @@ and fmt_module_expr ?(dock_struct = true) c ({ast= m; _} as xmod) =
           Option.some_if has_epi
             ( Cmts.fmt_after c pmod_loc
             $ fmt_attributes_and_docstrings c pmod_attributes ) }
-  | Pmod_functor _ ->
-      let xargs, me = sugar_pmod_functor c ~for_functor_kw:true xmod in
+  | Pmod_functor (args, me) ->
       let doc, atrs = doc_atrs pmod_attributes in
       { empty with
         bdy=
@@ -4007,10 +3991,11 @@ and fmt_module_expr ?(dock_struct = true) c ({ast= m; _} as xmod) =
                    ( str "functor"
                    $ fmt_attributes c ~pre:Blank atrs
                    $ fmt "@;<1 2>"
-                   $ list xargs "@;<1 2>" (fmt_functor_param c ctx)
+                   $ list args "@;<1 2>" (fmt_functor_param c ctx)
                    $ fmt "@;<1 2>->@;<1 2>"
-                   $ compose_module (fmt_module_expr c me) ~f:(hvbox 0) ) )
-            ) }
+                   $ compose_module
+                       (fmt_module_expr c (sub_mod ~ctx me))
+                       ~f:(hvbox 0) ) ) ) }
   | Pmod_ident lid ->
       { empty with
         opn= Some (open_hvbox 2)
@@ -4337,9 +4322,7 @@ and fmt_module_binding ?ext c ~rec_flag ~first {ast= pmb; _} =
   let ctx = Mb pmb in
   let ext = if first then ext else None in
   let keyword = if first then "module" else "and" in
-  let xargs, xbody =
-    sugar_pmod_functor c ~for_functor_kw:false (sub_mod ~ctx pmb.pmb_expr)
-  in
+  let xbody = sub_mod ~ctx pmb.pmb_expr in
   let xbody, xmty =
     match xbody.ast with
     | { pmod_desc= Pmod_constraint (body_me, body_mt)
@@ -4352,7 +4335,7 @@ and fmt_module_binding ?ext c ~rec_flag ~first {ast= pmb; _} =
   in
   Cmts.fmt c pmb.pmb_loc
     (fmt_module ~rec_:rec_flag ?ext c ctx keyword
-       ~rec_flag:(rec_flag && first) ~eqty:":" pmb.pmb_name xargs
+       ~rec_flag:(rec_flag && first) ~eqty:":" pmb.pmb_name pmb.pmb_args
        (Some xbody) xmty pmb.pmb_attributes )
 
 let fmt_toplevel_directive c ~semisemi dir =

lib/Source.ml

@@ -104,29 +104,6 @@ let extend_loc_to_include_attributes (loc : Location.t) (l : attributes) =
   else
     {loc with loc_end= {loc.loc_end with pos_cnum= loc_end.loc_end.pos_cnum}}
 
-let is_long_functor_syntax (t : t) ~(from : Location.t) = function
-  | Unit -> false
-  | Named ({loc= _; _}, _) -> (
-    (* since 4.12 the functor keyword is just before the loc of the functor
-       parameter *)
-    match
-      find_token_before t
-        ~filter:(function COMMENT _ | DOCSTRING _ -> false | _ -> true)
-        from.loc_start
-    with
-    | Some (Parser.FUNCTOR, _) -> true
-    | _ -> false )
-
-let is_long_pmod_functor t {pmod_desc; pmod_loc= from; _} =
-  match pmod_desc with
-  | Pmod_functor (fp, _) -> is_long_functor_syntax t ~from fp
-  | _ -> false
-
-let is_long_pmty_functor t {pmty_desc; pmty_loc= from; _} =
-  match pmty_desc with
-  | Pmty_functor (fp, _) -> is_long_functor_syntax t ~from fp
-  | _ -> false
-
 let string_literal t mode loc =
   Option.value_exn ~message:"Parse error while reading string literal"
     (Literal_lexer.string mode (string_at t loc))

lib/Source.mli

@@ -47,16 +47,6 @@ val find_token_before :
 
 val string_literal : t -> [`Normalize | `Preserve] -> Location.t -> string
 
-val is_long_pmod_functor : t -> module_expr -> bool
-(** [is_long_pmod_functor source mod_exp] holds if [mod_exp] is a
-    [Pmod_functor] expression that is expressed in long ('functor (M) ->')
-    form in source. *)
-
-val is_long_pmty_functor : t -> module_type -> bool
-(** [is_long_pmty_functor source mod_type] holds if [mod_type] is a
-    [Pmty_functor] type that is expressed in long ('functor (M) ->') form in
-    source. *)
-
 val begins_line : ?ignore_spaces:bool -> t -> Location.t -> bool
 
 val ends_line : t -> Location.t -> bool

lib/Sugar.ml

@@ -163,41 +163,6 @@ let sequence cmts xexp =
   in
   sequence_ xexp
 
-(* The sugar is different when used with the [functor] keyword. The syntax
-   M(A : A)(B : B) cannot handle [_] as module name. *)
-let rec functor_type cmts ~for_functor_kw ~source_is_long
-    ({ast= mty; _} as xmty) =
-  let ctx = Mty mty in
-  match mty with
-  | {pmty_desc= Pmty_functor (fp, body); pmty_loc; pmty_attributes}
-    when for_functor_kw
-         || (List.is_empty pmty_attributes && not (source_is_long mty)) ->
-      let body = sub_mty ~ctx body in
-      let xargs, xbody =
-        match pmty_attributes with
-        | [] -> functor_type cmts ~for_functor_kw ~source_is_long body
-        | _ -> ([], body)
-      in
-      (Location.mkloc fp pmty_loc :: xargs, xbody)
-  | _ -> ([], xmty)
-
-(* The sugar is different when used with the [functor] keyword. The syntax
-   M(A : A)(B : B) cannot handle [_] as module name. *)
-let rec functor_ cmts ~for_functor_kw ~source_is_long ({ast= me; _} as xme) =
-  let ctx = Mod me in
-  match me with
-  | {pmod_desc= Pmod_functor (fp, body); pmod_loc; pmod_attributes}
-    when for_functor_kw
-         || (List.is_empty pmod_attributes && not (source_is_long me)) ->
-      let body = sub_mod ~ctx body in
-      let xargs, xbody_me =
-        match pmod_attributes with
-        | [] -> functor_ cmts ~for_functor_kw ~source_is_long body
-        | _ -> ([], body)
-      in
-      (Location.mkloc fp pmod_loc :: xargs, xbody_me)
-  | _ -> ([], xme)
-
 let mod_with pmty =
   let rec mod_with_ ({ast= me; _} as xme) =
     let ctx = Mty me in

lib/Sugar.mli

@@ -51,26 +51,6 @@ val sequence :
 (** [sequence cmts exp] returns the list of expressions (with the optional
     extension) from a sequence of expressions [exp]. *)
 
-val functor_type :
-     Cmts.t
-  -> for_functor_kw:bool
-  -> source_is_long:(module_type -> bool)
-  -> module_type Ast.xt
-  -> functor_parameter loc list * module_type Ast.xt
-(** [functor_type cmts for_functor_kw m] returns the list of module types
-    applied to the functor of module type [m]. [for_functor_kw] indicates if
-    the keyword [functor] is used. *)
-
-val functor_ :
-     Cmts.t
-  -> for_functor_kw:bool
-  -> source_is_long:(module_expr -> bool)
-  -> module_expr Ast.xt
-  -> functor_parameter loc list * module_expr Ast.xt
-(** [functor_ cmts for_functor_kw m] returns the list of module types applied
-    to the functor of module [m]. [for_functor_kw] indicates if the keyword
-    [functor] is used. *)
-
 val mod_with :
      module_type Ast.xt
   -> (with_constraint list * Warnings.loc * attributes) list

test/passing/tests/comments.ml.ref

@@ -320,10 +320,7 @@ let y =
     (* b *)
     x
 
-module A (* A *) () =
-  (* B *)
-  (* C *)
-  B
+module A (* A *) () (* B *) = (* C *) B
 
 let kk = (* foo *) (module A : T)
 

vendor/parser-extended/ast_helper.ml

@@ -145,7 +145,7 @@ module Exp = struct
   let indexop_access ?loc ?attrs pia_lhs pia_kind pia_paren pia_rhs =
     mk ?loc ?attrs (Pexp_indexop_access {pia_lhs; pia_kind; pia_paren; pia_rhs})
   let override ?loc ?attrs a = mk ?loc ?attrs (Pexp_override a)
-  let letmodule ?loc ?attrs a b c= mk ?loc ?attrs (Pexp_letmodule (a, b, c))
+  let letmodule ?loc ?attrs a b c d = mk ?loc ?attrs (Pexp_letmodule (a, b, c, d))
   let letexception ?loc ?attrs a b = mk ?loc ?attrs (Pexp_letexception (a, b))
   let assert_ ?loc ?attrs a = mk ?loc ?attrs (Pexp_assert a)
   let lazy_ ?loc ?attrs a = mk ?loc ?attrs (Pexp_lazy a)
@@ -367,9 +367,10 @@ end
 
 module Md = struct
   let mk ?(loc = !default_loc) ?(attrs = [])
-        ?(docs = empty_docs) ?(text = []) name typ =
+        ?(docs = empty_docs) ?(text = []) name args typ =
     {
      pmd_name = name;
+     pmd_args = args;
      pmd_type = typ;
      pmd_attributes =
        add_text_attrs text (add_docs_attrs docs attrs);
@@ -403,9 +404,10 @@ end
 
 module Mb = struct
   let mk ?(loc = !default_loc) ?(attrs = [])
-        ?(docs = empty_docs) ?(text = []) name expr =
+        ?(docs = empty_docs) ?(text = []) name args expr =
     {
      pmb_name = name;
+     pmb_args = args;
      pmb_expr = expr;
      pmb_attributes =
        add_text_attrs text (add_docs_attrs docs attrs);