From cc8d97ad3ad5e76815da685743226cd5aab0c2ab Mon Sep 17 00:00:00 2001
From: Roman Leshchinskiy <rleshchinskiy@gmail.com>
Date: Thu, 9 Mar 2023 15:25:30 +0000
Subject: [PATCH] Don't copy when resolving aliases in try_modtypes (#1184)

We were calling `expand_module_alias` to get the unstrengthened,
non-lazy type (potential deep copy), then checking whether it's an ident
and then strengthening it, thus throwing most of it away. This patch
avoids doing that by working directly on the lazy representation. This
substantially reduces the peak heap size (and improves compile times,
albeit to a lesser degree) in some cases.

Note that https://github.com/ocaml-flambda/ocaml-jst/pull/119 changes
this area significantly so I didn't want to refactor too much here.

Co-authored-by: Roman Leshchinskiy <rleshchinskiy@janestreet.com>
---
 ocaml/typing/includemod.ml | 24 ++++++++++++++++--------
 ocaml/typing/mtype.mli     |  1 +
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/ocaml/typing/includemod.ml b/ocaml/typing/includemod.ml
index f85289382cb..eb931eda677 100644
--- a/ocaml/typing/includemod.ml
+++ b/ocaml/typing/includemod.ml
@@ -487,11 +487,12 @@ and try_modtypes ~in_eq ~loc env ~mark subst mty1 mty2 orig_shape =
       | exception Env.Error (Env.Missing_module (_, _, path)) ->
           Error Error.(Mt_core(Unbound_module_path path))
       | p1 ->
-          begin match expand_module_alias ~strengthen:false env p1 with
-          | Error e -> Error (Error.Mt_core e)
-          | Ok mty1 ->
-              match strengthened_modtypes ~in_eq ~loc ~aliasable:true env ~mark
-                      subst mty1 p1 mty2 orig_shape
+          begin match (Env.find_module_lazy p1 env).mdl_type with
+          | exception Not_found ->
+              Error (Error.Mt_core (Error.Unbound_module_path p1))
+          | mty1 ->
+              match strengthened_modtypes_lazy ~in_eq ~loc ~aliasable:true env
+                      ~mark subst mty1 p1 mty2 orig_shape
               with
               | Ok _ as x -> x
               | Error reason -> Error (Error.After_alias_expansion reason)
@@ -627,15 +628,22 @@ and functor_param ~in_eq ~loc env ~mark subst param1 param2 =
   | _, _ ->
       Error (Error.Incompatible_params (param1, param2)), env, subst
 
-and strengthened_modtypes ~in_eq ~loc ~aliasable env ~mark
+and strengthened_modtypes_lazy ~in_eq ~loc ~aliasable env ~mark
     subst mty1 path1 mty2 shape =
+  let open Subst.Lazy in
   match mty1, mty2 with
-  | Mty_ident p1, Mty_ident p2 when equal_modtype_paths env p1 subst p2 ->
+  | MtyL_ident p1, Mty_ident p2 when equal_modtype_paths env p1 subst p2 ->
       Ok (Tcoerce_none, shape)
   | _, _ ->
-      let mty1 = Mtype.strengthen ~aliasable env mty1 path1 in
+      let mty1 = Mtype.strengthen_lazy ~aliasable env mty1 path1 in
+      let mty1 = Subst.Lazy.force_modtype mty1 in
       modtypes ~in_eq ~loc env ~mark subst mty1 mty2 shape
 
+and strengthened_modtypes ~in_eq ~loc ~aliasable env ~mark
+    subst mty1 path1 mty2 shape =
+  strengthened_modtypes_lazy ~in_eq ~loc ~aliasable env ~mark
+    subst (Subst.Lazy.of_modtype mty1) path1 mty2 shape
+
 and strengthened_module_decl ~loc ~aliasable env ~mark
     subst md1 path1 md2 shape =
   match md1.md_type, md2.md_type with
diff --git a/ocaml/typing/mtype.mli b/ocaml/typing/mtype.mli
index 7a53298511e..34a9112730d 100644
--- a/ocaml/typing/mtype.mli
+++ b/ocaml/typing/mtype.mli
@@ -33,6 +33,7 @@ val scrape_for_type_of:
 val freshen: scope:int -> module_type -> module_type
         (* Return an alpha-equivalent copy of the given module type
            where bound identifiers are fresh. *)
+val strengthen_lazy: aliasable:bool -> Env.t -> Subst.Lazy.modtype -> Path.t -> Subst.Lazy.modtype
 val strengthen: aliasable:bool -> Env.t -> module_type -> Path.t -> module_type
         (* Strengthen abstract type components relative to the
            given path. *)