Add Strengthen extension for module types

Author: goldfirere

ocamldoc/odoc_sig.ml

@@ -1521,7 +1521,7 @@ module Analyser =
     and analyse_module_type_kind
       ?(erased = Name.Map.empty) env current_module_name module_type sig_module_type =
       match Extensions.Module_type.get_desc module_type with
-      | Extension _ -> .
+      | Extension (Emty_strengthen _) -> failwith "strengthen not implemented yet"
       | Regular desc -> match desc with
         Parsetree.Pmty_ident longident ->
           let name =
@@ -1621,7 +1621,7 @@ module Analyser =
     and analyse_module_kind
         ?(erased = Name.Map.empty) env current_module_name module_type sig_module_type =
       match Extensions.Module_type.get_desc module_type with
-      | Extension _ -> .
+      | Extension (Emty_strengthen _) -> failwith "strengthen not implemented yet"
       | Regular desc -> match 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

@@ -303,6 +303,12 @@ module MT = struct
         sub.attributes sub attrs;
         sub.extension sub x
     | Psig_attribute x -> sub.attribute sub x
+
+  let iter_extension sub : Extensions.Module_type.t -> _ = function
+    | Emty_strengthen { mty; mod_id } ->
+       iter sub mty;
+       iter_loc sub mod_id
+
 end
 
 
@@ -596,6 +602,7 @@ let default_iterator =
     signature = (fun this l -> List.iter (this.signature_item this) l);
     signature_item = MT.iter_signature_item;
     module_type = MT.iter;
+    module_type_extension = MT.iter_extension;
     with_constraint = MT.iter_with_constraint;
     class_declaration =
       (fun this -> CE.class_infos this (this.class_expr this));
@@ -656,9 +663,6 @@ 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_mapper.ml

@@ -284,7 +284,7 @@ module MT = struct
     | Extension emty -> begin
         Extensions_parsing.Module_type.wrap_desc ~loc ~attrs @@
         match sub.module_type_extension sub emty with
-        | _ -> .
+        | Emty_strengthen smty -> Extensions.Strengthen.mty_of ~loc smty
       end
     | Regular desc ->
     match desc with
@@ -342,6 +342,13 @@ module MT = struct
         let attrs = sub.attributes sub attrs in
         extension ~loc ~attrs (sub.extension sub x)
     | Psig_attribute x -> attribute ~loc (sub.attribute sub x)
+
+  let map_extension sub :
+        Extensions.Module_type.t -> Extensions.Module_type.t = function
+    | Emty_strengthen { mty; mod_id } ->
+       let mty = sub.module_type sub mty in
+       let mod_id = map_loc sub mod_id in
+       Emty_strengthen { mty; mod_id }
 end
 
 
@@ -666,6 +673,7 @@ let default_mapper =
     signature = (fun this l -> List.map (this.signature_item this) l);
     signature_item = MT.map_signature_item;
     module_type = MT.map;
+    module_type_extension = MT.map_extension;
     with_constraint = MT.map_with_constraint;
     class_declaration =
       (fun this -> CE.class_infos this (this.class_expr this));
@@ -729,9 +737,6 @@ 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/depend.ml

@@ -341,7 +341,7 @@ and add_binding_op bv bv' pbop =
 
 and add_modtype bv mty =
   match Extensions.Module_type.get_desc mty with
-  | Extension _ -> .
+  | Extension emty -> add_modtype_extension bv emty
   | Regular desc ->
   match desc with
     Pmty_ident l -> add bv l
@@ -373,6 +373,11 @@ and add_modtype bv mty =
   | Pmty_typeof m -> add_module_expr bv m
   | Pmty_extension e -> handle_extension e
 
+and add_modtype_extension bv : Extensions.Module_type.t -> _ = function
+  | Emty_strengthen { mty; mod_id } ->
+     add_modtype bv mty;
+     add_module_path bv mod_id
+
 and add_module_alias bv l =
   (* If we are in delayed dependencies mode, we delay the dependencies
        induced by "Lident s" *)
@@ -386,7 +391,7 @@ and add_module_alias bv l =
 
 and add_modtype_binding bv mty =
   match Extensions.Module_type.get_desc mty with
-  | Extension _ -> .
+  | Extension emty -> add_modtype_extension_binding bv emty
   | Regular desc ->
   match desc with
     Pmty_alias l ->
@@ -398,6 +403,13 @@ and add_modtype_binding bv mty =
   | _ ->
       add_modtype bv mty; bound
 
+and add_modtype_extension_binding bv : Extensions.Module_type.t -> _ = function
+  | Emty_strengthen { mty; mod_id } ->
+     (* treat like a [with] constraint *)
+     add_modtype bv mty;
+     add_module_path bv mod_id;
+     bound
+
 and add_signature bv sg =
   ignore (add_signature_binding bv sg)
 

parsing/extensions.ml

@@ -297,13 +297,13 @@ module Immutable_arrays = struct
 
   let of_pat expr = match expr.ppat_desc with
     | Ppat_array elts -> Iapat_immutable_array elts
-    | _ -> failwith "Malformed immutable array expression"
+    | _ -> failwith "Malformed immutable array pattern"
 end
 
 (** Strengthening *)
 module Strengthen = struct
   type nonrec module_type =
-    | Smt_strengthen of Parsetree.module_type * Longident.t Location.loc
+    { mty : Parsetree.module_type; mod_id : Longident.t Location.loc }
 
   let extension_string = Language_extension.to_string Strengthen
 
@@ -317,15 +317,16 @@ module Strengthen = struct
        ]}
   *)
 
-  let mty_of ~loc = function
-    | Smt_strengthen (mty, lid) ->
-      (* See Note [Wrapping with make_extension] *)
-      Module_type.make_extension ~loc [extension_string] @@
+  let mty_of ~loc { mty; mod_id } =
+    (* See Note [Wrapping with make_extension] *)
+    Module_type.make_extension ~loc [extension_string] @@
       Ast_helper.Mty.functor_ (Named (Location.mknoloc None, mty))
-        (Ast_helper.Mty.ident lid)
+        (Ast_helper.Mty.ident mod_id)
 
   let of_mty mty = match mty.pmty_desc with
-    | Pmty_functor(Named(_, mty), {pmty_desc = Pmty_ident
+    | Pmty_functor(Named(_, mty), {pmty_desc = Pmty_ident mod_id}) ->
+       { mty; mod_id }
+    | _ -> failwith "Malformed strengthen module type"
 end
 
 (******************************************************************************)

parsing/extensions.mli

@@ -82,7 +82,7 @@ end
 (** The ASTs for module type strengthening. *)
 module Strengthen : sig
   type module_type =
-    | Smt_strengthen of Parsetree.module_type * Longident.t Location.loc
+    { mty : Parsetree.module_type; mod_id : Longident.t Location.loc }
 
   val mty_of : loc:Location.t -> module_type -> Parsetree.module_type_desc
 end

parsing/pprintast.ml

@@ -1129,7 +1129,7 @@ and module_type ctxt f x =
       (attributes ctxt) x.pmty_attributes
   end else
     match Extensions.Module_type.get_desc x with
-    | Extension _ -> .
+    | Extension emty -> module_type_extension ctxt f emty
     | Regular desc -> match desc with
     | Pmty_functor (Unit, mt2) ->
         pp f "@[<hov2>functor () ->@ %a@]" (module_type ctxt) mt2
@@ -1149,6 +1149,12 @@ and module_type ctxt f x =
           (list (with_constraint ctxt) ~sep:"@ and@ ") l
     | _ -> module_type1 ctxt f x
 
+and module_type_extension ctxt f : Extensions.Module_type.t -> _ = function
+  | Emty_strengthen { mty; mod_id } ->
+      pp f "@[<hov2>%a@ with@ %a@]"
+        (module_type1 ctxt) mty
+        longident_loc mod_id
+
 and with_constraint ctxt f = function
   | Pwith_type (li, ({ptype_params= ls ;_} as td)) ->
       let ls = List.map fst ls in
@@ -1174,7 +1180,7 @@ and with_constraint ctxt f = function
 and module_type1 ctxt f x =
   if x.pmty_attributes <> [] then module_type ctxt f x
   else match Extensions.Module_type.get_desc x with
-    | Extension _ -> .
+    | Extension emty -> module_type_extension1 ctxt f emty
     | Regular desc -> match desc with
     | Pmty_ident li ->
         pp f "%a" longident_loc li;
@@ -1188,6 +1194,9 @@ and module_type1 ctxt f x =
     | Pmty_extension e -> extension ctxt f e
     | _ -> paren true (module_type ctxt) f x
 
+and module_type_extension1 ctxt f : Extensions.Module_type.t -> _ = function
+  | Emty_strengthen _ as emty -> paren true (module_type_extension ctxt) f emty
+
 and signature ctxt f x =  list ~sep:"@\n" (signature_item ctxt) f x
 
 and signature_item ctxt f x : unit =

typing/typemod.ml

@@ -836,7 +836,7 @@ let map_ext fn exts =
 
 let rec approx_modtype env smty =
   match Extensions.Module_type.get_desc smty with
-  | Extension _ -> .
+  | Extension emty -> approx_modtype_extension env emty
   | Regular desc ->
   match desc with
     Pmty_ident lid ->
@@ -895,6 +895,9 @@ let rec approx_modtype env smty =
   | Pmty_extension ext ->
       raise (Error_forward (Builtin_attributes.error_of_extension ext))
 
+and approx_modtype_extension _env : Extensions.Module_type.t -> _ = function
+  | Emty_strengthen { mty=_; mod_id=_ } -> failwith "strengthen not yet implemented"
+
 and approx_module_declaration env pmd =
   {
     Types.md_type = approx_modtype env pmd.pmd_type;
@@ -1374,7 +1377,7 @@ and transl_modtype_functor_arg env sarg =
 and transl_modtype_aux env smty =
   let loc = smty.pmty_loc in
   match Extensions.Module_type.get_desc smty with
-  | Extension _ -> .
+  | Extension emty -> transl_modtype_extension_aux env emty
   | Regular desc ->
   match desc with
     Pmty_ident lid ->
@@ -1437,6 +1440,9 @@ and transl_modtype_aux env smty =
   | Pmty_extension ext ->
       raise (Error_forward (Builtin_attributes.error_of_extension ext))
 
+and transl_modtype_extension_aux _env : Extensions.Module_type.t -> _ = function
+  | Emty_strengthen { mty=_ ; mod_id=_ } -> failwith "Strengthen not yet implemented"
+
 and transl_with ~loc env remove_aliases (rev_tcstrs,sg) constr =
   let lid, with_info = match constr with
     | Pwith_type (l,decl) ->l , With_type decl