flambda-backend: Record toplevel attributes in cms files (#2206)

* Record toplevel attributes in cms files

* code review

* We only need the first attribute

* fix build

Author: poechsel

driver/compile_common.ml

@@ -63,7 +63,10 @@ let typecheck_intf info ast =
   Profile.(record_call typing) @@ fun () ->
   let tsg =
     ast
-    |> Typemod.type_interface info.module_name info.env
+    |> Typemod.type_interface
+        ~sourcefile:info.source_file
+        info.module_name
+        info.env
     |> print_if info.ppf_dump Clflags.dump_typedtree Printtyped.interface
   in
   let sg = tsg.Typedtree.sig_type in
@@ -117,7 +120,7 @@ let typecheck_impl i parsetree =
   parsetree
   |> Profile.(record typing)
     (Typemod.type_implementation
-       i.source_file i.output_prefix i.module_name i.env)
+       ~sourcefile:i.source_file i.output_prefix i.module_name i.env)
   |> print_if i.ppf_dump Clflags.dump_typedtree
     Printtyped.implementation_with_coercion
   |> print_if i.ppf_dump Clflags.dump_shape

ocamldoc/odoc_analyse.ml

@@ -88,7 +88,7 @@ let process_implementation_file sourcefile =
     in
     let typedtree =
       Typemod.type_implementation
-        sourcefile prefixname compilation_unit env parsetree
+        ~sourcefile prefixname compilation_unit env parsetree
     in
     (Some (parsetree, typedtree), inputfile)
   with
@@ -122,7 +122,7 @@ let process_interface_file sourcefile =
     Pparse.file ~tool_name inputfile
       (no_docstring Parse.interface) Pparse.Signature
   in
-  let sg = Typemod.type_interface compilation_unit (initial_env()) ast in
+  let sg = Typemod.type_interface ~sourcefile compilation_unit (initial_env()) ast in
   Warnings.check_fatal ();
   (ast, sg, inputfile)
 

typing/typemod.ml

@@ -3307,7 +3307,32 @@ let gen_annot outputprefix sourcefile annots =
   Cmt2annot.gen_annot (Some (outputprefix ^ ".annot"))
     ~sourcefile:(Some sourcefile) ~use_summaries:false annots
 
-let type_implementation sourcefile outputprefix modulename initial_env ast =
+let cms_register_toplevel_attributes ~sourcefile ~uid ~f ast =
+  (* Cms files do not store the typetree. This can be a problem for Merlin has
+    it uses attributes - which is why we manually construct a mapping from uid
+    to attributes while typing.
+    Generally `Pstr_attribute` and `Psig_attribute` are not needed by Merlin,
+    except if it is the first element of the compilation unit structure or
+    signature. *)
+  let attr =
+    match ast with
+    | x :: _ -> f x
+    | [] -> None
+  in
+  match attr with
+  | None -> ()
+  | Some attr ->
+    Env.register_uid uid
+      ~loc:(Location.in_file sourcefile)
+      ~attributes:[ attr ]
+
+let cms_register_toplevel_struct_attributes ~sourcefile ~uid ast =
+  cms_register_toplevel_attributes ~sourcefile ~uid ast
+      ~f:(function
+        | { pstr_desc = Pstr_attribute attr; _ }  -> Some attr
+        | _ -> None)
+
+let type_implementation ~sourcefile outputprefix modulename initial_env ast =
   let error e =
     raise (Error (Location.in_file sourcefile, initial_env, e))
   in
@@ -3322,10 +3347,10 @@ let type_implementation sourcefile outputprefix modulename initial_env ast =
       let (str, sg, names, shape, finalenv) =
         Profile.record_call "infer" (fun () ->
           type_structure initial_env ast) in
-      let shape =
-        Shape.set_uid_if_none shape
-          (Uid.of_compilation_unit_id modulename)
-      in
+      let uid = Uid.of_compilation_unit_id modulename in
+      let shape = Shape.set_uid_if_none shape uid in
+      if !Clflags.binary_annotations_cms then
+        cms_register_toplevel_struct_attributes ~sourcefile ~uid ast;
       let simple_sg = Signature_names.simplify finalenv names sg in
       if !Clflags.print_types then begin
         remove_mode_and_jkind_variables finalenv sg;
@@ -3447,10 +3472,21 @@ let save_signature modname tsg outputprefix source_file initial_env cmi =
   Cms_format.save_cms  (outputprefix ^ ".cmsi") modname
     (Some source_file) None
 
-let type_interface modulename env ast =
+let cms_register_toplevel_signature_attributes ~sourcefile ~uid ast =
+  cms_register_toplevel_attributes ~sourcefile ~uid ast
+    ~f:(function
+        | { psig_desc = Psig_attribute attr; _ } -> Some attr
+        | _ -> None)
+
+
+let type_interface ~sourcefile modulename env ast =
   if !Clflags.as_parameter && Compilation_unit.is_packed modulename then begin
     raise(Error(Location.none, Env.empty, Cannot_pack_parameter))
   end;
+  if !Clflags.binary_annotations_cms then begin
+    let uid = Shape.Uid.of_compilation_unit_id modulename in
+    cms_register_toplevel_signature_attributes ~uid ~sourcefile ast
+  end;
   transl_signature env ast
 
 (* "Packaging" of several compilation units into one unit

typing/typemod.mli

@@ -39,10 +39,14 @@ val type_toplevel_phrase:
   Typedtree.structure * Types.signature * Signature_names.t * Shape.t *
   Env.t
 val type_implementation:
-  string -> string -> Compilation_unit.t -> Env.t ->
+  sourcefile:string -> string -> Compilation_unit.t -> Env.t ->
   Parsetree.structure -> Typedtree.implementation
 val type_interface:
-        Compilation_unit.t -> Env.t -> Parsetree.signature -> Typedtree.signature
+  sourcefile:string
+  -> Compilation_unit.t
+  -> Env.t
+  -> Parsetree.signature
+  -> Typedtree.signature
 val transl_signature:
         Env.t -> Parsetree.signature -> Typedtree.signature
 val check_nongen_signature: