Remember type path lookups within a compilation

changes/type-path-memo.polish.md

@@ -0,0 +1 @@
+- Build: remember type path lookups within a compilation, which avoids looking up the same types such as `option` and `string` hundreds of thousands of times per build.

compiler/ml/ctype.ml

@@ -1206,11 +1206,11 @@ let expand_abbrev_gen kind find_type_expansion env ty =
       ty'
     | None -> (
       match find_type_expansion path env with
-      | exception Not_found ->
+      | exception Not_found -> (
         (* another way to expand is to normalize the path itself *)
-        let path' = Env.normalize_path None env path in
-        if Path.same path path' then raise Cannot_expand
-        else newty2 level (Tconstr (path', args, abbrev))
+        match Env.normalize_type_path env path with
+        | None -> raise Cannot_expand
+        | Some path' -> newty2 level (Tconstr (path', args, abbrev)))
       | params, body, lv ->
         (* prerr_endline
            ("add a "^string_of_kind kind^" expansion for "^Path.name path);*)

compiler/ml/env.ml

@@ -549,11 +549,156 @@ type pers_struct = {
   ps_filename: string;
 }
 
+(* Request-local memo of type path lookups ([find_type_full] on regular
+   paths) and of [normalize_path] calls that leave a type path unchanged.
+
+   The typechecker resolves the same type paths over and over (expanding
+   [option], [Js.Null.t], ... in [expand_head], [update_level], [occur_univar],
+   ...), each time walking the [open] layers of [env.types] or the components
+   of a persistent module. Both lookups only read immutable environment tables
+   ([env.types] for [Pident] type paths, [env.components] for [Pdot] ones,
+   [env.modules] and [env.components] for normalization) plus request-wide
+   state that only grows: persistent structures (replacing one starts a new
+   memo epoch, see [add_persistent_structure]) and module components, which
+   are forced at most once. An entry is therefore
+   reused only when the tables it read are physically the same, and only
+   successful results are stored. A repeated call has no side effect the memo
+   would skip: the first, uncached call already loaded the [.cmi]s and forced
+   the lazy components it needed. Paths that contain a functor application are
+   never cached ([find_module_descr] applies the functor against [env]), and
+   the memo is bypassed while [.cmi] loading is disabled ([without_cmis] may
+   backtrack components). *)
+module Type_path_memo = struct
+  (* Powers of two. A table starts empty, and only allocates its first array
+     when it stores an entry, so that a process that doesn't typecheck (such as
+     a parse with -bs-ast) allocates none. It grows when it has stored as many
+     entries as it has slots since it last grew, which keeps a small
+     compilation from allocating the full size. *)
+  let initial_size = 64
+  let max_size = 4096
+
+  let string_hash s =
+    let n = String.length s in
+    if n = 0 then 0
+    else
+      (n * 31)
+      + (Char.code (String.unsafe_get s 0) * 7)
+      + Char.code (String.unsafe_get s (n - 1))
+
+  exception Not_cacheable
+
+  (* Consistent with [equal]: it ignores [pos] and ident flags, which no lookup
+     reads. *)
+  let rec hash (path : Path.t) =
+    match path with
+    | Pident id -> (id.Ident.stamp * 65599) + string_hash id.name
+    | Pdot (p, s, _pos) -> (hash p * 31) + string_hash s
+    | Papply _ -> raise_notrace Not_cacheable
+
+  (* Paths that every lookup in this file resolves identically: identifiers are
+     looked up by name and stamp, and components by name. *)
+  let rec equal (p1 : Path.t) (p2 : Path.t) =
+    p1 == p2
+    ||
+    match (p1, p2) with
+    | Pident id1, Pident id2 ->
+      id1.Ident.stamp = id2.Ident.stamp && String.equal id1.name id2.name
+    | Pdot (p1, s1, _), Pdot (p2, s2, _) -> String.equal s1 s2 && equal p1 p2
+    | _ -> false
+
+  type ('types, 'modules, 'components, 'result) entry = {
+    mutable epoch: unit ref;
+    mutable path: Path.t;
+    mutable types: 'types;
+    mutable modules: 'modules;
+    mutable components: 'components;
+    mutable result: 'result;
+  }
+
+  type ('types, 'modules, 'components, 'result) t = {
+    mutable entries:
+      ('types, 'modules, 'components, 'result) entry option array;
+    mutable stores: int;  (** Since [entries] was allocated. *)
+  }
+
+  let create () = {entries = [||]; stores = 0}
+
+  let clear t = Array.fill t.entries 0 (Array.length t.entries) None
+
+  (* The entry in the slot of a path whose [hash] is [h]. *)
+  let find t h =
+    let entries = t.entries in
+    let n = Array.length entries in
+    if n = 0 then None else Array.unsafe_get entries (h land (n - 1))
+
+  (* Stores an entry in the slot of a path whose [hash] is [h]. An entry is
+     only referenced from its slot, so the one it replaces is updated in place
+     instead of allocating another: a file that declares many types changes
+     [env.types] often, and stores the paths it looks up again each time. *)
+  let store t h ~epoch ~path ~types ~modules ~components ~result =
+    let n = Array.length t.entries in
+    if n = 0 then t.entries <- Array.make initial_size None
+    else if t.stores >= n && n < max_size then (
+      let old = t.entries in
+      let entries = Array.make (n * 2) None in
+      Array.iter
+        (function
+          | None -> ()
+          | Some e as slot ->
+            Array.unsafe_set entries (hash e.path land ((n * 2) - 1)) slot)
+        old;
+      t.entries <- entries;
+      t.stores <- 0);
+    t.stores <- t.stores + 1;
+    let entries = t.entries in
+    let i = h land (Array.length entries - 1) in
+    match Array.unsafe_get entries i with
+    | Some e ->
+      e.epoch <- epoch;
+      e.path <- path;
+      e.types <- types;
+      e.modules <- modules;
+      e.components <- components;
+      e.result <- result
+    | None ->
+      Array.unsafe_set entries i
+        (Some {epoch; path; types; modules; components; result})
+end
+
+(* The tables are allocated once per domain rather than per request: they can
+   grow large, and requests are small. An entry is valid only in the epoch that
+   stored it. Each request starts a fresh epoch (and a nested request restores
+   the outer one when it returns), so entries never cross requests; clearing
+   starts a new epoch too. *)
+let type_path_memo_key =
+  Domain.DLS.new_key (fun () ->
+      (Type_path_memo.create (), Type_path_memo.create ()))
+
+let type_path_memo_epoch_key = Request_local.new_key (fun () -> ref (ref ()))
+let type_path_memo_epoch () = !(Domain.DLS.get type_path_memo_epoch_key)
+
+let clear_type_path_memo_tables () =
+  let found, unchanged = Domain.DLS.get type_path_memo_key in
+  Type_path_memo.clear found;
+  Type_path_memo.clear unchanged
+
+let clear_type_path_memo () =
+  Domain.DLS.get type_path_memo_epoch_key := ref ();
+  clear_type_path_memo_tables ()
+
 let persistent_structures_key =
   Request_local.new_key (fun () ->
       (Hashtbl.create 17 : (string, pers_struct option) Hashtbl.t))
 let persistent_structures () = Domain.DLS.get persistent_structures_key
 
+(* [Hashtbl.add] shadows an earlier binding for [name], which the type path
+   memo may have read (directly, or as a missing module), so replacing one
+   starts a new memo epoch. *)
+let add_persistent_structure name ps =
+  let tbl = persistent_structures () in
+  if Hashtbl.mem tbl name then clear_type_path_memo ();
+  Hashtbl.add tbl name ps
+
 (* Consistency between persistent structures *)
 
 let crc_units_key = Request_local.new_key Consistbl.create
@@ -590,7 +735,7 @@ let check_consistency ps =
 
 let save_pers_struct crc ps =
   let modname = ps.ps_name in
-  Hashtbl.add (persistent_structures ()) modname (Some ps);
+  add_persistent_structure modname (Some ps);
   Consistbl.set (crc_units ()) modname crc ps.ps_filename;
   add_import modname
 
@@ -635,7 +780,7 @@ let acknowledge_pers_struct check modname {Persistent_signature.filename; cmi} =
   if ps.ps_name <> modname then
     error (Illegal_renaming (modname, ps.ps_name, filename));
   if check then check_consistency ps;
-  Hashtbl.add (persistent_structures ()) modname (Some ps);
+  add_persistent_structure modname (Some ps);
   ps
 
 let read_pers_struct check modname filename =
@@ -656,7 +801,7 @@ let find_pers_struct check name =
         match !Persistent_signature.load ~unit_name:name with
         | Some ps -> ps
         | None ->
-          Hashtbl.add (persistent_structures ()) name None;
+          add_persistent_structure name None;
           raise Not_found
       in
       add_import name;
@@ -701,6 +846,7 @@ let check_pers_struct name =
 
 let reset_cache () =
   current_unit () := "";
+  clear_type_path_memo ();
   Hashtbl.clear (persistent_structures ());
   clear_imports ();
   Hashtbl.clear (value_declarations ());
@@ -717,13 +863,16 @@ let reset_cache_toplevel () =
       (persistent_structures ()) []
   in
   List.iter (Hashtbl.remove (persistent_structures ())) l;
+  clear_type_path_memo ();
   Hashtbl.clear (value_declarations ());
   Hashtbl.clear (type_declarations ());
   Hashtbl.clear (module_declarations ());
   Hashtbl.clear (used_constructors ());
   Hashtbl.clear (prefixed_sg ())
 
-let set_unit_name name = current_unit () := name
+let set_unit_name name =
+  clear_type_path_memo ();
+  current_unit () := name
 
 let get_unit_name () = !(current_unit ())
 
@@ -770,11 +919,36 @@ let type_of_cstr path = function
     (d, ([], List.map snd (Datarepr.labels_of_type path d)))
   | _ -> assert false
 
+(* [find_type_full p env] for a regular type path, memoized as described at
+   [Type_path_memo]. *)
+let find_regular_type_full p env =
+  match !(can_load_cmis ()) with
+  | Cannot_load_cmis _ -> find_type_full p env
+  | Can_load_cmis -> (
+    match Type_path_memo.hash p with
+    | exception Type_path_memo.Not_cacheable -> find_type_full p env
+    | h -> (
+      let memo, _ = Domain.DLS.get type_path_memo_key in
+      let epoch = type_path_memo_epoch () in
+      match Type_path_memo.find memo h with
+      | Some e
+        when e.epoch == epoch
+             && (match p with
+               | Pident _ -> e.types == env.types
+               | Pdot _ | Papply _ -> e.components == env.components)
+             && Type_path_memo.equal e.path p ->
+        e.result
+      | _ ->
+        let result = find_type_full p env in
+        Type_path_memo.store memo h ~epoch ~path:p ~types:env.types ~modules:()
+          ~components:env.components ~result;
+        result))
+
 let find_type_full path env =
   match Path.constructor_typath path with
   | Regular p -> (
     try (Path_map.find p env.local_constraints, ([], []))
-    with Not_found -> find_type_full p env)
+    with Not_found -> find_regular_type_full p env)
   | Cstr (ty_path, s) ->
     let _, (cstrs, _) =
       try find_type_full ty_path env with Not_found -> assert false
@@ -881,6 +1055,41 @@ let normalize_path oloc env path =
     | Some loc ->
       raise (Error (Missing_module (loc, path, normalize_path true env path))))
 
+(* [Some (normalize_path None env path)] when normalization changes [path]
+   ([Path.same] is false), [None] when it leaves it unchanged. The unchanged
+   outcome is memoized as described at [Type_path_memo]; a changed path is
+   always computed afresh, so callers get a new path value as before. *)
+(* Not a closure in [normalize_type_path], which would be allocated on every
+   call, hits included. *)
+let normalize_type_path_uncached env path =
+  let path' = normalize_path None env path in
+  if Path.same path path' then None else Some path'
+
+let normalize_type_path env path =
+  match !(can_load_cmis ()) with
+  | Cannot_load_cmis _ -> normalize_type_path_uncached env path
+  | Can_load_cmis -> (
+    match Type_path_memo.hash path with
+    | exception Type_path_memo.Not_cacheable ->
+      normalize_type_path_uncached env path
+    | h -> (
+      let _, memo = Domain.DLS.get type_path_memo_key in
+      let epoch = type_path_memo_epoch () in
+      match Type_path_memo.find memo h with
+      | Some e
+        when e.epoch == epoch && e.modules == env.modules
+             && e.components == env.components
+             && Type_path_memo.equal e.path path ->
+        None
+      | _ ->
+        let result = normalize_type_path_uncached env path in
+        (match result with
+        | None ->
+          Type_path_memo.store memo h ~epoch ~path ~types:()
+            ~modules:env.modules ~components:env.components ~result:()
+        | Some _ -> ());
+        result))
+
 let normalize_path_prefix oloc env path =
   match path with
   | Pdot (p, s, pos) -> Pdot (normalize_path oloc env p, s, pos)
@@ -1953,7 +2162,25 @@ let reset_initial_for_request () =
 let with_fresh_initial action =
   let previous = initial_safe_string () in
   Fun.protect action ~finally:(fun () ->
-      Domain.DLS.set initial_safe_string_key previous)
+      Domain.DLS.set initial_safe_string_key previous;
+      (* The per-domain memo tables would otherwise keep the finished
+         request's environments alive until their slots are overwritten. Their
+         entries are already invalid outside the request's epoch. *)
+      clear_type_path_memo_tables ())
+
+module For_test = struct
+  let type_path_memo_entries () =
+    let found, unchanged = Domain.DLS.get type_path_memo_key in
+    let count (t : _ Type_path_memo.t) =
+      Array.fold_left
+        (fun n e ->
+          match e with
+          | None -> n
+          | Some _ -> n + 1)
+        0 t.entries
+    in
+    count found + count unchanged
+end
 
 let last_env_key = Request_local.new_key (fun () -> ref empty)
 let last_env () = Domain.DLS.get last_env_key

compiler/ml/env.mli

@@ -26,6 +26,8 @@ val empty : t
 val initial_safe_string : unit -> t
 val reset_initial_for_request : unit -> unit
 val with_fresh_initial : (unit -> 'a) -> 'a
+(* Runs a request with its own initial environment; on return, also empties
+   this domain's type path memo. *)
 
 type type_descriptions = constructor_description list * label_description list
 
@@ -65,6 +67,10 @@ val normalize_path : Location.t option -> t -> Path.t -> Path.t
    If the option is None, allow returning dangling paths.
    Otherwise raise a Missing_module error, and may add forgotten
    head as required global. *)
+val normalize_type_path : t -> Path.t -> Path.t option
+(* [normalize_type_path env path] is [Some (normalize_path None env path)] when
+   that differs from [path] according to [Path.same], and [None] otherwise. *)
+
 val normalize_path_prefix : Location.t option -> t -> Path.t -> Path.t
 (* Only normalize the prefix part of the path *)
 
@@ -285,3 +291,8 @@ module Persistent_signature : sig
       the .cmi file in the load path. This function can be overridden to load
       it from memory, for instance to build a self-contained toplevel. *)
 end
+
+module For_test : sig
+  val type_path_memo_entries : unit -> int
+  (** The type path memo entries held on this domain. *)
+end

compiler/ml/typecore.ml

@@ -814,9 +814,10 @@ let rec expand_path env p =
     match repr ty with
     | {desc = Tconstr (p, _, _)} -> expand_path env p
     | _ -> p (* PR#6394: recursive module may introduce incoherent manifest *))
-  | _ ->
-    let p' = Env.normalize_path None env p in
-    if Path.same p p' then p else expand_path env p'
+  | _ -> (
+    match Env.normalize_type_path env p with
+    | None -> p
+    | Some p' -> expand_path env p')
 
 let compare_type_path env tpath1 tpath2 =
   Path.same (expand_path env tpath1) (expand_path env tpath2)