Add a Module_strengthening extension (#142)

* Add Strengthen extension for module types

* Use Pmty_alias, not Pmty_ident.

* Antal's suggestions

* Comment improvement

Co-authored-by: Antal Spector-Zabusky <[email protected]>

---------

Co-authored-by: Antal Spector-Zabusky <[email protected]>

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.of_ast module_type with
-      | Some _ -> .
+      | Some (Emty_strengthen _) -> failwith "strengthen not implemented yet"
       | None ->
       match module_type.Parsetree.pmty_desc with
         Parsetree.Pmty_ident longident ->
@@ -1622,7 +1622,7 @@ module Analyser =
     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 _ -> .
+      | Some (Emty_strengthen _) -> failwith "strengthen not implemented yet"
       | None ->
       match module_type.Parsetree.pmty_desc with
       | Parsetree.Pmty_ident _longident ->

parsing/ast_iterator.ml

@@ -304,6 +304,11 @@ 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
 
 
@@ -597,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));
@@ -657,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

@@ -285,7 +285,7 @@ module MT = struct
     | Some 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
     | None ->
     match desc with
@@ -343,6 +343,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
 
 
@@ -667,6 +674,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));
@@ -730,9 +738,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.of_ast mty with
-  | Some _ -> .
+  | Some emty -> add_modtype_extension bv emty
   | None ->
   match mty.pmty_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.of_ast mty with
-  | Some _ -> .
+  | Some emty -> add_modtype_extension_binding bv emty
   | None ->
   match mty.pmty_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

@@ -296,7 +296,30 @@ 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
+
+(** Module strengthening *)
+module Strengthen = struct
+  type nonrec module_type =
+    { mty : Parsetree.module_type; mod_id : Longident.t Location.loc }
+
+  let extension_string = Language_extension.to_string Module_strengthening
+
+  (* Encoding: [S with M] becomes [functor (_ : S) -> (module M)], where
+     the [(module M)] is a [Pmty_alias].  This isn't syntax we can write, but
+     [(module M)] can be the inferred type for [M], so this should be fine. *)
+
+  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.alias mod_id)
+
+  let of_mty mty = match mty.pmty_desc with
+    | Pmty_functor(Named(_, mty), {pmty_desc = Pmty_alias mod_id}) ->
+       { mty; mod_id }
+    | _ -> failwith "Malformed strengthened module type"
 end
 
 (******************************************************************************)
@@ -350,9 +373,12 @@ module Module_type = struct
   module M = struct
     module AST = Extensions_parsing.Module_type
 
-    type t = |
+    type t =
+      | Emty_strengthen of Strengthen.module_type
 
-    let of_ast_internal (ext : Language_extension.t) _mty = match ext with
+    let of_ast_internal (ext : Language_extension.t) mty = match ext with
+      | Module_strengthening ->
+        Some (Emty_strengthen (Strengthen.of_mty mty))
       | _ -> None
   end
 

parsing/extensions.mli

@@ -79,6 +79,17 @@ module Immutable_arrays : sig
   val pat_of : loc:Location.t -> pattern -> Parsetree.pattern_desc
 end
 
+(** The ASTs for module type strengthening. *)
+module Strengthen : sig
+  type module_type =
+    { mty : Parsetree.module_type; mod_id : Longident.t Location.loc }
+
+  val mty_of : loc:Location.t -> module_type -> Parsetree.module_type_desc
+end
+
+(******************************************)
+(* General facility, which we export *)
+
 (** The module type of language extension ASTs, instantiated once for each
     syntactic category.  We tend to call the pattern-matching functions here
     with unusual indentation, not indenting the [None] branch further so as to
@@ -135,6 +146,9 @@ module type AST = sig
   val of_ast : ast -> t option
 end
 
+(******************************************)
+(* Individual syntactic categories *)
+
 (** Language extensions in expressions *)
 module Expression : sig
   type t =
@@ -154,7 +168,8 @@ end
 
 (** Language extensions in module types *)
 module Module_type : sig
-  type t = |
+  type t =
+    | Emty_strengthen of Strengthen.module_type
 
   include AST with type t := t and type ast := Parsetree.module_type
 end

parsing/pprintast.ml

@@ -1119,7 +1119,7 @@ and module_type ctxt f x =
       (attributes ctxt) x.pmty_attributes
   end else
     match Extensions.Module_type.of_ast x with
-    | Some _ -> .
+    | Some emty -> module_type_extension ctxt f emty
     | None ->
     match x.pmty_desc with
     | Pmty_functor (Unit, mt2) ->
@@ -1140,6 +1140,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
@@ -1165,7 +1171,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.of_ast x with
-    | Some _ -> .
+    | Some emty -> module_type_extension1 ctxt f emty
     | None ->
     match x.pmty_desc with
     | Pmty_ident li ->
@@ -1180,6 +1186,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

@@ -835,7 +835,7 @@ let map_ext fn exts =
 
 let rec approx_modtype env smty =
   match Extensions.Module_type.of_ast smty with
-  | Some _ -> .
+  | Some emty -> approx_modtype_extension env emty
   | None ->
   match smty.pmty_desc with
     Pmty_ident lid ->
@@ -894,6 +894,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;
@@ -1373,7 +1376,7 @@ 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 _ -> .
+  | Some emty -> transl_modtype_extension_aux env emty
   | None ->
   match smty.pmty_desc with
     Pmty_ident lid ->
@@ -1436,6 +1439,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

utils/language_extension.ml

@@ -4,6 +4,7 @@ type t =
   | Include_functor
   | Polymorphic_parameters
   | Immutable_arrays
+  | Module_strengthening
 
 let equal (a : t) (b : t) = (a = b)
 
@@ -13,6 +14,7 @@ let all =
   ; Include_functor
   ; Polymorphic_parameters
   ; Immutable_arrays
+  ; Module_strengthening
   ]
 
 let default_extensions =
@@ -27,13 +29,15 @@ let to_string = function
   | Include_functor -> "include_functor"
   | Polymorphic_parameters -> "polymorphic_parameters"
   | Immutable_arrays -> "immutable_arrays_experimental"
+  | Module_strengthening -> "module_strengthening"
 
 let of_string extn = match String.lowercase_ascii extn with
   | "comprehensions_experimental" -> Some Comprehensions
   | "local" -> Some Local
   | "include_functor" -> Some Include_functor
   | "polymorphic_parameters" -> Some Polymorphic_parameters
   | "immutable_arrays_experimental" -> Some Immutable_arrays
+  | "strengthening" -> Some Module_strengthening
   | _ -> None
 
 let of_string_exn extn =
@@ -48,7 +52,8 @@ let is_erasable = function
   | Comprehensions
   | Include_functor
   | Polymorphic_parameters
-  | Immutable_arrays ->
+  | Immutable_arrays
+  | Module_strengthening ->
       false
 
 module Universe = struct

utils/language_extension.mli

@@ -7,6 +7,7 @@ type t =
   | Include_functor
   | Polymorphic_parameters
   | Immutable_arrays
+  | Module_strengthening
 
 (** Equality on language extensions *)
 val equal : t -> t -> bool