flambda-backend: Remove variables from computed signature for `with m…

…odule` constriants. (#2764)

* Add failing test

* Remove vars from computed signature for [with module]

testsuite/tests/typing-zero-alloc/signatures.ml

@@ -1190,3 +1190,38 @@ Error: Signature mismatch:
        The former provides a weaker "zero_alloc" guarantee than the latter.
        Hint: Add a "zero_alloc" attribute to the implementation.
 |}]
+
+(************************)
+(* Test 14: with module *)
+
+(* [with module] constraints require us to remove variables just like [module
+   type of].  Regression test: the below used to hit an assert, as a result of
+   not removing the vars in the signature of [S_plain] in the implementation of
+   [N]. *)
+module type S = sig
+  module M : sig
+    val f : int -> int
+  end
+end
+
+module N : sig
+  module Plain : sig
+    val f : int -> int
+  end
+
+  module type S_plain = S with module M = Plain
+end = struct
+  module Plain = struct
+    let f x = x+1
+  end
+
+  module type S_plain = S with module M = Plain
+end
+[%%expect{|
+module type S = sig module M : sig val f : int -> int end end
+module N :
+  sig
+    module Plain : sig val f : int -> int end
+    module type S_plain = sig module M : sig val f : int -> int end end
+  end
+|}]

typing/typemod.ml

@@ -552,6 +552,59 @@ let params_are_constrained =
   in
   loop
 
+let rec remove_modality_and_zero_alloc_variables_sg env ~zap_modality sg =
+  let sg_item = function
+    | Sig_value (id, desc, vis) ->
+        let val_modalities =
+          desc.val_modalities
+          |> zap_modality |> Mode.Modality.Value.of_const
+        in
+        let val_zero_alloc =
+          Zero_alloc.create_const (Zero_alloc.get desc.val_zero_alloc)
+        in
+        let desc = {desc with val_modalities; val_zero_alloc} in
+        Sig_value (id, desc, vis)
+    | Sig_module (id, pres, md, re, vis) ->
+        let md_type =
+          remove_modality_and_zero_alloc_variables_mty env ~zap_modality
+            md.md_type
+        in
+        let md = {md with md_type} in
+        Sig_module (id, pres, md, re, vis)
+    | item -> item
+  in
+  List.map sg_item sg
+
+and remove_modality_and_zero_alloc_variables_mty env ~zap_modality mty =
+  match mty with
+  | Mty_ident _ | Mty_alias _ ->
+    (* module types with names can't have inferred modalities. *)
+    mty
+  | Mty_signature sg ->
+    Mty_signature
+      (remove_modality_and_zero_alloc_variables_sg env ~zap_modality sg)
+  | Mty_functor (param, mty) ->
+    let param : Types.functor_parameter =
+      match param with
+      | Named (id, mty) ->
+          let mty =
+            remove_modality_and_zero_alloc_variables_mty env
+              ~zap_modality:Mode.Modality.Value.to_const_exn mty
+          in
+          Named (id, mty)
+      | Unit -> Unit
+    in
+    let mty =
+      remove_modality_and_zero_alloc_variables_mty env ~zap_modality mty
+    in
+    Mty_functor (param, mty)
+  | Mty_strengthen (mty, path, alias) ->
+      let mty =
+        remove_modality_and_zero_alloc_variables_mty env
+        ~zap_modality:Mode.Modality.Value.to_const_exn mty
+      in
+      Mty_strengthen (mty, path, alias)
+
 type with_info =
   | With_type of Parsetree.type_declaration
   | With_typesubst of Parsetree.type_declaration
@@ -746,6 +799,10 @@ let merge_constraint initial_env loc sg lid constr =
         let sig_env = Env.add_signature sg_for_env outer_sig_env in
         let mty = md'.md_type in
         let mty = Mtype.scrape_for_type_of ~remove_aliases sig_env mty in
+        let mty =
+          remove_modality_and_zero_alloc_variables_mty sig_env
+            ~zap_modality:Mode.Modality.Value.zap_to_floor mty
+        in
         let md'' = { md' with md_type = mty } in
         let newmd = Mtype.strengthen_decl ~aliasable:false md'' path in
         ignore(Includemod.modtypes  ~mark:Mark_both ~loc sig_env
@@ -2120,59 +2177,6 @@ let remove_mode_and_jkind_variables env sg =
   let rm _env ty = Ctype.remove_mode_and_jkind_variables ty; None in
   List.find_map (nongen_signature_item env rm) sg |> ignore
 
-let rec remove_modality_and_zero_alloc_variables_sg env ~zap_modality sg =
-  let sg_item = function
-    | Sig_value (id, desc, vis) ->
-        let val_modalities =
-          desc.val_modalities
-          |> zap_modality |> Mode.Modality.Value.of_const
-        in
-        let val_zero_alloc =
-          Zero_alloc.create_const (Zero_alloc.get desc.val_zero_alloc)
-        in
-        let desc = {desc with val_modalities; val_zero_alloc} in
-        Sig_value (id, desc, vis)
-    | Sig_module (id, pres, md, re, vis) ->
-        let md_type =
-          remove_modality_and_zero_alloc_variables_mty env ~zap_modality
-            md.md_type
-        in
-        let md = {md with md_type} in
-        Sig_module (id, pres, md, re, vis)
-    | item -> item
-  in
-  List.map sg_item sg
-
-and remove_modality_and_zero_alloc_variables_mty env ~zap_modality mty =
-  match mty with
-  | Mty_ident _ | Mty_alias _ ->
-    (* module types with names can't have inferred modalities. *)
-    mty
-  | Mty_signature sg ->
-    Mty_signature
-      (remove_modality_and_zero_alloc_variables_sg env ~zap_modality sg)
-  | Mty_functor (param, mty) ->
-    let param : Types.functor_parameter =
-      match param with
-      | Named (id, mty) ->
-          let mty =
-            remove_modality_and_zero_alloc_variables_mty env
-              ~zap_modality:Mode.Modality.Value.to_const_exn mty
-          in
-          Named (id, mty)
-      | Unit -> Unit
-    in
-    let mty =
-      remove_modality_and_zero_alloc_variables_mty env ~zap_modality mty
-    in
-    Mty_functor (param, mty)
-  | Mty_strengthen (mty, path, alias) ->
-      let mty =
-        remove_modality_and_zero_alloc_variables_mty env
-        ~zap_modality:Mode.Modality.Value.to_const_exn mty
-      in
-      Mty_strengthen (mty, path, alias)
-
 (* Helpers for typing recursive modules *)
 
 let anchor_submodule name anchor =