Gather shapes

src/index-format/index_format.ml

@@ -81,7 +81,7 @@ let write ~file index =
       output_string oc magic_number;
       output_value oc (index : index))
 
-type file_content = Cmt of Cmt_format.cmt_infos | Index of index | Unknown
+type file_content = Cmt of Cmt_format.cmt_infos | Cms of Cms_format.cms_infos | Index of index | Unknown
 
 let read ~file =
   let ic = open_in_bin file in

src/index-format/index_format.mli

@@ -30,6 +30,7 @@ val add : Lid_set.t Uid_map.t -> Shape.Uid.t -> Lid_set.t -> Lid_set.t Uid_map.t
 
 type file_content =
   | Cmt of Cmt_format.cmt_infos
+  | Cms of Cms_format.cms_infos
   | Index of index
   | Unknown
 

src/ocaml-index/bin/ocaml_index.ml

@@ -15,12 +15,13 @@ let rewrite_root = ref false
 let store_shapes = ref false
 let do_not_use_cmt_loadpath = ref false
 
-type command = Aggregate | Dump | Stats
+type command = Aggregate | Dump | Stats | Gather_shapes
 
 let parse_command = function
   | "aggregate" -> Some Aggregate
   | "dump" -> Some Dump
   | "stats" -> Some Stats
+  | "gather-shapes" -> Some Gather_shapes
   | _ -> None
 
 let command = ref None
@@ -79,6 +80,8 @@ let () =
           Index_format.(
             read_exn ~file |> pp Format.std_formatter))
         !input_files
+  | Some Gather_shapes ->
+      Index.gather_shapes ~output_file:!output_file !input_files
   | Some Stats ->
       List.iter
         (fun file ->

src/ocaml-index/lib/cache.ml

@@ -0,0 +1,5 @@
+include File_cache.Make (struct
+    type t = Index_format.file_content
+    let read file = Index_format.read ~file
+    let cache_name = "Index_cache"
+  end)

src/ocaml-index/lib/index.ml

@@ -149,6 +149,24 @@ let index_of_artifact
   in
   { defs; approximated; cu_shape; stats; root_directory = None }
 
+let shape_of_artifact ~impl_shape ~modname =
+  let cu_shape = Hashtbl.create 1 in
+  let modname = Compilation_unit.name_as_string modname in
+  Option.iter (Hashtbl.add cu_shape modname) impl_shape;
+  {
+    defs = Shape.Uid.Map.empty;
+    approximated = Shape.Uid.Map.empty;
+    cu_shape;
+    stats = Stats.empty;
+    root_directory = None;
+  }
+
+let shape_of_cmt { Cmt_format.cmt_impl_shape; cmt_modname; _ }  =
+  shape_of_artifact ~impl_shape:cmt_impl_shape ~modname:cmt_modname
+
+let shape_of_cms { Cms_format.cms_impl_shape; cms_modname; _ }  =
+  shape_of_artifact ~impl_shape:cms_impl_shape ~modname:cms_modname
+
 let index_of_cmt ~root ~build_path cmt_infos =
   let {
     Cmt_format.cmt_loadpath;
@@ -259,3 +277,32 @@ let from_files ~store_shapes ~output_file ~root ~rewrite_root ~build_path
       initial_index files
   in
   write ~file:output_file final_index
+
+let gather_shapes ~output_file files =
+  let initial_index =
+    {
+      defs = Shape.Uid.Map.empty;
+      approximated = Shape.Uid.Map.empty;
+      cu_shape = Hashtbl.create 64;
+      stats = Stats.empty;
+      root_directory = None;
+    }
+  in
+  let final_index =
+    List.fold_left
+      (fun into file ->
+         let index =
+           match Cache.read file with
+           | Cmt cmt_infos -> Some (shape_of_cmt cmt_infos)
+           | Cms cmt_infos -> Some (shape_of_cms cmt_infos)
+           | Index index -> Some index
+           | Unknown | exception _ ->
+             Log.error "Not a valid file %S" file;
+             None
+         in
+         match index with
+         | None -> into
+         | Some index -> merge_index ~store_shapes:true index ~into)
+      initial_index files
+  in
+  write ~file:output_file final_index