Typer domain

src/frontend/ocamlmerlin/new/new_merlin.ml

@@ -109,15 +109,14 @@ let run =
               (float_of_int (60 * Mconfig.(config.merlin.cache_lifespan)))
             ();
           File_id.with_cache @@ fun () ->
+          let store = Mpipeline.Cache.get config in
+          Local_store.open_store store;
           let source = Msource.make (Misc.string_of_file stdin) in
-          let pipeline = Mpipeline.make config source in
+          let pipeline = Mpipeline.get config source in
           let json =
             let class_, message =
               Printexc.record_backtrace true;
-              match
-                Mpipeline.with_pipeline pipeline @@ fun () ->
-                command_action pipeline command_args
-              with
+              match command_action pipeline command_args with
               | result -> ("return", result)
               | exception Failure str ->
                 let trace = Printexc.get_backtrace () in
@@ -133,6 +132,7 @@ let run =
                   Location.print_main Format.str_formatter err;
                   ("error", `String (Format.flush_str_formatter ())))
             in
+            Local_store.close_store store;
             let cpu_time = Misc.time_spent () -. start_cpu in
             let gc_stats = Gc.quick_stat () in
             let heap_mbytes =

src/frontend/ocamlmerlin/ocamlmerlin_server.ml

@@ -56,16 +56,25 @@ module Server = struct
     | None -> Logger.log ~section:"server" ~title:"cannot setup listener" ""
     | Some server ->
       (* If the client closes its connection, don't let it kill us with a SIGPIPE. *)
+      let domain_typer = Domain.spawn Mpipeline.domain_typer in
       if Sys.unix then Sys.set_signal Sys.sigpipe Sys.Signal_ignore;
       loop (File_id.get Sys.executable_name) server;
+
+      Atomic.set Mpipeline.close_typer `True;
+      Domain.join domain_typer;
       Os_ipc.server_close server
 end
 
 let main () =
   (* Setup env for extensions *)
   Unix.putenv "__MERLIN_MASTER_PID" (string_of_int (Unix.getpid ()));
   match List.tl (Array.to_list Sys.argv) with
-  | "single" :: args -> exit (New_merlin.run ~new_env:None None args)
+  | "single" :: args ->
+    let domain_typer = Domain.spawn Mpipeline.domain_typer in
+    let vexit = New_merlin.run ~new_env:None None args in
+    Atomic.set Mpipeline.close_typer `True;
+    Domain.join domain_typer;
+    exit vexit
   | "old-protocol" :: args -> Old_merlin.run args
   | [ "server"; socket_path; socket_fd ] -> Server.start socket_path socket_fd
   | ("-help" | "--help" | "-h" | "server") :: _ ->

src/kernel/mocaml.mli

@@ -1,5 +1,5 @@
 (* An instance of load path, environment cache & btype unification log  *)
-type typer_state
+type typer_state = Local_store.store
 
 val new_state : unit -> typer_state
 val with_state : typer_state -> (unit -> 'a) -> 'a

src/kernel/mpipeline.ml

@@ -4,23 +4,22 @@ let { Logger.log } = Logger.for_section "Pipeline"
 
 let time_shift = ref 0.0
 
-let timed_lazy r x =
-  lazy
-    (let start = Misc.time_spent () in
-     let time_shift0 = !time_shift in
-     let update () =
-       let delta = Misc.time_spent () -. start in
-       let shift = !time_shift -. time_shift0 in
-       time_shift := time_shift0 +. delta;
-       r := !r +. delta -. shift
-     in
-     match Lazy.force x with
-     | x ->
-       update ();
-       x
-     | exception exn ->
-       update ();
-       Std.reraise exn)
+let timed r x =
+  let start = Misc.time_spent () in
+  let time_shift0 = !time_shift in
+  let update () =
+    let delta = Misc.time_spent () -. start in
+    let shift = !time_shift -. time_shift0 in
+    time_shift := time_shift0 +. delta;
+    r := !r +. delta -. shift
+  in
+  match x () with
+  | x ->
+    update ();
+    x
+  | exception exn ->
+    update ();
+    Std.reraise exn
 
 module Cache = struct
   let cache = ref []
@@ -65,7 +64,7 @@ module Cache = struct
 end
 
 module Typer = struct
-  type t = { errors : exn list lazy_t; result : Mtyper.result }
+  type t = { errors : exn list; result : Mtyper.result }
 end
 
 module Ppx = struct
@@ -82,10 +81,10 @@ type t =
   { config : Mconfig.t;
     state : Mocaml.typer_state;
     raw_source : Msource.t;
-    source : (Msource.t * Mreader.parsetree option) lazy_t;
-    reader : Reader.t lazy_t;
-    ppx : Ppx.t lazy_t;
-    typer : Typer.t lazy_t;
+    source : Msource.t * Mreader.parsetree option;
+    reader : Reader.t;
+    ppx : Ppx.t;
+    typer : Typer.t;
     pp_time : float ref;
     reader_time : float ref;
     ppx_time : float ref;
@@ -99,7 +98,7 @@ type t =
 let raw_source t = t.raw_source
 
 let input_config t = t.config
-let input_source t = fst (Lazy.force t.source)
+let input_source t = fst t.source
 
 let with_pipeline t f =
   Mocaml.with_state t.state @@ fun () ->
@@ -110,10 +109,10 @@ let get_lexing_pos t pos =
     ~filename:(Mconfig.filename t.config)
     pos
 
-let reader t = Lazy.force t.reader
+let reader t = t.reader
 
-let ppx t = Lazy.force t.ppx
-let typer t = Lazy.force t.typer
+let ppx t = t.ppx
+let typer t = t.typer
 
 let reader_config t = (reader t).config
 let reader_parsetree t = (reader t).result.Mreader.parsetree
@@ -131,7 +130,7 @@ let ppx_errors t = (ppx t).Ppx.errors
 let final_config t = (ppx t).Ppx.config
 
 let typer_result t = (typer t).Typer.result
-let typer_errors t = Lazy.force (typer t).Typer.errors
+let typer_errors t = (typer t).Typer.errors
 
 module Reader_phase = struct
   type t =
@@ -230,9 +229,8 @@ let process ?state ?(pp_time = ref 0.0) ?(reader_time = ref 0.0)
     | Some state -> state
   in
   let source =
-    timed_lazy pp_time
-      (lazy
-        (match Mconfig.(config.ocaml.pp) with
+    timed pp_time (fun () ->
+        match Mconfig.(config.ocaml.pp) with
         | None -> (raw_source, None)
         | Some { workdir; workval } -> (
           let source = Msource.text raw_source in
@@ -242,73 +240,67 @@ let process ?state ?(pp_time = ref 0.0) ?(reader_time = ref 0.0)
               ~source ~pp:workval
           with
           | `Source source -> (Msource.make source, None)
-          | (`Interface _ | `Implementation _) as ast -> (raw_source, Some ast))))
+          | (`Interface _ | `Implementation _) as ast -> (raw_source, Some ast)))
   in
   let reader =
-    timed_lazy reader_time
-      (lazy
-        (let (lazy ((_, pp_result) as source)) = source in
-         let config = Mconfig.normalize config in
-         Mocaml.setup_reader_config config;
-         let cache_disabling =
-           match (config.merlin.use_ppx_cache, pp_result) with
-           | false, _ -> Some "configuration"
-           | true, Some _ ->
-             (* The cache could be refined in the future to also act on the
+    timed reader_time (fun () ->
+        let ((_, pp_result) as source) = source in
+        let config = Mconfig.normalize config in
+        Mocaml.setup_reader_config config;
+        let cache_disabling =
+          match (config.merlin.use_ppx_cache, pp_result) with
+          | false, _ -> Some "configuration"
+          | true, Some _ ->
+            (* The cache could be refined in the future to also act on the
                 PP phase. For now, let's disable the whole cache when there's
                 a PP. *)
-             Some "source preprocessor usage"
-           | true, None -> None
-         in
-         let { Reader_with_cache.output = { result; cache_version };
-               cache_was_hit
-             } =
-           Reader_with_cache.apply ~cache_disabling
-             { source; for_completion; config }
-         in
-         reader_cache_hit := cache_was_hit;
-         let cache_version =
-           if Option.is_some cache_disabling then None else Some cache_version
-         in
-         { Reader.result; config; cache_version }))
+            Some "source preprocessor usage"
+          | true, None -> None
+        in
+        let { Reader_with_cache.output = { result; cache_version };
+              cache_was_hit
+            } =
+          Reader_with_cache.apply ~cache_disabling
+            { source; for_completion; config }
+        in
+        reader_cache_hit := cache_was_hit;
+        let cache_version =
+          if Option.is_some cache_disabling then None else Some cache_version
+        in
+        { Reader.result; config; cache_version })
   in
   let ppx =
-    timed_lazy ppx_time
-      (lazy
-        (let (lazy
-               { Reader.result = { Mreader.parsetree; _ };
-                 config;
-                 cache_version
-               }) =
-           reader
-         in
-         let caught = ref [] in
-         Msupport.catch_errors Mconfig.(config.ocaml.warnings) caught
-         @@ fun () ->
-         (* Currently the cache is invalidated even for source changes that don't
+    timed ppx_time (fun () ->
+        let { Reader.result = { Mreader.parsetree; _ }; config; cache_version }
+            =
+          reader
+        in
+        let caught = ref [] in
+        Msupport.catch_errors Mconfig.(config.ocaml.warnings) caught
+        @@ fun () ->
+        (* Currently the cache is invalidated even for source changes that don't
              change the parsetree. To avoid that, we'd have to digest the
              parsetree in the cache. *)
-         let cache_disabling, reader_cache =
-           match cache_version with
-           | Some v -> (None, Ppx_phase.Version v)
-           | None -> (Some "reader cache is disabled", Off)
-         in
-         let { Ppx_with_cache.output = parsetree; cache_was_hit } =
-           Ppx_with_cache.apply ~cache_disabling
-             { parsetree; config; reader_cache }
-         in
-         ppx_cache_hit := cache_was_hit;
-         { Ppx.config; parsetree; errors = !caught }))
+        let cache_disabling, reader_cache =
+          match cache_version with
+          | Some v -> (None, Ppx_phase.Version v)
+          | None -> (Some "reader cache is disabled", Off)
+        in
+        let { Ppx_with_cache.output = parsetree; cache_was_hit } =
+          Ppx_with_cache.apply ~cache_disabling
+            { parsetree; config; reader_cache }
+        in
+        ppx_cache_hit := cache_was_hit;
+        { Ppx.config; parsetree; errors = !caught })
   in
   let typer =
-    timed_lazy typer_time
-      (lazy
-        (let (lazy { Ppx.config; parsetree; _ }) = ppx in
-         Mocaml.setup_typer_config config;
-         let result = Mtyper.run config parsetree in
-         let errors = timed_lazy error_time (lazy (Mtyper.get_errors result)) in
-         typer_cache_stats := Mtyper.get_cache_stat result;
-         { Typer.errors; result }))
+    timed typer_time (fun () ->
+        let { Ppx.config; parsetree; _ } = ppx in
+        Mocaml.setup_typer_config config;
+        let result = Mtyper.run config parsetree in
+        let errors = timed error_time (fun () -> Mtyper.get_errors result) in
+        typer_cache_stats := Mtyper.get_cache_stat result;
+        { Typer.errors; result })
   in
   { config;
     state;
@@ -373,3 +365,45 @@ let cache_information t =
       ("cmt", cmt);
       ("cmi", cmi)
     ]
+
+let shared_config = Atomic.make None
+let shared_pipeline = Atomic.make None
+
+let close_typer = Atomic.make `False
+
+let domain_typer () =
+  let rec loop () =
+    if Atomic.get close_typer = `True then ()
+    else
+      match Atomic.get shared_config with
+      | None ->
+        Domain.cpu_relax ();
+        loop ()
+      | Some (config, source) as curr -> (
+        try
+          let pipeline = make config source in
+          if Atomic.compare_and_set shared_config curr None then
+            Atomic.set shared_pipeline (Some pipeline);
+          loop ()
+        with exn -> Atomic.set close_typer (`Exn exn))
+  in
+  loop ()
+
+let get config source =
+  Atomic.set shared_config (Some (config, source));
+
+  let rec loop count =
+    match Atomic.get shared_pipeline with
+    | None -> begin
+      match Atomic.get close_typer with
+      | `Exn exn -> raise exn
+      | `True -> assert false
+      | _ ->
+        Domain.cpu_relax ();
+        loop (count + 1)
+    end
+    | Some pipeline ->
+      Atomic.set shared_pipeline None;
+      pipeline
+  in
+  loop 0

src/kernel/mpipeline.mli

@@ -1,5 +1,7 @@
 type t
 val make : Mconfig.t -> Msource.t -> t
+val get : Mconfig.t -> Msource.t -> t
+
 val with_pipeline : t -> (unit -> 'a) -> 'a
 val for_completion : Msource.position -> t -> t
 
@@ -27,3 +29,11 @@ val typer_errors : t -> exn list
 
 val timing_information : t -> (string * float) list
 val cache_information : t -> Std.json
+
+module Cache : sig
+  val get : Mconfig.t -> Mocaml.typer_state
+end
+
+val close_typer : [ `True | `False | `Exn of exn ] Atomic.t
+
+val domain_typer : unit -> unit

src/ocaml/utils/local_store.ml

@@ -57,3 +57,12 @@ let with_store slots f =
     List.iter (fun (Slot s) -> s.value <- !(s.ref)) slots;
     global_bindings.is_bound <- false;
   )
+
+  let open_store slots = 
+    assert (not global_bindings.is_bound);
+    global_bindings.is_bound <- true;
+    List.iter (fun (Slot {ref;value}) -> ref := value) slots
+
+   let close_store slots = 
+    List.iter (fun (Slot s) -> s.value <- !(s.ref)) slots;
+    global_bindings.is_bound <- false    

src/ocaml/utils/local_store.mli

@@ -65,3 +65,7 @@ val reset : unit -> unit
 val is_bound : unit -> bool
 (** Returns [true] when a store is active (i.e. when called from the callback
     passed to {!with_store}), [false] otherwise. *)
+
+val open_store : store -> unit
+
+val close_store : store -> unit