Add support for extensions in module types

ocamldoc/Makefile

@@ -74,7 +74,7 @@ INCLUDES_NODEP=\
 DEPINCLUDES=$(INCLUDES_DEP)
 INCLUDES=$(INCLUDES_DEP) $(INCLUDES_NODEP)
 
-COMPFLAGS=$(INCLUDES) -absname -w +a-4-9-41-42-44-45-48-70 -warn-error +A \
+COMPFLAGS=$(INCLUDES) -absname -w +a-4-9-40-41-42-44-45-48-70 -warn-error +A \
   -safe-string -strict-sequence -strict-formats -bin-annot -principal
 
 LINKFLAGS=$(INCLUDES) -nostdlib

ocamldoc/odoc_sig.ml

@@ -1520,6 +1520,9 @@ module Analyser =
     (** Return a module_type_kind from a Parsetree.module_type and a Types.module_type *)
     and analyse_module_type_kind
       ?(erased = Name.Map.empty) env current_module_name module_type sig_module_type =
+      match Extensions.Module_type.of_ast module_type with
+      | Some _ -> .
+      | None ->
       match module_type.Parsetree.pmty_desc with
         Parsetree.Pmty_ident longident ->
           let name =
@@ -1618,6 +1621,9 @@ module Analyser =
     (** analyse of a Parsetree.module_type and a Types.module_type.*)
     and analyse_module_kind
         ?(erased = Name.Map.empty) env current_module_name module_type sig_module_type =
+      match Extensions.Module_type.of_ast module_type with
+      | Some _ -> .
+      | None ->
       match module_type.Parsetree.pmty_desc with
       | Parsetree.Pmty_ident _longident ->
           let k = analyse_module_type_kind env current_module_name module_type sig_module_type in

parsing/ast_iterator.ml

@@ -54,6 +54,7 @@ type iterator = {
   module_expr: iterator -> module_expr -> unit;
   module_type: iterator -> module_type -> unit;
   module_type_declaration: iterator -> module_type_declaration -> unit;
+  module_type_extension: iterator -> Extensions.Module_type.t -> unit;
   open_declaration: iterator -> open_declaration -> unit;
   open_description: iterator -> open_description -> unit;
   pat: iterator -> pattern -> unit;
@@ -246,9 +247,13 @@ let iter_functor_param sub = function
 module MT = struct
   (* Type expressions for the module language *)
 
-  let iter sub {pmty_desc = desc; pmty_loc = loc; pmty_attributes = attrs} =
+  let iter sub
+        ({pmty_desc = desc; pmty_loc = loc; pmty_attributes = attrs} as mty) =
     sub.location sub loc;
     sub.attributes sub attrs;
+    match Extensions.Module_type.of_ast mty with
+    | Some emty -> sub.module_type_extension sub emty
+    | None ->
     match desc with
     | Pmty_ident s -> iter_loc sub s
     | Pmty_alias s -> iter_loc sub s
@@ -652,6 +657,9 @@ let default_iterator =
          this.attributes this pmtd_attributes;
       );
 
+    module_type_extension = (fun _this emty -> match emty with
+      | _ -> .);
+
     module_binding =
       (fun this {pmb_name; pmb_expr; pmb_attributes; pmb_loc} ->
          iter_loc this pmb_name; this.module_expr this pmb_expr;

parsing/ast_iterator.mli

@@ -57,6 +57,7 @@ type iterator = {
   module_expr: iterator -> module_expr -> unit;
   module_type: iterator -> module_type -> unit;
   module_type_declaration: iterator -> module_type_declaration -> unit;
+  module_type_extension: iterator -> Extensions.Module_type.t -> unit;
   open_declaration: iterator -> open_declaration -> unit;
   open_description: iterator -> open_description -> unit;
   pat: iterator -> pattern -> unit;

parsing/ast_mapper.ml

@@ -61,6 +61,8 @@ type mapper = {
   module_type: mapper -> module_type -> module_type;
   module_type_declaration: mapper -> module_type_declaration
                            -> module_type_declaration;
+  module_type_extension: mapper
+    -> Extensions.Module_type.t -> Extensions.Module_type.t;
   open_declaration: mapper -> open_declaration -> open_declaration;
   open_description: mapper -> open_description -> open_description;
   pat: mapper -> pattern -> pattern;
@@ -274,10 +276,18 @@ let map_functor_param sub = function
 module MT = struct
   (* Type expressions for the module language *)
 
-  let map sub {pmty_desc = desc; pmty_loc = loc; pmty_attributes = attrs} =
+  let map sub
+        ({pmty_desc = desc; pmty_loc = loc; pmty_attributes = attrs} as mty) =
     let open Mty in
     let loc = sub.location sub loc in
     let attrs = sub.attributes sub attrs in
+    match Extensions.Module_type.of_ast mty with
+    | Some emty -> begin
+        Extensions_parsing.Module_type.wrap_desc ~loc ~attrs @@
+        match sub.module_type_extension sub emty with
+        | _ -> .
+      end
+    | None ->
     match desc with
     | Pmty_ident s -> ident ~loc ~attrs (map_loc sub s)
     | Pmty_alias s -> alias ~loc ~attrs (map_loc sub s)
@@ -720,6 +730,9 @@ let default_mapper =
            ~loc:(this.location this pmtd_loc)
       );
 
+    module_type_extension =
+      (fun _this emty -> match emty with _ -> .);
+
     module_binding =
       (fun this {pmb_name; pmb_expr; pmb_attributes; pmb_loc} ->
          Mb.mk (map_loc this pmb_name) (this.module_expr this pmb_expr)

parsing/ast_mapper.mli

@@ -90,6 +90,8 @@ type mapper = {
   module_type: mapper -> module_type -> module_type;
   module_type_declaration: mapper -> module_type_declaration
                            -> module_type_declaration;
+  module_type_extension: mapper ->
+    Extensions.Module_type.t -> Extensions.Module_type.t;
   open_declaration: mapper -> open_declaration -> open_declaration;
   open_description: mapper -> open_description -> open_description;
   pat: mapper -> pattern -> pattern;

parsing/depend.ml

@@ -340,6 +340,9 @@ and add_binding_op bv bv' pbop =
   add_pattern bv' pbop.pbop_pat
 
 and add_modtype bv mty =
+  match Extensions.Module_type.of_ast mty with
+  | Some _ -> .
+  | None ->
   match mty.pmty_desc with
     Pmty_ident l -> add bv l
   | Pmty_alias l -> add_module_path bv l
@@ -382,6 +385,9 @@ and add_module_alias bv l =
     | _ -> add_module_path bv l; bound (* cannot delay *)
 
 and add_modtype_binding bv mty =
+  match Extensions.Module_type.of_ast mty with
+  | Some _ -> .
+  | None ->
   match mty.pmty_desc with
     Pmty_alias l ->
       add_module_alias bv l

parsing/extensions.ml

@@ -345,3 +345,17 @@ module Pattern = struct
   include M
   include Make_of_ast(M)
 end
+
+module Module_type = struct
+  module M = struct
+    module AST = Extensions_parsing.Module_type
+
+    type t = |
+
+    let of_ast_internal (ext : Language_extension.t) _mty = match ext with
+      | _ -> None
+  end
+
+  include M
+  include Make_of_ast(M)
+end

parsing/extensions.mli

@@ -151,3 +151,10 @@ module Pattern : sig
 
   include AST with type t := t and type ast := Parsetree.pattern
 end
+
+(** Language extensions in module types *)
+module Module_type : sig
+  type t = |
+
+  include AST with type t := t and type ast := Parsetree.module_type
+end

parsing/extensions_parsing.ml

@@ -277,6 +277,30 @@ module Pattern = Make_AST(struct
        None
 end)
 
+(** Module types; embedded as [functor (_ : [%extension.EXTNAME]) -> BODY]. *)
+module Module_type = Make_AST(struct
+    type ast = module_type
+    type ast_desc = module_type_desc
+
+    let plural = "module types"
+
+    let location mty = mty.pmty_loc
+
+    let wrap_desc ~loc ~attrs = Ast_helper.Mty.mk ~loc ~attrs
+
+    let make_extension_node = Ast_helper.Mty.extension
+
+    let make_extension_use ~extension_node mty =
+      Pmty_functor(Named(Location.mknoloc None, extension_node), mty)
+
+    let match_extension_use mty =
+      match mty.pmty_desc with
+      | Pmty_functor(Named({txt = None},
+                           {pmty_desc = Pmty_extension ext}), mty) ->
+        Some (ext, mty)
+      | _ -> None
+end)
+
 (******************************************************************************)
 (** Generically lift and lower our custom language extension ASTs from/to OCaml
     ASTs. *)

parsing/extensions_parsing.mli

@@ -131,10 +131,12 @@ end
     adding these lazily as we need them. When you add another one, make
     sure also to add special handling in [Ast_iterator] and [Ast_mapper]. *)
 
-module Expression : AST with type ast      = Parsetree.expression
-                         and type ast_desc = Parsetree.expression_desc
-module Pattern    : AST with type ast      = Parsetree.pattern
-                         and type ast_desc = Parsetree.pattern_desc
+module Expression  : AST with type ast      = Parsetree.expression
+                          and type ast_desc = Parsetree.expression_desc
+module Pattern     : AST with type ast      = Parsetree.pattern
+                          and type ast_desc = Parsetree.pattern_desc
+module Module_type : AST with type ast      = Parsetree.module_type
+                          and type ast_desc = Parsetree.module_type_desc
 
 (** Each syntactic category will include a module that meets this signature.
     Then, the [Make_of_ast] functor produces the functions that actually

parsing/parsetree.mli

@@ -816,6 +816,14 @@ and class_declaration = class_expr class_infos
 and module_type =
     {
      pmty_desc: module_type_desc;
+     (** (Jane Street specific; delete when upstreaming.)
+         Consider using [Extensions.Module_type.of_ast] before matching on
+         this field directly, as the former will detect extension nodes
+         correctly. Extensions are encoded as
+         [Pmty_functor(Named(_, Pmty_extension _), _)];
+         if your pattern match avoids
+         matching that pattern, it is OK to skip [of_ast]. *)
+
      pmty_loc: Location.t;
      pmty_attributes: attributes;  (** [... [\@id1] [\@id2]] *)
     }

parsing/pprintast.ml

@@ -1118,6 +1118,9 @@ and module_type ctxt f x =
     pp f "((%a)%a)" (module_type ctxt) {x with pmty_attributes=[]}
       (attributes ctxt) x.pmty_attributes
   end else
+    match Extensions.Module_type.of_ast x with
+    | Some _ -> .
+    | None ->
     match x.pmty_desc with
     | Pmty_functor (Unit, mt2) ->
         pp f "@[<hov2>functor () ->@ %a@]" (module_type ctxt) mt2
@@ -1161,7 +1164,10 @@ and with_constraint ctxt f = function
 
 and module_type1 ctxt f x =
   if x.pmty_attributes <> [] then module_type ctxt f x
-  else match x.pmty_desc with
+  else match Extensions.Module_type.of_ast x with
+    | Some _ -> .
+    | None ->
+    match x.pmty_desc with
     | Pmty_ident li ->
         pp f "%a" longident_loc li;
     | Pmty_alias li ->

parsing/printast.ml

@@ -669,6 +669,7 @@ and module_type i ppf x =
   line i ppf "module_type %a\n" fmt_location x.pmty_loc;
   attributes i ppf x.pmty_attributes;
   let i = i+1 in
+  (* Print raw AST, without interpreting extensions *)
   match x.pmty_desc with
   | Pmty_ident li -> line i ppf "Pmty_ident %a\n" fmt_longident_loc li;
   | Pmty_alias li -> line i ppf "Pmty_alias %a\n" fmt_longident_loc li;

tools/lintapidiff.ml

@@ -123,6 +123,9 @@ module Ast = struct
   let rec add_item ~f path inherits map item =
     let rec add_module_type path ty (inherits, map) =
       let self = add_item ~f path inherits in
+      match Extensions.Module_type.of_ast ty with
+      | Some _ -> .
+      | None ->
       match ty.pmty_desc with
       | Pmty_signature lst -> List.fold_left self map lst
       | Pmty_functor ({txt;_}, _, m) ->

typing/typemod.ml

@@ -834,6 +834,9 @@ let map_ext fn exts =
    making them abstract otherwise. *)
 
 let rec approx_modtype env smty =
+  match Extensions.Module_type.of_ast smty with
+  | Some _ -> .
+  | None ->
   match smty.pmty_desc with
     Pmty_ident lid ->
       let path =
@@ -1369,6 +1372,9 @@ and transl_modtype_functor_arg env sarg =
 
 and transl_modtype_aux env smty =
   let loc = smty.pmty_loc in
+  match Extensions.Module_type.of_ast smty with
+  | Some _ -> .
+  | None ->
   match smty.pmty_desc with
     Pmty_ident lid ->
       let path = transl_modtype_longident loc env lid.txt in