Add support for bivariant params in module types

Signed-off-by: Nathan Rebours <nathan.rebours@ocamlpro.com>

Author: Nathan Rebours

astlib/encoding_504.ml

@@ -7,6 +7,7 @@ module Ext_name = struct
   let bivariant_param = "ppxlib.migration.bivariant_param_5_4"
   let bivariant_pstr = "ppxlib.migration.bivariant_str_item_5_4"
   let bivariant_psig = "ppxlib.migration.bivariant_sig_item_5_4"
+  let bivariant_pmty_with = "ppxlib.migration.bivariant_pmty_with_5_4"
 end
 
 let invalid_encoding ~loc name =
@@ -460,6 +461,50 @@ module To_503 = struct
         ] ->
         x
     | _ -> invalid_encoding ~loc "bivariant signature_item"
+
+  let encode_bivariant_pmty_with ~loc mty constraints =
+    let loc = { loc with Location.loc_ghost = true } in
+    let pmd_type =
+      {
+        pmty_loc = loc;
+        pmty_attributes = [];
+        pmty_desc = Pmty_with (mty, constraints);
+      }
+    in
+    let psig_desc =
+      Psig_module
+        {
+          pmd_name = { txt = None; loc };
+          pmd_type;
+          pmd_attributes = [];
+          pmd_loc = loc;
+        }
+    in
+    let ext =
+      ( { txt = Ext_name.bivariant_pmty_with; loc },
+        PSig [ { psig_loc = loc; psig_desc } ] )
+    in
+    Pmty_extension ext
+
+  let decode_bivariant_pmty_with ~loc payload =
+    match payload with
+    | PSig
+        [
+          {
+            psig_desc =
+              Psig_module
+                {
+                  pmd_name = { txt = None; _ };
+                  pmd_attributes = [];
+                  pmd_type =
+                    { pmty_attributes = []; pmty_desc = Pmty_with _ as x; _ };
+                  _;
+                };
+            _;
+          };
+        ] ->
+        x
+    | _ -> invalid_encoding ~loc "bivariant pmty_with"
 end
 
 module To_502 = struct

astlib/encoding_504.mli

@@ -4,6 +4,7 @@ module Ext_name : sig
   val ppat_labeled_tuple : string
   val bivariant_pstr : string
   val bivariant_psig : string
+  val bivariant_pmty_with : string
 end
 
 module To_503 : sig
@@ -64,8 +65,12 @@ module To_503 : sig
   val encode_bivariant_psig_class_type :
     loc:Location.t -> class_type_declaration list -> signature_item_desc
 
+  val encode_bivariant_pmty_with :
+    loc:Location.t -> module_type -> with_constraint list -> module_type_desc
+
   val decode_bivariant_pstr : loc:Location.t -> payload -> structure_item_desc
   val decode_bivariant_psig : loc:Location.t -> payload -> signature_item_desc
+  val decode_bivariant_pmty_with : loc:Location.t -> payload -> module_type_desc
 end
 
 module To_502 : sig

astlib/migrate_503_504.ml

@@ -775,13 +775,14 @@ and copy_module_type :
        Ast_503.Parsetree.pmty_loc;
        Ast_503.Parsetree.pmty_attributes;
      } ->
+  let loc = copy_location pmty_loc in
   {
-    Ast_504.Parsetree.pmty_desc = copy_module_type_desc pmty_desc;
-    Ast_504.Parsetree.pmty_loc = copy_location pmty_loc;
+    Ast_504.Parsetree.pmty_desc = copy_module_type_desc_with_loc ~loc pmty_desc;
+    Ast_504.Parsetree.pmty_loc = loc;
     Ast_504.Parsetree.pmty_attributes = copy_attributes pmty_attributes;
   }
 
-and copy_module_type_desc :
+and copy_module_type_desc_with_loc ~loc :
     Ast_503.Parsetree.module_type_desc -> Ast_504.Parsetree.module_type_desc =
   function
   | Ast_503.Parsetree.Pmty_ident x0 ->
@@ -796,11 +797,20 @@ and copy_module_type_desc :
         (copy_module_type x0, List.map copy_with_constraint x1)
   | Ast_503.Parsetree.Pmty_typeof x0 ->
       Ast_504.Parsetree.Pmty_typeof (copy_module_expr x0)
+  | Ast_503.Parsetree.Pmty_extension ({ txt; _ }, payload)
+    when String.equal txt Encoding_504.Ext_name.bivariant_pmty_with ->
+      let pmty_desc =
+        Encoding_504.To_503.decode_bivariant_pmty_with ~loc payload
+      in
+      copy_module_type_desc_with_loc ~loc pmty_desc
   | Ast_503.Parsetree.Pmty_extension x0 ->
       Ast_504.Parsetree.Pmty_extension (copy_extension x0)
   | Ast_503.Parsetree.Pmty_alias x0 ->
       Ast_504.Parsetree.Pmty_alias (copy_loc (copy_Longident_t ~loc:x0.loc) x0)
 
+and copy_module_type_desc pmty =
+  copy_module_type_desc_with_loc ~loc:Location.none pmty
+
 and copy_with_constraint :
     Ast_503.Parsetree.with_constraint -> Ast_504.Parsetree.with_constraint =
   function

astlib/migrate_504_503.ml

@@ -11,6 +11,7 @@ module Bivariant_param = struct
     | Class_decl of Ast_503.Parsetree.class_declaration
     | Class_desc of Ast_503.Parsetree.class_description
     | Class_type_decl of Ast_503.Parsetree.class_type_declaration
+    | With_constraint of Ast_503.Parsetree.with_constraint
 
   (* TODO: register exception printers to display those as location errors
    pointing to the right AST element and displaying a clear migration error
@@ -790,13 +791,14 @@ and copy_module_type :
        Ast_504.Parsetree.pmty_loc;
        Ast_504.Parsetree.pmty_attributes;
      } ->
+  let loc = copy_location pmty_loc in
   {
-    Ast_503.Parsetree.pmty_desc = copy_module_type_desc pmty_desc;
-    Ast_503.Parsetree.pmty_loc = copy_location pmty_loc;
+    Ast_503.Parsetree.pmty_desc = copy_module_type_desc_with_loc ~loc pmty_desc;
+    Ast_503.Parsetree.pmty_loc = loc;
     Ast_503.Parsetree.pmty_attributes = copy_attributes pmty_attributes;
   }
 
-and copy_module_type_desc :
+and copy_module_type_desc_with_loc ~loc :
     Ast_504.Parsetree.module_type_desc -> Ast_503.Parsetree.module_type_desc =
   function
   | Ast_504.Parsetree.Pmty_ident x0 ->
@@ -807,21 +809,46 @@ and copy_module_type_desc :
       Ast_503.Parsetree.Pmty_functor
         (copy_functor_parameter x0, copy_module_type x1)
   | Ast_504.Parsetree.Pmty_with (x0, x1) ->
-      Ast_503.Parsetree.Pmty_with
-        (copy_module_type x0, List.map copy_with_constraint x1)
+      let mty = copy_module_type x0 in
+      let contains_bivariant = ref false in
+      let constraints =
+        List.map
+          (fun c ->
+            match copy_with_constraint c with
+            | c' -> c'
+            | exception Bivariant_param.With_constraint c' ->
+                contains_bivariant := true;
+                c')
+          x1
+      in
+      if !contains_bivariant then
+        Encoding_504.To_503.encode_bivariant_pmty_with ~loc mty constraints
+      else
+        Ast_503.Parsetree.Pmty_with
+          (copy_module_type x0, List.map copy_with_constraint x1)
   | Ast_504.Parsetree.Pmty_typeof x0 ->
       Ast_503.Parsetree.Pmty_typeof (copy_module_expr x0)
   | Ast_504.Parsetree.Pmty_extension x0 ->
       Ast_503.Parsetree.Pmty_extension (copy_extension x0)
   | Ast_504.Parsetree.Pmty_alias x0 ->
       Ast_503.Parsetree.Pmty_alias (copy_loc copy_Longident_t x0)
 
+and copy_module_type_desc pmty =
+  copy_module_type_desc_with_loc ~loc:Location.none pmty
+
 and copy_with_constraint :
     Ast_504.Parsetree.with_constraint -> Ast_503.Parsetree.with_constraint =
   function
   | Ast_504.Parsetree.Pwith_type (x0, x1) ->
-      Ast_503.Parsetree.Pwith_type
-        (copy_loc copy_Longident_t x0, copy_type_declaration x1)
+      let lident_loc = copy_loc copy_Longident_t x0 in
+      let td, contains_bivariant =
+        match copy_type_declaration x1 with
+        | td -> (td, false)
+        | exception Bivariant_param.Type_decl td -> (td, true)
+      in
+      let res = Ast_503.Parsetree.Pwith_type (lident_loc, td) in
+      if contains_bivariant then raise (Bivariant_param.With_constraint res)
+      else res
   | Ast_504.Parsetree.Pwith_module (x0, x1) ->
       Ast_503.Parsetree.Pwith_module
         (copy_loc copy_Longident_t x0, copy_loc copy_Longident_t x1)
@@ -832,8 +859,15 @@ and copy_with_constraint :
       Ast_503.Parsetree.Pwith_modtypesubst
         (copy_loc copy_Longident_t x0, copy_module_type x1)
   | Ast_504.Parsetree.Pwith_typesubst (x0, x1) ->
-      Ast_503.Parsetree.Pwith_typesubst
-        (copy_loc copy_Longident_t x0, copy_type_declaration x1)
+      let lident_loc = copy_loc copy_Longident_t x0 in
+      let td, contains_bivariant =
+        match copy_type_declaration x1 with
+        | td -> (td, false)
+        | exception Bivariant_param.Type_decl td -> (td, true)
+      in
+      let res = Ast_503.Parsetree.Pwith_typesubst (lident_loc, td) in
+      if contains_bivariant then raise (Bivariant_param.With_constraint res)
+      else res
   | Ast_504.Parsetree.Pwith_modsubst (x0, x1) ->
       Ast_503.Parsetree.Pwith_modsubst
         (copy_loc copy_Longident_t x0, copy_loc copy_Longident_t x1)

test/encoding/504/migrations/run.t

@@ -75,6 +75,8 @@ Bivariant are also correctly handled in a signature context:
   > type +-'a extensible += B
   > class [+-'a] c : object end
   > class type [+-'a] d = object end
+  > module M : S with type +-'a t = 'a u
+  > module M2 : S with type +-'a t := 'a u
   > EOF
 
   $ ./id_driver.exe bivariant.mli
@@ -84,6 +86,12 @@ Bivariant are also correctly handled in a signature context:
                                                  | B ]
   [%%ppxlib.migration.bivariant_sig_item_5_4 : class ['a] c : object  end]
   [%%ppxlib.migration.bivariant_sig_item_5_4 : class type ['a] d = object  end]
+  module M :
+  [%ppxlib.migration.bivariant_pmty_with_5_4 :
+    module _ : S with type 'a  t =  'a u]
+  module M2 :
+  [%ppxlib.migration.bivariant_pmty_with_5_4 :
+    module _ : S with type 'a  t :=  'a u]
 
   $ ./id_driver.exe bivariant.mli --use-compiler-pp
   type +-'a t
@@ -92,3 +100,5 @@ Bivariant are also correctly handled in a signature context:
     | B 
   class [+-'a] c : object  end
   class type [+-'a] d = object  end
+  module M : S with type +-'a  t =  'a u
+  module M2 : S with type +-'a  t :=  'a u