update locations for destructive substitutions

typing/subst.ml

@@ -31,13 +31,15 @@ type t =
     modules: Path.t PathMap.t;
     modtypes: (Ident.t, module_type) Tbl.t;
     for_saving: bool;
+    loc: Location.t option;
   }
 
 let identity =
   { types = PathMap.empty;
     modules = PathMap.empty;
     modtypes = Tbl.empty;
     for_saving = false;
+    loc = None;
   }
 
 let add_type_path id p s = { s with types = PathMap.add id (Path p) s.types }
@@ -53,8 +55,13 @@ let add_modtype id ty s = { s with modtypes = Tbl.add id ty s.modtypes }
 
 let for_saving s = { s with for_saving = true }
 
+let change_locs s loc = { s with loc = Some loc }
+
 let loc s x =
-  if s.for_saving && not !Clflags.keep_locs then Location.none else x
+  match s.loc with
+  | Some l -> l
+  | None ->
+    if s.for_saving && not !Clflags.keep_locs then Location.none else x
 
 let remove_loc =
   let open Ast_mapper in
@@ -474,6 +481,11 @@ let merge_tbls f m1 m2 =
 let merge_path_maps f m1 m2 =
   PathMap.fold (fun k d accu -> PathMap.add k (f d) accu) m1 m2
 
+let keep_latest_loc l1 l2 =
+  match l2 with
+  | None -> l1
+  | Some _ -> l2
+
 let type_replacement s = function
   | Path p -> Path (type_path s p)
   | Type_function { params; body } ->
@@ -489,4 +501,5 @@ let compose s1 s2 =
     modules = merge_path_maps (module_path s2) s1.modules s2.modules;
     modtypes = merge_tbls (modtype s2) s1.modtypes s2.modtypes;
     for_saving = s1.for_saving || s2.for_saving;
+    loc = keep_latest_loc s1.loc s2.loc;
   }

typing/subst.mli

@@ -42,6 +42,7 @@ val add_module_path: Path.t -> Path.t -> t -> t
 val add_modtype: Ident.t -> module_type -> t -> t
 val for_saving: t -> t
 val reset_for_saving: unit -> unit
+val change_locs: t -> Location.t -> t
 
 val module_path: t -> Path.t -> Path.t
 val type_path: t -> Path.t -> Path.t

typing/typemod.ml

@@ -473,13 +473,15 @@ let merge_constraint initial_env remove_aliases loc sg constr =
             then raise(Error(loc, initial_env, With_cannot_remove_constrained_type));
             fun s path -> Subst.add_type_function path ~params ~body s
        in
-       let sub = List.fold_left how_to_extend_subst Subst.identity !real_ids in
+       let sub = Subst.change_locs Subst.identity loc in
+       let sub = List.fold_left how_to_extend_subst sub !real_ids in
        Subst.signature sub sg
     | (_, _, Twith_modsubst (real_path, _)) ->
+       let sub = Subst.change_locs Subst.identity loc in
        let sub =
          List.fold_left
            (fun s path -> Subst.add_module_path path real_path s)
-           Subst.identity
+           sub
            !real_ids
        in
        Subst.signature sub sg