occurrences: incrementally index the buffer in Mtyper

src/analysis/ast_iterators.ml

@@ -139,3 +139,7 @@ let iter_on_usages ~f (local_defs : Mtyper.typedtree) =
   begin match local_defs with
   | `Interface signature -> iter.signature iter signature
   | `Implementation structure -> iter.structure iter structure end
+
+let iterator_on_usages ~f =
+  let occ_iter = Cmt_format.iter_on_occurrences ~f in
+  iter_only_visible occ_iter

src/analysis/index_occurrences.ml

@@ -0,0 +1,96 @@
+open Std
+module Lid_set = Index_format.Lid_set
+let {Logger. log} = Logger.for_section "index-occurrences"
+
+let set_fname ~file (loc : Location.t) =
+  let pos_fname = file in
+  { loc with
+      loc_start = { loc.loc_start with pos_fname };
+      loc_end   = { loc.loc_end with pos_fname }}
+
+let decl_of_path_or_lid env namespace path lid =
+  match (namespace : Shape.Sig_component_kind.t) with
+  | Constructor ->
+    begin match Env.find_constructor_by_name lid env with
+    | exception Not_found -> None
+    | {cstr_uid; cstr_loc; _ } ->
+      Some { Env_lookup.uid = cstr_uid; loc = cstr_loc; namespace }
+    end
+  | Label ->
+    begin match Env.find_label_by_name lid env with
+    | exception Not_found -> None
+    | {lbl_uid; lbl_loc; _ } ->
+      Some { Env_lookup.uid = lbl_uid; loc = lbl_loc; namespace }
+    end
+  | _ -> Env_lookup.by_path path namespace env
+
+let iterator ~current_buffer_path ~index ~stamp ~reduce_for_uid =
+  let add uid loc =
+    Stamped_hashtable.add index ~stamp (uid, loc) ()
+  in
+  let f ~namespace env path (lid : Longident.t Location.loc)  =
+    log ~title:"index_buffer" "Path: %a" Logger.fmt (Fun.flip Path.print path);
+    let not_ghost { Location.loc = { loc_ghost; _ }; _ } = not loc_ghost in
+    let lid = { lid with loc = set_fname ~file:current_buffer_path lid.loc } in
+    let index_decl () =
+      begin match decl_of_path_or_lid env namespace path lid.txt with
+      | exception _ |  None -> log ~title:"index_buffer" "Declaration not found"
+      | Some decl ->
+        log ~title:"index_buffer" "Found declaration: %a"
+          Logger.fmt (Fun.flip Location.print_loc decl.loc);
+        add decl.uid lid
+      end
+    in
+     if not_ghost lid then
+      match Env.shape_of_path ~namespace env path with
+      | exception Not_found -> ()
+      | path_shape ->
+        log ~title:"index_buffer" "Shape of path: %a"
+          Logger.fmt (Fun.flip Shape.print path_shape);
+        let result = reduce_for_uid env path_shape in
+        begin match Locate.uid_of_result ~traverse_aliases:false result with
+        | Some uid, false ->
+          log ~title:"index_buffer" "Found %a (%a) wiht uid %a"
+            Logger.fmt (Fun.flip Pprintast.longident lid.txt)
+            Logger.fmt (Fun.flip Location.print_loc lid.loc)
+            Logger.fmt (Fun.flip Shape.Uid.print uid);
+          add uid lid
+        | Some uid, true ->
+          log ~title:"index_buffer" "Shape is approximative, found uid: %a"
+            Logger.fmt (Fun.flip Shape.Uid.print uid);
+          index_decl ()
+        | None, _ ->
+          log ~title:"index_buffer" "Reduction failed: missing uid";
+          index_decl ()
+        end
+  in
+  Ast_iterators.iterator_on_usages ~f
+
+let items ~index ~stamp (config : Mconfig.t) items =
+  let module Shape_reduce =
+    Shape_reduce.Make (struct
+      let fuel = 10
+
+      let read_unit_shape ~unit_name =
+          log ~title:"read_unit_shape" "inspecting %s" unit_name;
+          let cmt = Format.sprintf "%s.cmt" unit_name in
+          match Cmt_cache.read (Load_path.find_normalized cmt) with
+          | { cmt_infos = { cmt_impl_shape; _ }; _ } ->
+            log ~title:"read_unit_shape" "shapes loaded for %s" unit_name;
+            cmt_impl_shape
+          | exception _ ->
+            log ~title:"read_unit_shape" "failed to find %s" unit_name;
+            None
+    end)
+  in
+  let current_buffer_path =
+    Filename.concat config.query.directory config.query.filename
+  in
+  let reduce_for_uid = Shape_reduce.reduce_for_uid in
+  let iterator = iterator ~current_buffer_path ~index ~stamp ~reduce_for_uid in
+  match items with
+  | `Impl items ->
+      List.iter ~f:(iterator.structure_item iterator) items
+  | `Intf items ->
+      List.iter ~f:(iterator.signature_item iterator) items
+

src/analysis/occurrences.ml

@@ -5,114 +5,14 @@ let {Logger. log} = Logger.for_section "occurrences"
 
 type t = { locs: Warnings.loc list; status: Query_protocol.occurrences_status }
 
+let () = Mtyper.set_index_items Index_occurrences.items
+
 let set_fname ~file (loc : Location.t) =
   let pos_fname = file in
   { loc with
       loc_start = { loc.loc_start with pos_fname };
       loc_end   = { loc.loc_end with pos_fname }}
 
-let decl_of_path_or_lid env namespace path lid =
-  match (namespace : Shape.Sig_component_kind.t) with
-  | Constructor ->
-    begin match Env.find_constructor_by_name lid env with
-    | exception Not_found -> None
-    | {cstr_uid; cstr_loc; _ } ->
-      Some { Env_lookup.uid = cstr_uid; loc = cstr_loc; namespace }
-    end
-  | Label ->
-    begin match Env.find_label_by_name lid env with
-    | exception Not_found -> None
-    | {lbl_uid; lbl_loc; _ } ->
-      Some { Env_lookup.uid = lbl_uid; loc = lbl_loc; namespace }
-    end
-  | _ -> Env_lookup.by_path path namespace env
-
-let index_buffer_ ~current_buffer_path ~local_defs () =
-  let {Logger. log} = Logger.for_section "index" in
-  let defs = Hashtbl.create 64 in
-  let add tbl uid locs =
-    try
-      let locations = Hashtbl.find tbl uid in
-      Hashtbl.replace tbl uid (Lid_set.union locs locations)
-    with Not_found -> Hashtbl.add tbl uid locs
-  in
-  let module Shape_reduce =
-    Shape_reduce.Make (struct
-      let fuel = 10
-
-      let read_unit_shape ~unit_name =
-          log ~title:"read_unit_shape" "inspecting %s" unit_name;
-          let cmt = Format.sprintf "%s.cmt" unit_name in
-          match Cmt_cache.read (Load_path.find_normalized cmt) with
-          | { cmt_infos = { cmt_impl_shape; _ }; _ } ->
-            log ~title:"read_unit_shape" "shapes loaded for %s" unit_name;
-            cmt_impl_shape
-          | exception _ ->
-            log ~title:"read_unit_shape" "failed to find %s" unit_name;
-            None
-    end)
-  in
-  let f ~namespace env path (lid : Longident.t Location.loc)  =
-    log ~title:"index_buffer" "Path: %a" Logger.fmt (Fun.flip Path.print path);
-    let not_ghost { Location.loc = { loc_ghost; _ }; _ } = not loc_ghost in
-    let lid = { lid with loc = set_fname ~file:current_buffer_path lid.loc } in
-    let index_decl () =
-      begin match decl_of_path_or_lid env namespace path lid.txt with
-      | exception _ |  None -> log ~title:"index_buffer" "Declaration not found"
-      | Some decl ->
-        log ~title:"index_buffer" "Found declaration: %a"
-          Logger.fmt (Fun.flip Location.print_loc decl.loc);
-        add defs decl.uid (Lid_set.singleton lid)
-      end
-    in
-     if not_ghost lid then
-      match Env.shape_of_path ~namespace env path with
-      | exception Not_found -> ()
-      | path_shape ->
-        log ~title:"index_buffer" "Shape of path: %a"
-          Logger.fmt (Fun.flip Shape.print path_shape);
-        let result =  Shape_reduce.reduce_for_uid env path_shape in
-        begin match Locate.uid_of_result ~traverse_aliases:false result with
-        | Some uid, false ->
-          log ~title:"index_buffer" "Found %a (%a) wiht uid %a"
-            Logger.fmt (Fun.flip Pprintast.longident lid.txt)
-            Logger.fmt (Fun.flip Location.print_loc lid.loc)
-            Logger.fmt (Fun.flip Shape.Uid.print uid);
-          add defs uid (Lid_set.singleton lid)
-        | Some uid, true ->
-          log ~title:"index_buffer" "Shape is approximative, found uid: %a"
-            Logger.fmt (Fun.flip Shape.Uid.print uid);
-          index_decl ()
-        | None, _ ->
-          log ~title:"index_buffer" "Reduction failed: missing uid";
-          index_decl ()
-        end
-  in
-  Ast_iterators.iter_on_usages ~f local_defs;
-  defs
-
-let index_buffer =
-  (* Right now, we only cache the last used index. We could do better by caching
-     the index for every known buffer. *)
-  let cache = ref None in
-  fun ~scope ~current_buffer_path ~stamp ~local_defs () ->
-    let {Logger. log} = Logger.for_section "index" in
-    match !cache with
-    | Some (path, stamp', scope', value) when
-        String.equal path current_buffer_path
-        && Int.equal stamp' stamp
-        && scope' = scope ->
-      log ~title:"index_cache" "Reusing cached value for path %s and stamp %i."
-        path stamp';
-      value
-    | _ ->
-      log ~title:"index_cache" "No valid cache found, reindexing.";
-      let result =
-        index_buffer_ ~current_buffer_path ~local_defs ()
-      in
-      cache := Some (current_buffer_path, stamp, scope, result);
-      result
-
 (* A longident can have the form: A.B.x Right now we are only interested in
    values, but we will eventually want to index all occurrences of modules in
    such longidents. However there is an issue with that: we only have the
@@ -210,6 +110,13 @@ end = struct
     | None -> cache_and_return (stat t file)
 end
 
+let get_buffer_locs result uid =
+  Stamped_hashtable.fold
+    (fun (uid', loc) () acc ->
+      if Shape.Uid.equal uid uid' then Lid_set.add loc acc else acc)
+    (Mtyper.get_index result)
+    Lid_set.empty
+
 let locs_of ~config ~env ~typer_result ~pos ~scope path =
   log ~title:"occurrences" "Looking for occurences of %s (pos: %s)"
     path
@@ -253,11 +160,7 @@ let locs_of ~config ~env ~typer_result ~pos ~scope path =
       Logger.fmt (fun fmt -> Shape.Uid.print fmt def_uid)
       Logger.fmt (fun fmt -> Location.print_loc fmt def_loc);
     log ~title:"locs_of" "Indexing current buffer";
-    let buffer_index =
-      let stamp = Mtyper.get_stamp typer_result in
-      index_buffer ~scope ~current_buffer_path ~stamp ~local_defs ()
-    in
-    let buffer_locs = Hashtbl.find_opt buffer_index def_uid in
+    let buffer_locs = get_buffer_locs typer_result def_uid in
     let external_locs =
       if scope = `Buffer then []
       else List.filter_map config.merlin.index_files ~f:(fun file ->
@@ -291,11 +194,7 @@ let locs_of ~config ~env ~typer_result ~pos ~scope path =
           (Lid_set.union acc_locs locs, String.Set.union acc_files files))
         (external_locs)
     in
-    let locs =
-      match buffer_locs with
-      | Some buffer_locs -> Lid_set.union buffer_locs external_locs
-      | None -> external_locs
-    in
+    let locs = Lid_set.union buffer_locs external_locs in
     let locs =
       log ~title:"occurrences" "Found %i locs" (Lid_set.cardinal locs);
       Lid_set.elements locs

src/kernel/mtyper.ml

@@ -3,6 +3,23 @@ open Local_store
 
 let {Logger. log} = Logger.for_section "Mtyper"
 
+let index_changelog =
+  Local_store.s_table Stamped_hashtable.create_changelog ()
+
+type index_tbl =
+  (Shape.Uid.t * Longident.t Location.loc, unit) Stamped_hashtable.t
+
+(* Forward ref to be filled by analysis.Occurrences *)
+let index_items :
+  (index:index_tbl
+    -> stamp:int
+    -> Mconfig.t
+    -> [ `Impl of Typedtree.structure_item list
+      | `Intf of Typedtree.signature_item list ]
+    -> unit) ref =
+  ref (fun ~index:_ ~stamp:_ _config _item -> ())
+let set_index_items f = index_items := f
+
 type ('p,'t) item = {
   parsetree_item: 'p;
   typedtree_items: 't list * Types.signature_item list;
@@ -33,6 +50,7 @@ type 'a cache_result = {
   snapshot : Types.snapshot;
   ident_stamp : int;
   value : 'a;
+  index : (Shape.Uid.t * Longident.t Location.loc, unit) Stamped_hashtable.t;
 }
 
 let cache : typedtree_items option cache_result option ref = s_ref None
@@ -49,7 +67,8 @@ let get_cache config =
   | Some ({ snapshot; _ } as c) when Types.is_valid snapshot -> c
   | Some _ | None ->
     let env, snapshot, ident_stamp = fresh_env config in
-    { env; snapshot; ident_stamp; value = None  }
+    let index = Stamped_hashtable.create !index_changelog 256 in
+    { env; snapshot; ident_stamp; value = None; index  }
 
 let return_and_cache status =
   cache := Some ({ status with value = Some status.value });
@@ -62,6 +81,7 @@ type result = {
   initial_stamp : int;
   stamp : int;
   typedtree : typedtree_items;
+  index : (Shape.Uid.t * Longident.t Location.loc, unit) Stamped_hashtable.t;
   cache_stat : typer_cache_stats
 }
 
@@ -119,7 +139,7 @@ let rec type_signature caught env = function
   | [] -> []
 
 let type_implementation config caught parsetree =
-  let { env; snapshot; ident_stamp; value = prefix; _ } = get_cache config in
+  let { env; snapshot; ident_stamp; value = prefix; index; _ } = get_cache config in
   let prefix, parsetree, cache_stats =
     match prefix with
     | Some (`Implementation items) -> compatible_prefix items parsetree
@@ -135,12 +155,19 @@ let type_implementation config caught parsetree =
   Btype.backtrack snap';
   Warnings.restore warn';
   Env.cleanup_functor_caches ~stamp:stamp';
+  let stamp = List.length prefix - 1 in
+  Stamped_hashtable.backtrack !index_changelog ~stamp;
   let suffix = type_structure caught env' parsetree in
+  let () =
+    List.iteri ~f:(fun i { typedtree_items = (items, _); _ } ->
+      let stamp  = stamp + i + 1 in
+      !index_items ~index ~stamp config (`Impl items)) suffix
+  in
   let value = `Implementation (List.rev_append prefix suffix) in
-  return_and_cache { env; snapshot; ident_stamp; value }, cache_stats
+  return_and_cache { env; snapshot; ident_stamp; value; index }, cache_stats
 
 let type_interface config caught parsetree =
-  let { env; snapshot; ident_stamp; value = prefix; _ } = get_cache config in
+  let { env; snapshot; ident_stamp; value = prefix; index; _ } = get_cache config in
   let prefix, parsetree, cache_stats =
     match prefix with
     | Some (`Interface items) -> compatible_prefix items parsetree
@@ -156,9 +183,16 @@ let type_interface config caught parsetree =
   Btype.backtrack snap';
   Warnings.restore warn';
   Env.cleanup_functor_caches ~stamp:stamp';
+  let stamp = List.length prefix in
+  Stamped_hashtable.backtrack !index_changelog ~stamp;
   let suffix = type_signature caught env' parsetree in
+  let () =
+    List.iteri ~f:(fun i { typedtree_items = (items, _); _ } ->
+      let stamp  = stamp + i + 1 in
+      !index_items ~index ~stamp config (`Intf items)) suffix
+  in
   let value = `Interface (List.rev_append prefix suffix) in
-  return_and_cache { env; snapshot; ident_stamp; value}, cache_stats
+  return_and_cache { env; snapshot; ident_stamp; value; index}, cache_stats
 
 let run config parsetree =
   if not (Env.check_state_consistency ()) then (
@@ -186,6 +220,7 @@ let run config parsetree =
     initial_stamp = cached_result.ident_stamp;
     stamp;
     typedtree = cached_result.value;
+    index = cached_result.index;
     cache_stat;
   }
 
@@ -224,6 +259,8 @@ let get_typedtree t =
     let sig_items, sig_type = split_items l in
     `Interface {Typedtree. sig_items; sig_type; sig_final_env = get_env t}
 
+let get_index t = t.index
+
 let get_stamp t = t.stamp
 
 let node_at ?(skip_recovered=false) t pos_cursor =