Skip to content

Index lazily #195

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
liam923 merged 2 commits into from
Nov 7, 2025
Merged

Index lazily #195

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6a18258
Refactor index caching. Get rid of the stamping.
voodoos Nov 6, 2025
6e7fc02
Only index the buffer if the query needs lit
voodoos Nov 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 11 additions & 7 deletions src/analysis/index_occurrences.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
7 changes: 4 additions & 3 deletions src/analysis/occurrences.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
118 changes: 62 additions & 56 deletions src/kernel/mtyper.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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 =
Expand All @@ -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
Expand All @@ -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 };
Expand All @@ -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
}

Expand All @@ -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
Expand All @@ -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;
Expand All @@ -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);
Expand All @@ -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;
Expand All @@ -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 =
Expand All @@ -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;
Expand All @@ -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 = _;
Expand All @@ -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 =
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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

Expand Down
10 changes: 4 additions & 6 deletions src/kernel/mtyper.mli
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down