From 1fe268928b9cf9649b968d17186fe325296e85d3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ulysse=20G=C3=A9rard?= <thevoodoos@gmail.com>
Date: Wed, 8 Jan 2025 17:46:16 +0100
Subject: [PATCH] Immediately grow the final index instead of building and
 merging

---
 src/ocaml-index/lib/index.ml | 60 +++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 28 deletions(-)

diff --git a/src/ocaml-index/lib/index.ml b/src/ocaml-index/lib/index.ml
index 554880cc4..c3618994b 100644
--- a/src/ocaml-index/lib/index.ml
+++ b/src/ocaml-index/lib/index.ml
@@ -24,6 +24,13 @@ let merge m m' =
     (fun _uid locs locs' -> Some (Lid_set.union locs locs'))
     m m'
 
+let add_one uid lid map =
+  Shape.Uid.Map.update uid
+    (function
+      | None -> Some (Lid_set.singleton lid)
+      | Some set -> Some (Lid_set.add lid set))
+    map
+
 (** Cmt files contains a table of declarations' Uids associated to a typedtree
     fragment. [add_locs_from_fragments] gather locations from these *)
 let gather_locs_from_fragments ~root ~rewrite_root map fragments =
@@ -36,7 +43,7 @@ let gather_locs_from_fragments ~root ~rewrite_root map fragments =
     | Some lid ->
       let lid = to_located_lid lid in
       let lid = if rewrite_root then add_root ~root lid else lid in
-      Shape.Uid.Map.add uid (Lid_set.singleton lid) acc
+      add_one uid lid acc
   in
   Shape.Uid.Tbl.fold add_loc fragments map
 
@@ -72,8 +79,8 @@ let init_load_path_once ~do_not_use_cmt_loadpath =
       Load_path.(init ~auto_include:no_auto_include ~visible ~hidden);
       loaded := true)
 
-let index_of_cmt ~root ~rewrite_root ~build_path ~do_not_use_cmt_loadpath
-    cmt_infos =
+let index_of_cmt ~into ~root ~rewrite_root ~build_path ~do_not_use_cmt_loadpath
+    ~store_shapes cmt_infos =
   let { Cmt_format.cmt_loadpath;
         cmt_impl_shape;
         cmt_modname;
@@ -89,8 +96,7 @@ let index_of_cmt ~root ~rewrite_root ~build_path ~do_not_use_cmt_loadpath
   init_load_path_once ~do_not_use_cmt_loadpath ~dirs:build_path cmt_loadpath;
   let module Reduce = Shape_reduce.Make (Reduce_conf) in
   let defs =
-    gather_locs_from_fragments ~root ~rewrite_root Shape.Uid.Map.empty
-      cmt_uid_to_decl
+    gather_locs_from_fragments ~root ~rewrite_root into.defs cmt_uid_to_decl
   in
   (* The list [cmt_ident_occurrences] associate each ident usage location in the
      module with its (partially) reduced shape. We finish the reduction and
@@ -105,30 +111,31 @@ let index_of_cmt ~root ~rewrite_root ~build_path ~do_not_use_cmt_loadpath
           | result -> result
         in
         match Locate.uid_of_result ~traverse_aliases:false resolved with
-        | Some uid, false -> (add acc_defs uid (Lid_set.singleton lid), acc_apx)
-        | Some uid, true -> (acc_defs, add acc_apx uid (Lid_set.singleton lid))
+        | Some uid, false -> (add_one uid lid acc_defs, acc_apx)
+        | Some uid, true -> (acc_defs, add_one uid lid acc_apx)
         | None, _ -> acc)
-      (defs, Shape.Uid.Map.empty)
-      cmt_ident_occurrences
+      (defs, into.approximated) cmt_ident_occurrences
   in
-  let cu_shape = Hashtbl.create 1 in
-  Option.iter (Hashtbl.add cu_shape cmt_modname) cmt_impl_shape;
+  let cu_shape = into.cu_shape in
+  if store_shapes then
+    Option.iter (Hashtbl.add cu_shape cmt_modname) cmt_impl_shape;
   let stats =
     match cmt_sourcefile with
-    | None -> Stats.empty
+    | None -> into.stats
     | Some src -> (
       let rooted_src = with_root ?root src in
       try
         let stats = Unix.stat rooted_src in
         let src = if rewrite_root then rooted_src else src in
-        Stats.singleton src
+        Stats.add src
           { mtime = stats.st_mtime;
             size = stats.st_size;
             source_digest = cmt_source_digest
           }
-      with Unix.Unix_error _ -> Stats.empty)
+          into.stats
+      with Unix.Unix_error _ -> into.stats)
   in
-  { defs; approximated; cu_shape; stats; root_directory = None }
+  { defs; approximated; cu_shape; stats; root_directory = into.root_directory }
 
 let merge_index ~store_shapes ~into index =
   let defs = merge index.defs into.defs in
@@ -154,19 +161,16 @@ let from_files ~store_shapes ~output_file ~root ~rewrite_root ~build_path
     @@ fun () ->
     List.fold_left
       (fun into file ->
-        let index =
-          match Cmt_cache.read file with
-          | cmt_item ->
-            index_of_cmt ~root ~rewrite_root ~build_path
-              ~do_not_use_cmt_loadpath cmt_item.cmt_infos
-          | exception _ -> (
-            match read ~file with
-            | Index index -> index
-            | _ ->
-              Log.error "Unknown file type: %s" file;
-              exit 1)
-        in
-        merge_index ~store_shapes index ~into)
+        match Cmt_cache.read file with
+        | cmt_item ->
+          index_of_cmt ~into ~root ~rewrite_root ~build_path ~store_shapes
+            ~do_not_use_cmt_loadpath cmt_item.cmt_infos
+        | exception _ -> (
+          match read ~file with
+          | Index index -> merge_index ~store_shapes index ~into
+          | _ ->
+            Log.error "Unknown file type: %s" file;
+            exit 1))
       initial_index files
   in
   write ~file:output_file final_index