fix(graindoc): Improve location lookup so re-exports do not crash it (#1280)

chore(compiler): Add export ident to export_declarations

Co-authored-by: Philip Blair <philip@pblair.org>

Author: phated

compiler/graindoc/docblock.re

@@ -32,6 +32,69 @@ let () =
     }
   });
 
+let enumerate_exports = stmts => {
+  let id_tbl = ref(Ident.empty);
+
+  let rec pattern_ids = ({pat_desc, pat_loc}: Typedtree.pattern) => {
+    switch (pat_desc) {
+    | TPatVar(id, _) => [(id, pat_loc)]
+    | TPatAlias(subpat, id, _) => [(id, pat_loc), ...pattern_ids(subpat)]
+    | TPatTuple(pats)
+    | TPatArray(pats)
+    | TPatConstruct(_, _, pats) => List.concat(List.map(pattern_ids, pats))
+    | TPatRecord(elts, _) =>
+      List.concat(List.map(((_, _, pat)) => pattern_ids(pat), elts))
+    | _ => []
+    };
+  };
+
+  module ExportIterator =
+    TypedtreeIter.MakeIterator({
+      include TypedtreeIter.DefaultIteratorArgument;
+
+      let enter_toplevel_stmt =
+          ({ttop_desc, ttop_attributes}: Typedtree.toplevel_stmt) => {
+        switch (ttop_desc) {
+        | TTopData(decls) =>
+          List.iter(
+            ({data_id, data_loc}: Typedtree.data_declaration) => {
+              id_tbl := Ident.add(data_id, data_loc, id_tbl^)
+            },
+            decls,
+          )
+        | TTopExport(decls) =>
+          List.iter(
+            ({tex_id, tex_loc}: Typedtree.export_declaration) => {
+              id_tbl := Ident.add(tex_id, tex_loc, id_tbl^)
+            },
+            decls,
+          )
+        | TTopForeign({tvd_id, tvd_loc}) =>
+          id_tbl := Ident.add(tvd_id, tvd_loc, id_tbl^)
+        | TTopLet(_, _, vbinds) =>
+          List.iter(
+            ({vb_pat}: Typedtree.value_binding) => {
+              List.iter(
+                ((id, loc)) => {id_tbl := Ident.add(id, loc, id_tbl^)},
+                pattern_ids(vb_pat),
+              )
+            },
+            vbinds,
+          )
+        | _ => ()
+        };
+      };
+    });
+
+  List.iter(ExportIterator.iter_toplevel_stmt, stmts);
+
+  id_tbl^;
+};
+
+let location_for_ident = (~exports, ident) => {
+  snd(Ident.find_name(Ident.name(ident), exports));
+};
+
 let module_name_of_location = (loc: Grain_parsing.Location.t) => {
   Grain_utils.Filepath.String.filename_to_module_name(
     loc.loc_start.pos_fname,
@@ -83,12 +146,12 @@ let lookup_type_expr = (~idx, type_exprs) => {
 };
 
 let for_value_description =
-    (~comments, ~ident: Ident.t, vd: Types.value_description) => {
-  let module_name = module_name_of_location(vd.val_loc);
+    (~comments, ~ident: Ident.t, ~loc, vd: Types.value_description) => {
+  let module_name = module_name_of_location(loc);
   let name = title_for_api(~module_name, ident);
   let type_sig = Printtyp.string_of_value_description(~ident, vd);
   let comment =
-    Comments.Doc.ending_on(~lnum=vd.val_loc.loc_start.pos_lnum - 1, comments);
+    Comments.Doc.ending_on(~lnum=loc.loc_start.pos_lnum - 1, comments);
 
   let (description, attributes) =
     switch (comment) {
@@ -118,15 +181,12 @@ let for_value_description =
 };
 
 let for_type_declaration =
-    (~comments, ~ident: Ident.t, td: Types.type_declaration) => {
-  let module_name = module_name_of_location(td.type_loc);
+    (~comments, ~ident: Ident.t, ~loc, td: Types.type_declaration) => {
+  let module_name = module_name_of_location(loc);
   let name = title_for_api(~module_name, ident);
   let type_sig = Printtyp.string_of_type_declaration(~ident, td);
   let comment =
-    Comments.Doc.ending_on(
-      ~lnum=td.type_loc.loc_start.pos_lnum - 1,
-      comments,
-    );
+    Comments.Doc.ending_on(~lnum=loc.loc_start.pos_lnum - 1, comments);
 
   let (description, attributes) =
     switch (comment) {
@@ -138,35 +198,37 @@ let for_type_declaration =
 };
 
 let for_signature_item =
-    (~env: Env.t, ~comments, sig_item: Types.signature_item) => {
+    (
+      ~comments,
+      ~exports: Ident.tbl(Grain_parsing.Location.t),
+      sig_item: Types.signature_item,
+    ) => {
   switch (sig_item) {
-  | TSigValue(ident, ovd) =>
-    // Fetch original location as signatures don't contain real locations
-    let vd = Env.find_value(ovd.val_fullpath, env);
-    let val_loc = vd.val_loc;
-    let docblock =
-      for_value_description(~comments, ~ident, {...ovd, val_loc});
+  | TSigValue(ident, vd) =>
+    let loc = location_for_ident(~exports, ident);
+    let docblock = for_value_description(~comments, ~ident, ~loc, vd);
     Some(docblock);
-  | TSigType(ident, otd, _rec) =>
-    // Fetch original location as signatures don't contain real locations
-    let td = Env.find_type(otd.type_path, env);
-    let type_loc = td.type_loc;
-    let docblock =
-      for_type_declaration(~comments, ~ident, {...otd, type_loc});
+  | TSigType(ident, td, _rec) =>
+    let loc = location_for_ident(~exports, ident);
+    let docblock = for_type_declaration(~comments, ~ident, ~loc, td);
     Some(docblock);
   | _ => None
   };
 };
 
 let signature_item_in_range =
-    (~env: Env.t, sig_item: Types.signature_item, range: Grain_utils.Range.t) => {
+    (
+      ~exports: Ident.tbl(Grain_parsing.Location.t),
+      sig_item: Types.signature_item,
+      range: Grain_utils.Range.t,
+    ) => {
   switch (sig_item) {
   | TSigValue(ident, vd) =>
-    let vd = Env.find_value(vd.val_fullpath, env);
-    Grain_utils.Range.inRange(vd.val_loc.loc_start.pos_lnum, range);
+    let loc = location_for_ident(~exports, ident);
+    Grain_utils.Range.inRange(loc.loc_start.pos_lnum, range);
   | TSigType(ident, td, _rec) =>
-    let td = Env.find_type(td.type_path, env);
-    Grain_utils.Range.inRange(td.type_loc.loc_start.pos_lnum, range);
+    let loc = location_for_ident(~exports, ident);
+    Grain_utils.Range.inRange(loc.loc_start.pos_lnum, range);
   | _ => false
   };
 };

compiler/graindoc/graindoc.re

@@ -78,7 +78,8 @@ let generate_docs =
     (~current_version, ~output=?, program: Typedtree.typed_program) => {
   let comments = Comments.to_ordered(program.comments);
 
-  let env = program.env;
+  let exports = Docblock.enumerate_exports(program.statements);
+
   let signature_items = program.signature.cmi_sign;
 
   let buf = Buffer.create(0);
@@ -156,7 +157,7 @@ let generate_docs =
   };
 
   let add_docblock = sig_item => {
-    let docblock = Docblock.for_signature_item(~env, ~comments, sig_item);
+    let docblock = Docblock.for_signature_item(~comments, ~exports, sig_item);
     switch (docblock) {
     | Some(docblock) =>
       Buffer.add_buffer(
@@ -197,7 +198,7 @@ let generate_docs =
         );
         List.iter(
           sig_item =>
-            if (Docblock.signature_item_in_range(~env, sig_item, range)) {
+            if (Docblock.signature_item_in_range(~exports, sig_item, range)) {
               add_docblock(sig_item);
             },
           signature_items,

compiler/src/typed/typedtree.re

@@ -509,6 +509,7 @@ type import_declaration = {
 
 [@deriving sexp]
 type export_declaration = {
+  tex_id: Ident.t,
   tex_path: Path.t,
   [@sexp_drop_if sexp_locs_disabled]
   tex_loc: Location.t,

compiler/src/typed/typedtree.rei

@@ -472,6 +472,7 @@ type import_declaration = {
 
 [@deriving sexp]
 type export_declaration = {
+  tex_id: Ident.t,
   tex_path: Path.t,
   [@sexp_drop_if sexp_locs_disabled]
   tex_loc: Location.t,

compiler/src/typed/typemod.re

@@ -602,7 +602,10 @@ let type_module = (~toplevel=false, funct_body, anchor, env, sstr /*scope*/) =>
             };
           let name = Identifier.IdentName(name);
           let (p, {val_fullpath} as desc) = Env.lookup_value(name, env);
-          (TSigValue(id, desc), {tex_path: val_fullpath, tex_loc: loc});
+          (
+            TSigValue(id, desc),
+            {tex_id: id, tex_path: val_fullpath, tex_loc: loc},
+          );
         },
         exports,
       );