Remove expansion from "unexpanded" module type of

We've been keeping each `module type of` expression's expansion around, even in
the nominally unexpanded version of the type. This was needed in case the
expression was invalidated so that we had the expansion on hand as a fallback.
However, it led to exponential blowup in the case of deeply nested `module type
of`s.

Happily, we can do without the expansion if we do a little more work when
substituting: We already track in the `subst` which module paths have been
invalidated. When we need to invalidate a module path, we can stash the
expression for the module's type in that `subst`. This is effectively the same
data that we were keeping in the "unexpanded" expansion. And then when we need
to expand an invalidated `module type of`, we can use the stashed expression.

Author: lukemaurer

src/document/generator.ml

@@ -1510,8 +1510,7 @@ module Make (Syntax : SYNTAX) = struct
     and umty_hidden : Odoc_model.Lang.ModuleType.U.expr -> bool = function
       | Path p -> Paths.Path.(is_hidden (p :> t))
       | With (_, expr) -> umty_hidden expr
-      | TypeOf { t_desc = ModPath m; _ }
-      | TypeOf { t_desc = StructInclude m; _ } ->
+      | TypeOf (ModPath m) | TypeOf (StructInclude m) ->
           Paths.Path.(is_hidden (m :> t))
       | Signature _ -> false
 
@@ -1559,7 +1558,7 @@ module Make (Syntax : SYNTAX) = struct
       | With (_, expr) when is_elidable_with_u expr ->
           Syntax.Mod.open_tag ++ O.txt " ... " ++ Syntax.Mod.close_tag
       | With (subs, expr) -> mty_with subs expr
-      | TypeOf { t_desc; _ } -> mty_typeof t_desc
+      | TypeOf t -> mty_typeof t
 
     and mty : Odoc_model.Lang.ModuleType.expr -> text =
      fun m ->

src/loader/cmt.ml

@@ -544,7 +544,7 @@ and read_include env parent incl =
   let decl_modty =
     match unwrap_module_expr_desc incl.incl_mod.mod_desc with
     | Tmod_ident(p, _) ->
-      Some (ModuleType.U.TypeOf {t_desc = ModuleType.StructInclude (Env.Path.read_module env p); t_expansion=None })
+      Some (ModuleType.U.TypeOf (ModuleType.StructInclude (Env.Path.read_module env p)))
     | _ ->
       let mty = read_module_expr env parent container incl.incl_mod in
       umty_of_mty mty

src/model/lang.ml

@@ -92,19 +92,12 @@ and ModuleType : sig
     | Signature of Signature.t
     | Functor of FunctorParameter.t * simple_expansion
 
-  type typeof_t = {
-    t_desc : type_of_desc;
-    t_expansion : simple_expansion option;
-  }
-
   module U : sig
     type expr =
       | Path of Path.ModuleType.t
       | Signature of Signature.t
       | With of substitution list * expr
-      | TypeOf of typeof_t
-
-    (* Nb. this may have an expansion! *)
+      | TypeOf of type_of_desc
   end
 
   type path_t = {
@@ -118,6 +111,11 @@ and ModuleType : sig
     w_expr : U.expr;
   }
 
+  type typeof_t = {
+    t_desc : type_of_desc;
+    t_expansion : simple_expansion option;
+  }
+
   type expr =
     | Path of path_t
     | Signature of Signature.t
@@ -534,7 +532,7 @@ let umty_of_mty : ModuleType.expr -> ModuleType.U.expr option = function
   | Signature sg -> Some (Signature sg)
   | Path { p_path; _ } -> Some (Path p_path)
   | Functor _ -> None
-  | TypeOf t -> Some (TypeOf t)
+  | TypeOf t -> Some (TypeOf t.t_desc)
   | With { w_substitutions; w_expr; _ } -> Some (With (w_substitutions, w_expr))
 
 (** Query the top-comment of a signature. This is [s.doc] most of the time with

src/model_desc/lang_desc.ml

@@ -173,7 +173,7 @@ and moduletype_u_expr =
           ( "With",
             (t, e),
             Pair (List moduletype_substitution, moduletype_u_expr) )
-    | TypeOf x -> C ("TypeOf", x, moduletype_typeof_t))
+    | TypeOf x -> C ("TypeOf", x, moduletype_type_of_desc))
 
 and moduletype_t =
   let open Lang.ModuleType in

src/xref2/compile.ml

@@ -592,8 +592,7 @@ and module_type_map_subs env id cexpr subs =
     | Path (`Resolved p) -> Some (`ModuleType p)
     | Path _ -> None
     | With (_, e) -> find_parent e
-    | TypeOf { t_desc = ModPath (`Resolved p); _ }
-    | TypeOf { t_desc = StructInclude (`Resolved p); _ } ->
+    | TypeOf (ModPath (`Resolved p)) | TypeOf (StructInclude (`Resolved p)) ->
         Some (`Module p)
     | TypeOf _ -> None
   in
@@ -635,13 +634,13 @@ and u_module_type_expr :
         in
         let result : ModuleType.U.expr = With (subs', expr') in
         result
-    | TypeOf { t_desc; t_expansion } ->
-        let t_desc =
-          match t_desc with
+    | TypeOf t ->
+        let t =
+          match t with
           | ModPath p -> ModPath (module_path env p)
           | StructInclude p -> StructInclude (module_path env p)
         in
-        TypeOf { t_desc; t_expansion }
+        TypeOf t
   in
   inner expr
 

src/xref2/component.ml

@@ -199,17 +199,12 @@ and ModuleType : sig
     | Signature of Signature.t
     | Functor of FunctorParameter.t * simple_expansion
 
-  type typeof_t = {
-    t_desc : type_of_desc;
-    t_expansion : simple_expansion option;
-  }
-
   module U : sig
     type expr =
       | Path of Cpath.module_type
       | Signature of Signature.t
       | With of substitution list * expr
-      | TypeOf of typeof_t
+      | TypeOf of type_of_desc
   end
 
   type path_t = {
@@ -223,6 +218,11 @@ and ModuleType : sig
     w_expr : U.expr;
   }
 
+  type typeof_t = {
+    t_desc : type_of_desc;
+    t_expansion : simple_expansion option;
+  }
+
   type expr =
     | Path of path_t
     | Signature of Signature.t
@@ -455,7 +455,7 @@ and Substitution : sig
     type_replacement : (TypeExpr.t * TypeDecl.Equation.t) PathTypeMap.t;
     module_type_replacement : ModuleType.expr ModuleTypeMap.t;
     path_invalidating_modules : Ident.path_module list;
-    module_type_of_invalidating_modules : Ident.path_module list;
+    module_type_of_invalidating_modules : ModuleType.expr PathModuleMap.t;
     unresolve_opaque_paths : bool;
   }
 end =
@@ -744,7 +744,7 @@ module Fmt = struct
     | With (subs, e) ->
         Format.fprintf ppf "%a with [%a]" u_module_type_expr e substitution_list
           subs
-    | TypeOf { t_desc; _ } -> module_type_type_of_desc ppf t_desc
+    | TypeOf t -> module_type_type_of_desc ppf t
 
   and module_type_expr ppf mt =
     let open ModuleType in
@@ -2121,14 +2121,13 @@ module Of_Lang = struct
     | With (w, e) ->
         let w' = List.map (with_module_type_substitution ident_map) w in
         With (w', u_module_type_expr ident_map e)
-    | TypeOf { t_desc; t_expansion } ->
-        let t_desc =
-          match t_desc with
+    | TypeOf t ->
+        let t =
+          match t with
           | ModPath p -> ModuleType.ModPath (module_path ident_map p)
           | StructInclude p -> StructInclude (module_path ident_map p)
         in
-        let t_expansion = Opt.map (simple_expansion ident_map) t_expansion in
-        TypeOf { t_desc; t_expansion }
+        TypeOf t
 
   and module_type_expr ident_map m =
     let open Odoc_model in

src/xref2/component.mli

@@ -182,17 +182,12 @@ and ModuleType : sig
     | Signature of Signature.t
     | Functor of FunctorParameter.t * simple_expansion
 
-  type typeof_t = {
-    t_desc : type_of_desc;
-    t_expansion : simple_expansion option;
-  }
-
   module U : sig
     type expr =
       | Path of Cpath.module_type
       | Signature of Signature.t
       | With of substitution list * expr
-      | TypeOf of typeof_t
+      | TypeOf of type_of_desc
   end
 
   type path_t = {
@@ -206,6 +201,11 @@ and ModuleType : sig
     w_expr : U.expr;
   }
 
+  type typeof_t = {
+    t_desc : type_of_desc;
+    t_expansion : simple_expansion option;
+  }
+
   type expr =
     | Path of path_t
     | Signature of Signature.t
@@ -426,7 +426,7 @@ and Substitution : sig
     type_replacement : (TypeExpr.t * TypeDecl.Equation.t) PathTypeMap.t;
     module_type_replacement : ModuleType.expr ModuleTypeMap.t;
     path_invalidating_modules : Ident.path_module list;
-    module_type_of_invalidating_modules : Ident.path_module list;
+    module_type_of_invalidating_modules : ModuleType.expr PathModuleMap.t;
     unresolve_opaque_paths : bool;
   }
 end

src/xref2/expand_tools.ml

@@ -25,7 +25,9 @@ let handle_expansion env id expansion =
           Subst.add_module (arg.id :> Ident.path_module) p rp Subst.identity
         in
         let subst =
-          Subst.mto_invalidate_module (arg.id :> Ident.path_module) subst
+          Subst.mto_invalidate_module
+            (arg.id :> Ident.path_module)
+            arg.expr subst
         in
         (env', Subst.module_type_expr subst expr)
   in

src/xref2/lang_of.ml

@@ -772,18 +772,8 @@ and u_module_type_expr map identifier = function
       With
         ( List.map (mty_substitution map identifier) subs,
           u_module_type_expr map identifier expr )
-  | TypeOf { t_desc = ModPath p; t_expansion } ->
-      TypeOf
-        {
-          t_desc = ModPath (Path.module_ map p);
-          t_expansion = Opt.map (simple_expansion map identifier) t_expansion;
-        }
-  | TypeOf { t_desc = StructInclude p; t_expansion } ->
-      TypeOf
-        {
-          t_desc = StructInclude (Path.module_ map p);
-          t_expansion = Opt.map (simple_expansion map identifier) t_expansion;
-        }
+  | TypeOf (ModPath p) -> TypeOf (ModPath (Path.module_ map p))
+  | TypeOf (StructInclude p) -> TypeOf (StructInclude (Path.module_ map p))
 
 and module_type_expr map identifier = function
   | Component.ModuleType.Path { p_path; p_expansion } ->

src/xref2/link.ml

@@ -748,10 +748,8 @@ and u_module_type_expr :
       | Error e ->
           Errors.report ~what:(`Module_type_U cexpr) ~tools_error:e `Resolve;
           unresolved)
-  | TypeOf { t_desc = StructInclude p; t_expansion } ->
-      TypeOf { t_desc = StructInclude (module_path env p); t_expansion }
-  | TypeOf { t_desc = ModPath p; t_expansion } ->
-      TypeOf { t_desc = ModPath (module_path env p); t_expansion }
+  | TypeOf (StructInclude p) -> TypeOf (StructInclude (module_path env p))
+  | TypeOf (ModPath p) -> TypeOf (ModPath (module_path env p))
 
 and module_type_expr :
     Env.t -> Id.Signature.t -> ModuleType.expr -> ModuleType.expr =