refactor(stdune): improve path tables (#8052)

Use a more memory efficient path table. Instead of using the variant for
the key, combine 3 tables all for the individual paths.

This makes the empty table a little more bloated (3x bigger), but gives
us a saving of 2 words for every single key we store.

Signed-off-by: Rudi Grinberg <[email protected]>

otherlibs/stdune/src/path.ml

@@ -290,6 +290,8 @@ module Local : sig
 
   include Path_intf.S with type t := t
 
+  module Table : Hashtbl.S with type key = t
+
   val root : t
 
   val is_root : t -> bool
@@ -369,6 +371,8 @@ end
 module External : sig
   include Path_intf.S
 
+  module Table : Hashtbl.S with type key = t
+
   val relative : t -> string -> t
 
   val mkdir_p : ?perms:int -> t -> unit
@@ -741,8 +745,6 @@ module T : sig
       Path of the same kind are compared using the standard lexical order *)
   val compare : t -> t -> Ordering.t
 
-  val equal : t -> t -> bool
-
   val hash : t -> int
 
   val in_build_dir : Local.t -> t
@@ -766,8 +768,6 @@ end = struct
     | In_build_dir _, In_source_tree _ -> Gt
     | In_build_dir x, In_build_dir y -> Local.compare x y
 
-  let equal (x : t) (y : t) = x = y
-
   let hash = Poly.hash
 
   let in_build_dir s = In_build_dir s
@@ -784,8 +784,6 @@ end
 
 include T
 
-let hash (t : t) = Poly.hash t
-
 let build_dir = in_build_dir Local.root
 
 let is_root = function
@@ -1167,16 +1165,6 @@ let touch ?(create = true) p =
     (* OCaml PR#8857 *)
     create ()
 
-let compare x y =
-  match (x, y) with
-  | External x, External y -> External.compare x y
-  | External _, _ -> Lt
-  | _, External _ -> Gt
-  | In_source_tree x, In_source_tree y -> Local.compare x y
-  | In_source_tree _, _ -> Lt
-  | _, In_source_tree _ -> Gt
-  | In_build_dir x, In_build_dir y -> Local.compare x y
-
 let extension t =
   match t with
   | External t -> External.extension t
@@ -1220,7 +1208,77 @@ let source s = in_source_tree s
 
 let build s = in_build_dir s
 
-module Table = Hashtbl.Make (T)
+module Table = struct
+  type key = t
+
+  type 'a t =
+    { source : 'a Source0.Table.t
+    ; build : 'a Build.Table.t
+    ; external_ : 'a External.Table.t
+    }
+
+  let create () =
+    { source = Source0.Table.create 0
+    ; build = Build.Table.create 0
+    ; external_ = External.Table.create 0
+    }
+
+  let clear { source; build; external_ } =
+    Source0.Table.clear source;
+    Build.Table.clear build;
+    External.Table.clear external_
+
+  let[@inline] mem { source; build; external_ } key =
+    match key with
+    | In_source_tree p -> Source0.Table.mem source p
+    | In_build_dir p -> Build.Table.mem build p
+    | External p -> External.Table.mem external_ p
+
+  let[@inline] set { source; build; external_ } k v =
+    match k with
+    | In_source_tree p -> Source0.Table.set source p v
+    | In_build_dir p -> Build.Table.set build p v
+    | External p -> External.Table.set external_ p v
+
+  let[@inline] remove { source; build; external_ } k =
+    match k with
+    | In_source_tree p -> Source0.Table.remove source p
+    | In_build_dir p -> Build.Table.remove build p
+    | External p -> External.Table.remove external_ p
+
+  let iter { source; build; external_ } ~f =
+    Source0.Table.iter source ~f;
+    Build.Table.iter build ~f;
+    External.Table.iter external_ ~f
+
+  let[@inline] find { source; build; external_ } = function
+    | In_source_tree p -> Source0.Table.find source p
+    | In_build_dir p -> Build.Table.find build p
+    | External p -> External.Table.find external_ p
+
+  let filteri_inplace { source; build; external_ } ~f =
+    Source0.Table.filteri_inplace source ~f:(fun [@inline] ~key ~data ->
+        f ~key:(In_source_tree key) ~data);
+    Build.Table.filteri_inplace build ~f:(fun [@inline] ~key ~data ->
+        f ~key:(In_build_dir key) ~data);
+    External.Table.filteri_inplace external_ ~f:(fun [@inline] ~key ~data ->
+        f ~key:(External key) ~data)
+
+  let filter_inplace { source; build; external_ } ~f =
+    Source0.Table.filteri_inplace source ~f:(fun [@inline] ~key:_ ~data ->
+        f data);
+    Build.Table.filteri_inplace build ~f:(fun [@inline] ~key:_ ~data -> f data);
+    External.Table.filteri_inplace external_ ~f:(fun [@inline] ~key:_ ~data ->
+        f data)
+
+  let to_dyn f { source; build; external_ } =
+    let open Dyn in
+    record
+      [ ("source", Source0.Table.to_dyn f source)
+      ; ("build", Build.Table.to_dyn f build)
+      ; ("external_", External.Table.to_dyn f external_)
+      ]
+end
 
 module L = struct
   (* TODO more efficient implementation *)

otherlibs/stdune/src/path.mli

@@ -84,6 +84,8 @@ module External : sig
   val of_filename_relative_to_initial_cwd : string -> t
 
   val append_local : t -> Local.t -> t
+
+  module Table : Hashtbl.S with type key = t
 end
 
 (** In the source section of the current workspace. *)
@@ -113,6 +115,8 @@ module Source : sig
   val descendant : t -> of_:t -> t option
 
   val to_local : t -> Local.t
+
+  module Table : Hashtbl.S with type key = t
 end
 
 module Permissions : sig
@@ -221,6 +225,8 @@ module Build : sig
   val lstat : t -> Unix.stats
 
   val unlink_no_err : t -> unit
+
+  module Table : Hashtbl.S with type key = t
 end
 
 type t = private
@@ -230,6 +236,39 @@ type t = private
 
 include Path_intf.S with type t := t
 
+module Table : sig
+  (** Specialized tables for path. We do implement all of [Hashtbl_intf.S] -
+      only what we use in dune. *)
+
+  type path := t
+
+  type key = path
+
+  type 'a t
+
+  val create : unit -> 'a t
+
+  val clear : 'a t -> unit
+
+  val mem : 'a t -> path -> bool
+
+  val set : 'a t -> path -> 'a -> unit
+
+  val remove : 'a t -> path -> unit
+
+  val iter : 'a t -> f:('a -> unit) -> unit
+
+  val find : 'a t -> path -> 'a option
+
+  val filteri_inplace : 'a t -> f:(key:path -> data:'a -> bool) -> unit
+
+  val filter_inplace : 'a t -> f:('a -> bool) -> unit
+
+  val to_dyn : ('a -> Dyn.t) -> 'a t -> Dyn.t
+end
+
+val equal : t -> t -> bool
+
 val as_outside_build_dir_exn : t -> Outside_build_dir.t
 
 val destruct_build_dir :

otherlibs/stdune/src/path_intf.ml

@@ -42,7 +42,7 @@ module type S = sig
     val of_listing : dir:elt -> filenames:string list -> t
   end
 
-  module Table : Hashtbl.S with type key = t
+  val equal : t -> t -> bool
 
   val relative : ?error_loc:Loc0.t -> t -> string -> t
 

src/dune_engine/action_builder.ml

@@ -82,7 +82,8 @@ let alias a = dep (Dep.alias a)
 let contents =
   let read_file =
     Memo.exec
-      (Memo.create "Action_builder.contents"
+      (Memo.create_with_store "Action_builder.contents"
+         ~store:(module Path.Table)
          ~input:(module Path)
          ~cutoff:String.equal
          (fun p -> Build_system.read_file p ~f:Io.read_file))

src/dune_engine/build_system.ml

@@ -1069,7 +1069,10 @@ end = struct
          | `Disabled -> None
          | `Enabled -> Some (Tuple.T2.equal Digest.equal target_kind_equal)
        in
-       Memo.create "build-file" ~input:(module Path) ?cutoff build_file_impl)
+       Memo.create_with_store "build-file"
+         ~store:(module Path.Table)
+         ~input:(module Path)
+         ?cutoff build_file_impl)
 
   let build_file path = Memo.exec (Lazy.force build_file_memo) path >>| fst
 

src/dune_engine/cached_digest.ml

@@ -70,7 +70,7 @@ let cache =
   lazy
     (match P.load db_file with
     | None ->
-      { checked_key = 0; table = Path.Table.create 1024; max_timestamp = 0. }
+      { checked_key = 0; table = Path.Table.create (); max_timestamp = 0. }
     | Some cache ->
       cache.checked_key <- cache.checked_key + 1;
       cache)

src/dune_engine/load_rules.ml

@@ -885,7 +885,12 @@ end = struct
 
   let load_dir =
     let load_dir_impl dir = load_dir_impl ~dir in
-    let memo = Memo.create "load-dir" ~input:(module Path) load_dir_impl in
+    let memo =
+      Memo.create_with_store "load-dir"
+        ~store:(module Path.Table)
+        ~input:(module Path)
+        load_dir_impl
+    in
     fun ~dir -> Memo.exec memo dir
 
   let is_under_directory_target p =

src/dune_engine/rule_cache.ml

@@ -55,7 +55,7 @@ module Workspace_local = struct
         (* This mutable table is safe: it's only used by [execute_rule_impl] to
            decide whether to rebuild a rule or not; [execute_rule_impl] ensures
            that the targets are produced deterministically. *)
-        | None -> Path.Table.create 1024)
+        | None -> Path.Table.create ())
 
     let dump () =
       if !needs_dumping && Path.build_dir_exists () then (