fix(graindoc): Avoid singletons when building ordered comments (#1208)

Author: phated

compiler/graindoc/docblock.re

@@ -105,12 +105,13 @@ let string_of_type_expr = out_type => {
   Option.map(to_string, out_type);
 };
 
-let for_value_description = (~ident: Ident.t, vd: Types.value_description) => {
+let for_value_description =
+    (~comments, ~ident: Ident.t, vd: Types.value_description) => {
   let module_name = module_name_of_location(vd.val_loc);
   let name = title_for_api(~module_name, ident);
   let type_sig = string_of_value_description(~ident, vd);
   let comment =
-    Comments.Doc.ending_on_lnum(vd.val_loc.loc_start.pos_lnum - 1);
+    Comments.Doc.ending_on(~lnum=vd.val_loc.loc_start.pos_lnum - 1, comments);
 
   let (description, attributes) =
     switch (comment) {
@@ -137,12 +138,16 @@ let for_value_description = (~ident: Ident.t, vd: Types.value_description) => {
   {module_name, name, type_sig, description, attributes};
 };
 
-let for_type_declaration = (~ident: Ident.t, td: Types.type_declaration) => {
+let for_type_declaration =
+    (~comments, ~ident: Ident.t, td: Types.type_declaration) => {
   let module_name = module_name_of_location(td.type_loc);
   let name = title_for_api(~module_name, ident);
   let type_sig = string_of_type_declaration(~ident, td);
   let comment =
-    Comments.Doc.ending_on_lnum(td.type_loc.loc_start.pos_lnum - 1);
+    Comments.Doc.ending_on(
+      ~lnum=td.type_loc.loc_start.pos_lnum - 1,
+      comments,
+    );
 
   let (description, attributes) =
     switch (comment) {
@@ -153,15 +158,16 @@ let for_type_declaration = (~ident: Ident.t, td: Types.type_declaration) => {
   {module_name, name, type_sig, description, attributes};
 };
 
-let for_signature_item = (~env: Env.t, sig_item: Types.signature_item) => {
+let for_signature_item =
+    (~env: Env.t, ~comments, sig_item: Types.signature_item) => {
   switch (sig_item) {
   | TSigValue(ident, vd) =>
     let vd = Env.find_value(vd.val_fullpath, env);
-    let docblock = for_value_description(~ident, vd);
+    let docblock = for_value_description(~comments, ~ident, vd);
     Some(docblock);
   | TSigType(ident, td, _rec) =>
     let td = Env.find_type(td.type_path, env);
-    let docblock = for_type_declaration(~ident, td);
+    let docblock = for_type_declaration(~comments, ~ident, td);
     Some(docblock);
   | _ => None
   };

compiler/graindoc/graindoc.re

@@ -101,13 +101,13 @@ let compile_typed = (opts: params) => {
 
 let generate_docs =
     ({current_version, output}: params, program: Typedtree.typed_program) => {
-  Comments.setup_comments(program.comments);
+  let comments = Comments.to_ordered(program.comments);
 
   let env = program.env;
   let signature_items = program.signature.cmi_sign;
 
   let buf = Buffer.create(0);
-  let module_comment = Comments.Doc.find_module();
+  let module_comment = Comments.Doc.find_module(comments);
   switch (module_comment) {
   | Some((_, desc, attrs)) =>
     // TODO: Should we fail if more than one `@module` attribute?
@@ -181,7 +181,7 @@ let generate_docs =
   };
 
   let add_docblock = sig_item => {
-    let docblock = Docblock.for_signature_item(~env, sig_item);
+    let docblock = Docblock.for_signature_item(~env, ~comments, sig_item);
     switch (docblock) {
     | Some(docblock) =>
       Buffer.add_buffer(
@@ -192,7 +192,7 @@ let generate_docs =
     };
   };
 
-  let section_comments = Comments.Doc.find_sections();
+  let section_comments = Comments.Doc.find_sections(comments);
   if (List.length(section_comments) == 0) {
     List.iter(add_docblock, signature_items);
   } else {

compiler/src/diagnostics/comments.re

@@ -280,18 +280,35 @@ module Attribute = {
 type description = option(string);
 type attributes = list(Attribute.t);
 
-module IntMap = Map.Make(Int);
+module type OrderedComments = {
+  type comment = (Typedtree.comment, description, attributes);
 
-type comments = {
-  mutable by_start_lnum:
-    IntMap.t((Typedtree.comment, description, attributes)),
-  mutable by_end_lnum:
-    IntMap.t((Typedtree.comment, description, attributes)),
+  type comments;
+
+  let comments: comments;
+
+  let find_starting_on_lnum: int => option(comment);
+  let find_ending_on_lnum: int => option(comment);
+
+  let iter: ((int, comment) => unit) => unit;
 };
 
-let comments = {by_start_lnum: IntMap.empty, by_end_lnum: IntMap.empty};
+module MakeOrderedComments =
+       (Raw: {let comments: list(Typedtree.comment);})
+       : OrderedComments => {
+  module IntMap = Map.Make(Int);
+
+  type comment = (Typedtree.comment, description, attributes);
+
+  type comments = {
+    mutable by_start_lnum:
+      IntMap.t((Typedtree.comment, description, attributes)),
+    mutable by_end_lnum:
+      IntMap.t((Typedtree.comment, description, attributes)),
+  };
+
+  let comments = {by_start_lnum: IntMap.empty, by_end_lnum: IntMap.empty};
 
-let setup_comments = (raw_comments: list(Typedtree.comment)) => {
   List.iter(
     (comment: Typedtree.comment) => {
       let (start_lnum, end_lnum, data) =
@@ -310,10 +327,23 @@ let setup_comments = (raw_comments: list(Typedtree.comment)) => {
         IntMap.add(start_lnum, data, comments.by_start_lnum);
       comments.by_end_lnum = IntMap.add(end_lnum, data, comments.by_end_lnum);
     },
-    raw_comments,
+    Raw.comments,
   );
+
+  let find_starting_on_lnum = lnum =>
+    IntMap.find_opt(lnum, comments.by_start_lnum);
+  let find_ending_on_lnum = lnum =>
+    IntMap.find_opt(lnum, comments.by_end_lnum);
+
+  let iter = fn => IntMap.iter(fn, comments.by_start_lnum);
 };
 
+let to_ordered = (comments): (module OrderedComments) =>
+  (module
+   MakeOrderedComments({
+     let comments = comments;
+   }));
+
 let start_line = (comment: Typedtree.comment) => {
   switch (comment) {
   | Line({cmt_loc})
@@ -333,17 +363,17 @@ let end_line = (comment: Typedtree.comment) => {
 };
 
 module Doc = {
-  let starting_on_lnum = lnum => {
-    let data = IntMap.find_opt(lnum, comments.by_start_lnum);
+  let starting_on = (~lnum, module C: OrderedComments) => {
+    let data = C.find_starting_on_lnum(lnum);
     switch (data) {
     | Some((Doc({cmt_content}), _, _)) => data
     | _ => None
     };
   };
 
-  let ending_on_lnum = lnum => {
+  let ending_on = (~lnum, module C: OrderedComments) => {
     let rec ending_on_lnum_help = (lnum, check_prev) => {
-      let data = IntMap.find_opt(lnum, comments.by_end_lnum);
+      let data = C.find_ending_on_lnum(lnum);
       switch (data) {
       | Some((Doc({cmt_content}), _, _)) => data
       // Hack to handle code that has an attribute on the line before, such as `@disableGC`
@@ -354,14 +384,12 @@ module Doc = {
     ending_on_lnum_help(lnum, true);
   };
 
-  let find_module = () => {
+  let find_module = (module C: OrderedComments) => {
     let module_comments = ref([]);
-    IntMap.iter(
-      (_, (_comment, _desc, attrs) as comment) =>
-        if (List.exists(Attribute.is_module, attrs)) {
-          module_comments := [comment, ...module_comments^];
-        },
-      comments.by_start_lnum,
+    C.iter((_, (_comment, _desc, attrs) as comment) =>
+      if (List.exists(Attribute.is_module, attrs)) {
+        module_comments := [comment, ...module_comments^];
+      }
     );
     if (List.length(module_comments^) > 1) {
       failwith("More than one @module block is not supported");
@@ -370,14 +398,12 @@ module Doc = {
     };
   };
 
-  let find_sections = () => {
+  let find_sections = (module C: OrderedComments) => {
     let section_comments = ref([]);
-    IntMap.iter(
-      (_, (_comment, _desc, attrs) as comment) =>
-        if (List.exists(Attribute.is_section, attrs)) {
-          section_comments := [comment, ...section_comments^];
-        },
-      comments.by_start_lnum,
+    C.iter((_, (_comment, _desc, attrs) as comment) =>
+      if (List.exists(Attribute.is_section, attrs)) {
+        section_comments := [comment, ...section_comments^];
+      }
     );
     List.rev(section_comments^);
   };