Use Eio for parallelizing some analysis commands

analysis/bin/dune

@@ -8,4 +8,4 @@
  (package analysis)
  (modes byte exe)
  (name main)
- (libraries analysis))
+ (libraries analysis eio_main))

analysis/src/Commands.ml

@@ -178,100 +178,107 @@ let typeDefinition ~path ~pos ~debug =
     | Some location -> location |> Protocol.stringifyLocation)
 
 let references ~path ~pos ~debug =
-  let allLocs =
-    match Cmt.loadFullCmtFromPath ~path with
-    | None -> []
-    | Some full -> (
-      match References.getLocItem ~full ~pos ~debug with
-      | None -> []
-      | Some locItem ->
-        let allReferences = References.allReferencesForLocItem ~full locItem in
-        allReferences
-        |> List.fold_left
-             (fun acc {References.uri = uri2; locOpt} ->
-               let loc =
-                 match locOpt with
-                 | Some loc -> loc
-                 | None -> Uri.toTopLevelLoc uri2
-               in
-               Protocol.stringifyLocation
-                 {uri = Uri.toString uri2; range = Utils.cmtLocToRange loc}
-               :: acc)
-             [])
-  in
-  print_endline
-    (if allLocs = [] then Protocol.null
-     else "[\n" ^ (allLocs |> String.concat ",\n") ^ "\n]")
+  Eio_main.run (fun env ->
+      let allLocs =
+        match Cmt.loadFullCmtFromPath ~path with
+        | None -> []
+        | Some full -> (
+          match References.getLocItem ~full ~pos ~debug with
+          | None -> []
+          | Some locItem ->
+            References.allReferencesForLocItem ~domain_mgr:env#domain_mgr ~full
+              locItem
+            |> List.fold_left
+                 (fun acc {References.uri = uri2; locOpt} ->
+                   let loc =
+                     match locOpt with
+                     | Some loc -> loc
+                     | None -> Uri.toTopLevelLoc uri2
+                   in
+                   Protocol.stringifyLocation
+                     {uri = Uri.toString uri2; range = Utils.cmtLocToRange loc}
+                   :: acc)
+                 [])
+      in
+      print_endline
+        (if allLocs = [] then Protocol.null
+         else "[\n" ^ (allLocs |> String.concat ",\n") ^ "\n]"))
 
 let rename ~path ~pos ~newName ~debug =
   let result =
-    match Cmt.loadFullCmtFromPath ~path with
-    | None -> Protocol.null
-    | Some full -> (
-      match References.getLocItem ~full ~pos ~debug with
-      | None -> Protocol.null
-      | Some locItem ->
-        let allReferences = References.allReferencesForLocItem ~full locItem in
-        let referencesToToplevelModules =
-          allReferences
-          |> Utils.filterMap (fun {References.uri = uri2; locOpt} ->
-                 if locOpt = None then Some uri2 else None)
-        in
-        let referencesToItems =
-          allReferences
-          |> Utils.filterMap (function
-               | {References.uri = uri2; locOpt = Some loc} -> Some (uri2, loc)
-               | {locOpt = None} -> None)
-        in
-        let fileRenames =
-          referencesToToplevelModules
-          |> List.map (fun uri ->
-                 let path = Uri.toPath uri in
-                 let dir = Filename.dirname path in
-                 let newPath =
-                   Filename.concat dir (newName ^ Filename.extension path)
-                 in
-                 let newUri = Uri.fromPath newPath in
-                 Protocol.
-                   {
-                     oldUri = uri |> Uri.toString;
-                     newUri = newUri |> Uri.toString;
-                   })
-        in
-        let textDocumentEdits =
-          let module StringMap = Misc.StringMap in
-          let textEditsByUri =
-            referencesToItems
-            |> List.map (fun (uri, loc) -> (Uri.toString uri, loc))
-            |> List.fold_left
-                 (fun acc (uri, loc) ->
-                   let textEdit =
+    Eio_main.run (fun env ->
+        match Cmt.loadFullCmtFromPath ~path with
+        | None -> Protocol.null
+        | Some full -> (
+          match References.getLocItem ~full ~pos ~debug with
+          | None -> Protocol.null
+          | Some locItem ->
+            let allReferences =
+              References.allReferencesForLocItem ~domain_mgr:env#domain_mgr
+                ~full locItem
+            in
+            let referencesToToplevelModules =
+              allReferences
+              |> Utils.filterMap (fun {References.uri = uri2; locOpt} ->
+                     if locOpt = None then Some uri2 else None)
+            in
+            let referencesToItems =
+              allReferences
+              |> Utils.filterMap (function
+                   | {References.uri = uri2; locOpt = Some loc} ->
+                     Some (uri2, loc)
+                   | {locOpt = None} -> None)
+            in
+            let fileRenames =
+              referencesToToplevelModules
+              |> List.map (fun uri ->
+                     let path = Uri.toPath uri in
+                     let dir = Filename.dirname path in
+                     let newPath =
+                       Filename.concat dir (newName ^ Filename.extension path)
+                     in
+                     let newUri = Uri.fromPath newPath in
                      Protocol.
-                       {range = Utils.cmtLocToRange loc; newText = newName}
-                   in
-                   match StringMap.find_opt uri acc with
-                   | None -> StringMap.add uri [textEdit] acc
-                   | Some prevEdits ->
-                     StringMap.add uri (textEdit :: prevEdits) acc)
-                 StringMap.empty
-          in
-          StringMap.fold
-            (fun uri edits acc ->
-              let textDocumentEdit =
-                Protocol.{textDocument = {uri; version = None}; edits}
+                       {
+                         oldUri = uri |> Uri.toString;
+                         newUri = newUri |> Uri.toString;
+                       })
+            in
+            let textDocumentEdits =
+              let module StringMap = Misc.StringMap in
+              let textEditsByUri =
+                referencesToItems
+                |> List.map (fun (uri, loc) -> (Uri.toString uri, loc))
+                |> List.fold_left
+                     (fun acc (uri, loc) ->
+                       let textEdit =
+                         Protocol.
+                           {range = Utils.cmtLocToRange loc; newText = newName}
+                       in
+                       match StringMap.find_opt uri acc with
+                       | None -> StringMap.add uri [textEdit] acc
+                       | Some prevEdits ->
+                         StringMap.add uri (textEdit :: prevEdits) acc)
+                     StringMap.empty
               in
-              textDocumentEdit :: acc)
-            textEditsByUri []
-        in
-        let fileRenamesString =
-          fileRenames |> List.map Protocol.stringifyRenameFile
-        in
-        let textDocumentEditsString =
-          textDocumentEdits |> List.map Protocol.stringifyTextDocumentEdit
-        in
-        "[\n"
-        ^ (fileRenamesString @ textDocumentEditsString |> String.concat ",\n")
-        ^ "\n]")
+              StringMap.fold
+                (fun uri edits acc ->
+                  let textDocumentEdit =
+                    Protocol.{textDocument = {uri; version = None}; edits}
+                  in
+                  textDocumentEdit :: acc)
+                textEditsByUri []
+            in
+            let fileRenamesString =
+              fileRenames |> List.map Protocol.stringifyRenameFile
+            in
+            let textDocumentEditsString =
+              textDocumentEdits |> List.map Protocol.stringifyTextDocumentEdit
+            in
+            "[\n"
+            ^ (fileRenamesString @ textDocumentEditsString
+              |> String.concat ",\n")
+            ^ "\n]"))
   in
   print_endline result
 

analysis/src/References.ml

@@ -1,6 +1,6 @@
 open SharedTypes
 
-let debugReferences = ref true
+let debugReferences = ref false
 let maybeLog m = if !debugReferences then Log.log ("[ref] " ^ m)
 
 let checkPos (line, char)
@@ -433,7 +433,34 @@ type references = {
   locOpt: Location.t option; (* None: reference to a toplevel module *)
 }
 
-let forLocalStamp ~full:{file; extra; package} stamp (tip : Tip.t) =
+(* Single helper for parallel work distribution over a list. *)
+let parallel_map ~domain_mgr ~items ~f =
+  let len = List.length items in
+  let doms = Domain.recommended_domain_count () in
+  let chunk_size = max 1 (if doms <= 1 then len else (len + doms - 1) / doms) in
+  let rec chunks_of size lst =
+    if size <= 0 then [lst]
+    else
+      match lst with
+      | [] -> []
+      | _ ->
+        let rec take n acc rest =
+          if n = 0 then (List.rev acc, rest)
+          else
+            match rest with
+            | [] -> (List.rev acc, [])
+            | x :: xs -> take (n - 1) (x :: acc) xs
+        in
+        let chunk, rest = take size [] lst in
+        chunk :: chunks_of size rest
+  in
+  let chunks = chunks_of chunk_size items in
+  Eio.Fiber.List.map
+    (fun chunk -> Eio.Domain_manager.run domain_mgr (fun () -> f chunk))
+    chunks
+  |> List.concat
+
+let forLocalStamp ~domain_mgr ~full:{file; extra; package} stamp (tip : Tip.t) =
   let env = QueryEnv.fromFile file in
   match
     match tip with
@@ -485,30 +512,34 @@ let forLocalStamp ~full:{file; extra; package} stamp (tip : Tip.t) =
             in
             maybeLog ("Now checking path " ^ pathToString path);
             let thisModuleName = file.moduleName in
-            let externals =
+            let names =
               package.projectFiles |> FileSet.elements
               |> List.filter (fun name -> name <> file.moduleName)
-              |> List.map (fun moduleName ->
-                     Cmt.fullsFromModule ~package ~moduleName
-                     |> List.map (fun {file; extra} ->
-                            match
-                              Hashtbl.find_opt extra.externalReferences
-                                thisModuleName
-                            with
-                            | None -> []
-                            | Some refs ->
-                              let locs =
-                                refs
-                                |> Utils.filterMap (fun (p, t, locs) ->
-                                       if p = path && t = tip then Some locs
-                                       else None)
-                              in
-                              locs
-                              |> List.map (fun loc ->
-                                     {uri = file.uri; locOpt = Some loc})))
-              |> List.concat |> List.concat
             in
-            alternativeReferences @ externals)
+            let results =
+              parallel_map ~domain_mgr ~items:names ~f:(fun chunk ->
+                  let process_module moduleName =
+                    Cmt.fullsFromModule ~package ~moduleName
+                    |> List.map (fun {file; extra} ->
+                           match
+                             Hashtbl.find_opt extra.externalReferences
+                               thisModuleName
+                           with
+                           | None -> []
+                           | Some refs ->
+                             let locs =
+                               refs
+                               |> Utils.filterMap (fun (p, t, locs) ->
+                                      if p = path && t = tip then Some locs
+                                      else None)
+                             in
+                             locs
+                             |> List.map (fun loc ->
+                                    {uri = file.uri; locOpt = Some loc}))
+                  in
+                  chunk |> List.map process_module |> List.concat |> List.concat)
+            in
+            alternativeReferences @ results)
           else (
             maybeLog "Not visible";
             [])
@@ -517,26 +548,36 @@ let forLocalStamp ~full:{file; extra; package} stamp (tip : Tip.t) =
         (locs |> List.map (fun loc -> {uri = file.uri; locOpt = Some loc}))
         externals)
 
-let allReferencesForLocItem ~full:({file; package} as full) locItem =
+let allReferencesForLocItem ~domain_mgr ~full:({file; package} as full) locItem
+    =
   match locItem.locType with
   | TopLevelModule moduleName ->
     let otherModulesReferences =
-      package.projectFiles |> FileSet.elements
-      |> Utils.filterMap (fun name ->
-             match ProcessCmt.fileForModule ~package name with
-             | None -> None
-             | Some file -> Cmt.fullFromUri ~uri:file.uri)
-      |> List.map (fun full ->
-             match Hashtbl.find_opt full.extra.fileReferences moduleName with
-             | None -> []
-             | Some locs ->
-               locs |> LocationSet.elements
-               |> List.map (fun loc ->
-                      {
-                        uri = Uri.fromPath loc.Location.loc_start.pos_fname;
-                        locOpt = Some loc;
-                      }))
-      |> List.flatten
+      let names = package.projectFiles |> FileSet.elements in
+      let per_chunk =
+        parallel_map ~domain_mgr ~items:names ~f:(fun chunk ->
+            let process_name name =
+              match ProcessCmt.fileForModule ~package name with
+              | None -> []
+              | Some file -> (
+                match Cmt.fullFromUri ~uri:file.uri with
+                | None -> []
+                | Some full -> (
+                  match
+                    Hashtbl.find_opt full.extra.fileReferences moduleName
+                  with
+                  | None -> []
+                  | Some locs ->
+                    locs |> LocationSet.elements
+                    |> List.map (fun loc ->
+                           {
+                             uri = Uri.fromPath loc.Location.loc_start.pos_fname;
+                             locOpt = Some loc;
+                           })))
+            in
+            chunk |> List.map process_name |> List.concat)
+      in
+      per_chunk
     in
     let targetModuleReferences =
       match Hashtbl.find_opt package.pathsForModule moduleName with
@@ -547,13 +588,13 @@ let allReferencesForLocItem ~full:({file; package} as full) locItem =
     in
     List.append targetModuleReferences otherModulesReferences
   | Typed (_, _, NotFound) | LModule NotFound | Constant _ -> []
-  | TypeDefinition (_, _, stamp) -> forLocalStamp ~full stamp Type
+  | TypeDefinition (_, _, stamp) -> forLocalStamp ~domain_mgr ~full stamp Type
   | Typed (_, _, (LocalReference (stamp, tip) | Definition (stamp, tip)))
   | LModule (LocalReference (stamp, tip) | Definition (stamp, tip)) ->
     maybeLog
       ("Finding references for " ^ Uri.toString file.uri ^ " and stamp "
      ^ string_of_int stamp ^ " and tip " ^ Tip.toString tip);
-    forLocalStamp ~full stamp tip
+    forLocalStamp ~domain_mgr ~full stamp tip
   | LModule (GlobalReference (moduleName, path, tip))
   | Typed (_, _, GlobalReference (moduleName, path, tip)) -> (
     match ProcessCmt.fileForModule ~package moduleName with
@@ -570,4 +611,4 @@ let allReferencesForLocItem ~full:({file; package} as full) locItem =
             ("Finding references for (global) " ^ Uri.toString env.file.uri
            ^ " and stamp " ^ string_of_int stamp ^ " and tip "
            ^ Tip.toString tip);
-          forLocalStamp ~full stamp tip)))
+          forLocalStamp ~domain_mgr ~full stamp tip)))

analysis/src/dune

@@ -2,4 +2,4 @@
  (name analysis)
  (flags
   (-w "+6+26+27+32+33+39"))
- (libraries unix str ext ml jsonlib syntax reanalyze))
+ (libraries unix str ext ml jsonlib syntax reanalyze eio eio_main))