Merge pull request #195 from oxcaml/index-lazily

Index lazily

Author: liam923

src/analysis/index_occurrences.ml

@@ -44,8 +44,8 @@ let should_ignore_lid (lid : Longident.t Location.loc) =
   *)
   Location.is_none lid.loc
 
-let iterator ~current_buffer_path ~index ~stamp ~reduce_for_uid =
-  let add uid loc = Stamped_hashtable.add index ~stamp (uid, loc) () in
+let iterator ~current_buffer_path ~index ~reduce_for_uid =
+  let add uid loc = index := Shape.Uid.Map.add_to_list uid loc !index 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 lid = { lid with loc = set_fname ~file:current_buffer_path lid.loc } in
@@ -89,7 +89,7 @@ let iterator ~current_buffer_path ~index ~stamp ~reduce_for_uid =
   in
   Ast_iterators.iterator_on_usages ~include_hidden:true ~f
 
-let items ~index ~stamp (config : Mconfig.t) items =
+let items index (config : Mconfig.t) items =
   let module Shape_reduce = Shape_reduce.Make (struct
     let fuel () = Misc_stdlib.Maybe_bounded.of_int 10
 
@@ -123,7 +123,11 @@ let items ~index ~stamp (config : Mconfig.t) items =
     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
+  let index = ref index in
+  let iterator = iterator ~current_buffer_path ~index ~reduce_for_uid in
+  let () =
+    match items with
+    | `Impl items -> List.iter ~f:(iterator.structure_item iterator) items
+    | `Intf items -> List.iter ~f:(iterator.signature_item iterator) items
+  in
+  !index

src/analysis/occurrences.ml

@@ -186,10 +186,11 @@ end = struct
 end
 
 let get_buffer_locs result uid =
-  Stamped_hashtable.fold
-    (fun (uid', loc) () acc ->
+  Shape.Uid.Map.fold
+    (fun uid' lids acc ->
       if Shape.Uid.equal uid uid' then
-        Lid_set.add (Index_format.Lid.of_lid loc) acc
+        List.fold_left lids ~init:acc ~f:(fun acc lid ->
+            Lid_set.add (Index_format.Lid.of_lid lid) acc)
       else acc)
     (Mtyper.get_index result) Lid_set.empty
 

src/kernel/mtyper.ml

@@ -3,21 +3,17 @@ 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
+type index = Longident.t Location.loc list Shape.Uid.Map.t
 
 (* Forward ref to be filled by analysis.Occurrences *)
 let index_items :
-    (index:index_tbl ->
-    stamp:int ->
+    (index ->
     Mconfig.t ->
     [ `Impl of Typedtree.structure_item list
     | `Intf of Typedtree.signature_item list ] ->
-    unit)
+    index)
     ref =
-  ref (fun ~index:_ ~stamp:_ _config _item -> ())
+  ref (fun acc _config _item -> acc)
 let set_index_items f = index_items := f
 
 type ('p, 't) item =
@@ -30,7 +26,8 @@ type ('p, 't) item =
     part_rev_sg : Types.signature_item list;
     part_errors : exn list;
     part_checks : Typecore.delayed_check list;
-    part_warnings : Warnings.state
+    part_warnings : Warnings.state;
+    part_index : index lazy_t
   }
 
 type typedtree =
@@ -53,8 +50,7 @@ type 'a cache_result =
     snapshot : Types.snapshot;
     ident_stamp : int;
     uid_stamp : int;
-    value : 'a;
-    index : (Shape.Uid.t * Longident.t Location.loc, unit) Stamped_hashtable.t
+    value : 'a
   }
 
 let cache : typedtree_items option cache_result option ref = s_ref None
@@ -72,8 +68,7 @@ let get_cache config =
   | Some ({ snapshot; _ } as c) when Types.is_valid snapshot -> c
   | Some _ | None ->
     let env, snapshot, ident_stamp, uid_stamp = fresh_env config in
-    let index = Stamped_hashtable.create !index_changelog 256 in
-    { env; snapshot; ident_stamp; uid_stamp; value = None; index }
+    { env; snapshot; ident_stamp; uid_stamp; value = None }
 
 let return_and_cache status =
   cache := Some { status with value = Some status.value };
@@ -87,7 +82,6 @@ type result =
     stamp : int;
     initial_uid_stamp : int;
     typedtree : typedtree_items;
-    index : (Shape.Uid.t * Longident.t Location.loc, unit) Stamped_hashtable.t;
     cache_stat : typer_cache_stats
   }
 
@@ -111,11 +105,14 @@ let compatible_prefix_rev result_items tree_items =
   in
   aux [] (result_items, tree_items)
 
-let[@tail_mod_cons] rec type_structure caught env sg = function
+let[@tail_mod_cons] rec type_structure config caught env index sg = function
   | parsetree_item :: rest ->
     let items, sg', part_env =
       Typemod.merlin_type_structure env sg [ parsetree_item ]
     in
+    let part_index =
+      lazy (!index_items (Lazy.force index) config (`Impl items.str_items))
+    in
     let typedtree_items =
       (items.Typedtree.str_items, items.Typedtree.str_type)
     in
@@ -130,14 +127,15 @@ let[@tail_mod_cons] rec type_structure caught env sg = function
         part_uid = Shape.Uid.get_current_stamp ();
         part_errors = !caught;
         part_checks = !Typecore.delayed_checks;
-        part_warnings = Warnings.backup ()
+        part_warnings = Warnings.backup ();
+        part_index
       }
     in
-    item :: type_structure caught part_env part_rev_sg rest
+    item :: type_structure config caught part_env part_index part_rev_sg rest
   | [] -> []
 
-let[@tail_mod_cons] rec type_signature caught env sg psg_modalities psg_loc =
-  function
+let[@tail_mod_cons] rec type_signature config caught env index sg psg_modalities
+    psg_loc = function
   | parsetree_item :: rest ->
     let { Typedtree.sig_final_env = part_env;
           sig_items;
@@ -150,6 +148,9 @@ let[@tail_mod_cons] rec type_signature caught env sg psg_modalities psg_loc =
            [ parsetree_item ])
     in
     let part_rev_sg = List.rev_append sig_type sg in
+    let part_index =
+      lazy (!index_items (Lazy.force index) config (`Intf sig_items))
+    in
     let item =
       { parsetree_item;
         typedtree_items = (sig_items, sig_type);
@@ -160,26 +161,34 @@ let[@tail_mod_cons] rec type_signature caught env sg psg_modalities psg_loc =
         part_uid = Shape.Uid.get_current_stamp ();
         part_errors = !caught;
         part_checks = !Typecore.delayed_checks;
-        part_warnings = Warnings.backup ()
+        part_warnings = Warnings.backup ();
+        part_index
       }
     in
     item
-    :: type_signature caught part_env part_rev_sg psg_modalities psg_loc rest
+    :: type_signature config caught part_env part_index part_rev_sg
+         psg_modalities psg_loc rest
   | [] -> []
 
 let type_implementation config caught parsetree =
-  let { env; snapshot; ident_stamp; uid_stamp; value = cached_value; index; _ }
-      =
+  let { env; snapshot; ident_stamp; uid_stamp; value = cached_value; _ } =
     get_cache config
   in
   let rev_prefix, parsetree_suffix, cache_stats =
     match cached_value with
     | Some (Implementation_items items) -> compatible_prefix_rev items parsetree
     | Some (Interface_items _) | None -> ([], parsetree, Miss)
   in
-  let env', sg', snap', stamp', uid_stamp', warn' =
+  let env', sg', snap', stamp', uid_stamp', warn', index' =
     match rev_prefix with
-    | [] -> (env, [], snapshot, ident_stamp, uid_stamp, Warnings.backup ())
+    | [] ->
+      ( env,
+        [],
+        snapshot,
+        ident_stamp,
+        uid_stamp,
+        Warnings.backup (),
+        lazy Shape.Uid.Map.empty )
     | x :: _ ->
       caught := x.part_errors;
       Typecore.delayed_checks := x.part_checks;
@@ -188,30 +197,21 @@ let type_implementation config caught parsetree =
         x.part_snapshot,
         x.part_stamp,
         x.part_uid,
-        x.part_warnings )
+        x.part_warnings,
+        x.part_index )
   in
   Btype.backtrack snap';
   Warnings.restore warn';
   Env.cleanup_functor_caches ~stamp:stamp';
-  let stamp = List.length rev_prefix - 1 in
-  Stamped_hashtable.backtrack !index_changelog ~stamp;
   Env.cleanup_usage_tables ~stamp:uid_stamp';
   Shape.Uid.restore_stamp uid_stamp';
-  let suffix = type_structure caught env' sg' parsetree_suffix 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 suffix = type_structure config caught env' index' sg' parsetree_suffix in
   let value = Implementation_items (List.rev_append rev_prefix suffix) in
-  ( return_and_cache { env; snapshot; ident_stamp; uid_stamp; value; index },
+  ( return_and_cache { env; snapshot; ident_stamp; uid_stamp; value },
     cache_stats )
 
 let type_interface config caught (parsetree : Parsetree.signature) =
-  let { env; snapshot; ident_stamp; uid_stamp; value = cached_value; index; _ }
-      =
+  let { env; snapshot; ident_stamp; uid_stamp; value = cached_value; _ } =
     get_cache config
   in
   let rev_prefix, parsetree_suffix, cache_stats =
@@ -230,9 +230,16 @@ let type_interface config caught (parsetree : Parsetree.signature) =
     | Some (Interface_items _) | Some (Implementation_items _) | None ->
       ([], parsetree.psg_items, Miss)
   in
-  let env', sg', snap', stamp', uid_stamp', warn' =
+  let env', sg', snap', stamp', uid_stamp', warn', index' =
     match rev_prefix with
-    | [] -> (env, [], snapshot, ident_stamp, uid_stamp, Warnings.backup ())
+    | [] ->
+      ( env,
+        [],
+        snapshot,
+        ident_stamp,
+        uid_stamp,
+        Warnings.backup (),
+        lazy Shape.Uid.Map.empty )
     | x :: _ ->
       caught := x.part_errors;
       Typecore.delayed_checks := x.part_checks;
@@ -241,25 +248,17 @@ let type_interface config caught (parsetree : Parsetree.signature) =
         x.part_snapshot,
         x.part_stamp,
         x.part_uid,
-        x.part_warnings )
+        x.part_warnings,
+        x.part_index )
   in
   Btype.backtrack snap';
   Warnings.restore warn';
   Env.cleanup_functor_caches ~stamp:stamp';
-  let stamp = List.length rev_prefix in
-  Stamped_hashtable.backtrack !index_changelog ~stamp;
   Env.cleanup_usage_tables ~stamp:uid_stamp';
   Shape.Uid.restore_stamp uid_stamp';
   let suffix =
-    type_signature caught env' sg' parsetree.psg_modalities parsetree.psg_loc
-      parsetree_suffix
-  in
-  let () =
-    List.iteri
-      ~f:(fun i { typedtree_items = items, _; _ } ->
-        let stamp = stamp + i + 1 in
-        !index_items ~index ~stamp config (`Intf items))
-      suffix
+    type_signature config caught env' index' sg' parsetree.psg_modalities
+      parsetree.psg_loc parsetree_suffix
   in
   (* transl an empty signature to get the sig_modalities and sig_sloc *)
   let ({ sig_final_env = _;
@@ -281,7 +280,7 @@ let type_interface config caught (parsetree : Parsetree.signature) =
         sig_sloc
       }
   in
-  ( return_and_cache { env; snapshot; ident_stamp; uid_stamp; value; index },
+  ( return_and_cache { env; snapshot; ident_stamp; uid_stamp; value },
     cache_stats )
 
 let run config parsetree =
@@ -313,7 +312,6 @@ let run config parsetree =
     stamp;
     initial_uid_stamp = cached_result.uid_stamp;
     typedtree = cached_result.value;
-    index = cached_result.index;
     cache_stat
   }
 
@@ -362,7 +360,15 @@ let get_typedtree t =
         sig_sloc
       }
 
-let get_index t = t.index
+let get_index t =
+  let of_items items =
+    List.last items
+    |> Option.value_map ~default:Shape.Uid.Map.empty
+         ~f:(fun { part_index; _ } -> Lazy.force part_index)
+  in
+  match t.typedtree with
+  | Implementation_items items -> of_items items
+  | Interface_items { items; _ } -> of_items items
 
 let get_stamp t = t.stamp
 

src/kernel/mtyper.mli

@@ -14,16 +14,14 @@ type typedtree =
 
 type typer_cache_stats = Miss | Hit of { reused : int; typed : int }
 
-type index_tbl =
-  (Shape.Uid.t * Longident.t Location.loc, unit) Stamped_hashtable.t
+type index = Longident.t Location.loc list Shape.Uid.Map.t
 
 val set_index_items :
-  (index:index_tbl ->
-  stamp:int ->
+  (index ->
   Mconfig.t ->
   [ `Impl of Typedtree.structure_item list
   | `Intf of Typedtree.signature_item list ] ->
-  unit) ->
+  index) ->
   unit
 
 val run : Mconfig.t -> Mreader.parsetree -> result
@@ -32,7 +30,7 @@ val get_env : ?pos:Msource.position -> result -> Env.t
 
 val get_typedtree : result -> typedtree
 
-val get_index : result -> index_tbl
+val get_index : result -> index
 
 val get_stamp : result -> int