Use Eio for parallelizing some analysis commands

analysis/bin/main.ml

@@ -119,7 +119,7 @@ let main () =
   | [_; "cache-delete"; rootPath] -> (
     Cfg.readProjectConfigCache := false;
     let uri = Uri.fromPath rootPath in
-    match Packages.findRoot ~uri (Hashtbl.create 0) with
+    match Packages.findRoot ~uri with
     | Some (`Bs rootPath) -> (
       match BuildSystem.getLibBs rootPath with
       | None -> print_endline "\"ERR\""

analysis/src/AnalysisCache.ml

@@ -0,0 +1,7 @@
+(* Helpers for domain-local caches *)
+
+let make_hashtbl (size : int) : ('k, 'v) Hashtbl.t Domain.DLS.key =
+  Domain.DLS.new_key (fun () -> Hashtbl.create size)
+
+let get_hashtbl (key : ('k, 'v) Hashtbl.t Domain.DLS.key) : ('k, 'v) Hashtbl.t =
+  Domain.DLS.get key

analysis/src/Cmt.ml

@@ -1,5 +1,12 @@
 open SharedTypes
 
+module FullCache = struct
+  let key : (string, full) Hashtbl.t Domain.DLS.key =
+    AnalysisCache.make_hashtbl 64
+
+  let get () : (string, full) Hashtbl.t = AnalysisCache.get_hashtbl key
+end
+
 let fullForCmt ~moduleName ~package ~uri cmt =
   match Shared.tryReadCmt cmt with
   | None -> None
@@ -8,45 +15,64 @@ let fullForCmt ~moduleName ~package ~uri cmt =
     let extra = ProcessExtra.getExtra ~file ~infos in
     Some {file; extra; package}
 
-let fullFromUri ~uri =
+let fullFromUriWithPackage ~package ~uri =
   let path = Uri.toPath uri in
-  match Packages.getPackage ~uri with
-  | None -> None
-  | Some package -> (
-    let moduleName =
-      BuildSystem.namespacedName package.namespace (FindFiles.getName path)
-    in
-    let incremental =
-      if !Cfg.inIncrementalTypecheckingMode then
-        let incrementalCmtPath =
-          package.rootPath ^ "/lib/bs/___incremental" ^ "/" ^ moduleName
-          ^
-          match Files.classifySourceFile path with
-          | Resi -> ".cmti"
-          | _ -> ".cmt"
-        in
-        fullForCmt ~moduleName ~package ~uri incrementalCmtPath
-      else None
+  let moduleName =
+    BuildSystem.namespacedName package.namespace (FindFiles.getName path)
+  in
+  let cached_full cmt_path =
+    let cache = FullCache.get () in
+    match Hashtbl.find_opt cache cmt_path with
+    | Some v -> Some v
+    | None -> (
+      match fullForCmt ~moduleName ~package ~uri cmt_path with
+      | Some v as res ->
+        Hashtbl.replace cache cmt_path v;
+        res
+      | None -> None)
+  in
+  if !Cfg.inIncrementalTypecheckingMode then
+    let incrementalCmtPath =
+      package.rootPath ^ "/lib/bs/___incremental" ^ "/" ^ moduleName
+      ^
+      match Files.classifySourceFile path with
+      | Resi -> ".cmti"
+      | _ -> ".cmt"
     in
-    match incremental with
-    | Some cmtInfo ->
-      if Debug.verbose () then Printf.printf "[cmt] Found incremental cmt\n";
-      Some cmtInfo
+    match cached_full incrementalCmtPath with
+    | Some _ as x -> x
     | None -> (
+      (* Fallback to non-incremental *)
       match Hashtbl.find_opt package.pathsForModule moduleName with
       | Some paths ->
         let cmt = getCmtPath ~uri paths in
-        fullForCmt ~moduleName ~package ~uri cmt
+        cached_full cmt
       | None ->
         prerr_endline ("can't find module " ^ moduleName);
-        None))
+        None)
+  else
+    match Hashtbl.find_opt package.pathsForModule moduleName with
+    | Some paths ->
+      let cmt = getCmtPath ~uri paths in
+      cached_full cmt
+    | None ->
+      prerr_endline ("can't find module " ^ moduleName);
+      None
+
+let fullFromUri ~uri =
+  match Packages.getPackage ~uri with
+  | None -> None
+  | Some package -> fullFromUriWithPackage ~package ~uri
 
 let fullsFromModule ~package ~moduleName =
-  if Hashtbl.mem package.pathsForModule moduleName then
-    let paths = Hashtbl.find package.pathsForModule moduleName in
+  match Hashtbl.find_opt package.pathsForModule moduleName with
+  | None -> []
+  | Some paths ->
     let uris = getUris paths in
-    uris |> List.filter_map (fun uri -> fullFromUri ~uri)
-  else []
+    uris
+    |> List.filter_map (fun uri ->
+           let cmt = getCmtPath ~uri paths in
+           fullForCmt ~moduleName ~package ~uri cmt)
 
 let loadFullCmtFromPath ~path =
   let uri = Uri.fromPath path in

analysis/src/Packages.ml

@@ -1,5 +1,17 @@
 open SharedTypes
 
+(* Domain-local caches for packages and URI->root mapping. *)
+module LocalCache = struct
+  let packages_key : (string, package) Hashtbl.t Domain.DLS.key =
+    AnalysisCache.make_hashtbl 1
+
+  let roots_key : (Uri.t, string) Hashtbl.t Domain.DLS.key =
+    AnalysisCache.make_hashtbl 30
+
+  let packages () = AnalysisCache.get_hashtbl packages_key
+  let roots () = AnalysisCache.get_hashtbl roots_key
+end
+
 (* Creates the `pathsForModule` hashtbl, which maps a `moduleName` to it's `paths` (the ml/re, mli/rei, cmt, and cmti files) *)
 let makePathsForModule ~projectFilesAndPaths ~dependenciesFilesAndPaths =
   let pathsForModule = Hashtbl.create 30 in
@@ -200,14 +212,11 @@ let newBsPackage ~rootPath =
       Log.log ("Unable to read " ^ bsconfigJson);
       None)
 
-let findRoot ~uri packagesByRoot =
+let findRoot ~uri =
   let path = Uri.toPath uri in
   let rec loop path =
     if path = "/" then None
-    else if
-      SharedTypes.StateSync.with_lock (fun () ->
-          Hashtbl.mem packagesByRoot path)
-    then Some (`Root path)
+    else if Hashtbl.mem (LocalCache.packages ()) path then Some (`Root path)
     else if
       Files.exists (Filename.concat path "rescript.json")
       || Files.exists (Filename.concat path "bsconfig.json")
@@ -219,29 +228,29 @@ let findRoot ~uri packagesByRoot =
   loop (if Sys.is_directory path then path else Filename.dirname path)
 
 let getPackage ~uri =
-  let open SharedTypes in
-  match
-    SharedTypes.StateSync.with_lock (fun () ->
-        if Hashtbl.mem state.rootForUri uri then
-          let root = Hashtbl.find state.rootForUri uri in
-          Some (Hashtbl.find state.packagesByRoot root)
-        else None)
-  with
-  | Some pkg -> Some pkg
+  let roots = LocalCache.roots () in
+  let packages = LocalCache.packages () in
+  match Hashtbl.find_opt roots uri with
+  | Some root -> Hashtbl.find_opt packages root
   | None -> (
-    match findRoot ~uri state.packagesByRoot with
+    match findRoot ~uri with
     | None ->
       Log.log "No root directory found";
       None
-    | Some (`Root rootPath) ->
-      SharedTypes.StateSync.with_lock (fun () ->
-          Hashtbl.replace state.rootForUri uri rootPath;
-          Some (Hashtbl.find state.packagesByRoot rootPath))
+    | Some (`Root rootPath) -> (
+      Hashtbl.replace roots uri rootPath;
+      match Hashtbl.find_opt packages rootPath with
+      | Some pkg -> Some pkg
+      | None -> (
+        match newBsPackage ~rootPath with
+        | Some pkg ->
+          Hashtbl.replace packages rootPath pkg;
+          Some pkg
+        | None -> None))
     | Some (`Bs rootPath) -> (
       match newBsPackage ~rootPath with
       | None -> None
       | Some package ->
-        SharedTypes.StateSync.with_lock (fun () ->
-            Hashtbl.replace state.rootForUri uri package.rootPath;
-            Hashtbl.replace state.packagesByRoot package.rootPath package;
-            Some package)))
+        Hashtbl.replace roots uri package.rootPath;
+        Hashtbl.replace packages package.rootPath package;
+        Some package))

analysis/src/ProcessCmt.ml

@@ -773,27 +773,16 @@ let fileForCmtInfos ~moduleName ~uri
   | _ -> File.create moduleName uri
 
 let fileForCmt ~moduleName ~cmt ~uri =
-  (* Double-checked locking: fast path under lock; if missing, compute without
-     holding the lock, then insert under lock if still absent. *)
-  match
-    SharedTypes.StateSync.with_lock (fun () ->
-        Hashtbl.find_opt state.cmtCache cmt)
-  with
+  let local = SharedTypes.CmtCache.get () in
+  match Hashtbl.find_opt local cmt with
   | Some file -> Some file
   | None -> (
     match Shared.tryReadCmt cmt with
     | None -> None
     | Some infos ->
       let file = fileForCmtInfos ~moduleName ~uri infos in
-      let cached =
-        SharedTypes.StateSync.with_lock (fun () ->
-            match Hashtbl.find_opt state.cmtCache cmt with
-            | Some f -> Some f
-            | None ->
-              Hashtbl.replace state.cmtCache cmt file;
-              Some file)
-      in
-      cached)
+      Hashtbl.replace local cmt file;
+      Some file)
 
 let fileForModule moduleName ~package =
   match Hashtbl.find_opt package.pathsForModule moduleName with

analysis/src/References.ml

@@ -219,7 +219,9 @@ let alternateDeclared ~(file : File.t) ~package (declared : _ Declared.t) tip =
       maybeLog
         ("alternateDeclared for " ^ file.moduleName ^ " has both resi and res");
       let alternateUri = if Uri.isInterface file.uri then res else resi in
-      match Cmt.fullFromUri ~uri:(Uri.fromPath alternateUri) with
+      match
+        Cmt.fullFromUriWithPackage ~package ~uri:(Uri.fromPath alternateUri)
+      with
       | None -> None
       | Some {file; extra} -> (
         let env = QueryEnv.fromFile file in
@@ -568,7 +570,7 @@ let allReferencesForLocItem ~domain_mgr ~full:({file; package} as full) locItem
               match ProcessCmt.fileForModule ~package name with
               | None -> []
               | Some file -> (
-                match Cmt.fullFromUri ~uri:file.uri with
+                match Cmt.fullFromUriWithPackage ~package ~uri:file.uri with
                 | None -> []
                 | Some full -> (
                   match
@@ -612,7 +614,7 @@ let allReferencesForLocItem ~domain_mgr ~full:({file; package} as full) locItem
       match exportedForTip ~env ~path ~package ~tip with
       | None -> []
       | Some (env, _name, stamp) -> (
-        match Cmt.fullFromUri ~uri:env.file.uri with
+        match Cmt.fullFromUriWithPackage ~package ~uri:env.file.uri with
         | None -> []
         | Some full ->
           maybeLog

analysis/src/SharedTypes.ml

@@ -537,34 +537,15 @@ let initExtra () =
     locItems = [];
   }
 
-type state = {
-  packagesByRoot: (string, package) Hashtbl.t;
-  rootForUri: (Uri.t, string) Hashtbl.t;
-  cmtCache: (filePath, File.t) Hashtbl.t;
-}
+module CmtCache = struct
+  let key : (filePath, File.t) Hashtbl.t Domain.DLS.key =
+    AnalysisCache.make_hashtbl 30
 
-(* There's only one state, so it can as well be global *)
-let state =
-  {
-    packagesByRoot = Hashtbl.create 1;
-    rootForUri = Hashtbl.create 30;
-    cmtCache = Hashtbl.create 30;
-  }
-
-module StateSync = struct
-  let mutex : Mutex.t = Mutex.create ()
-
-  let with_lock f =
-    Mutex.lock mutex;
-    match f () with
-    | v ->
-      Mutex.unlock mutex;
-      v
-    | exception exn ->
-      Mutex.unlock mutex;
-      raise exn
+  let get () : (filePath, File.t) Hashtbl.t = AnalysisCache.get_hashtbl key
 end
 
+module StringMap = Map.Make (String)
+
 let locKindToString = function
   | LocalReference (_, tip) -> "(LocalReference " ^ Tip.toString tip ^ ")"
   | GlobalReference _ -> "GlobalReference"