From d45ddfaac5067d7d6c08703653a31187d35914a5 Mon Sep 17 00:00:00 2001
From: Oscar Spencer <oscar@grain-lang.org>
Date: Fri, 26 Apr 2024 15:57:05 -0400
Subject: [PATCH] feat(graindoc): Cache module comments (#2102)

---
 compiler/graindoc/docblock.re | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/compiler/graindoc/docblock.re b/compiler/graindoc/docblock.re
index 0ff1517af0..a38719e5f4 100644
--- a/compiler/graindoc/docblock.re
+++ b/compiler/graindoc/docblock.re
@@ -319,23 +319,27 @@ let lookup_type_expr = (~idx, type_exprs) => {
   Option.bind(type_exprs, te => List.nth_opt(te, idx));
 };
 
+let saved_comments = Hashtbl.create(64);
+
 let get_comments_from_loc = (loc: Grain_parsing.Location.t) => {
   open Compile;
 
-  let comments =
-    switch (
-      compile_file(
-        ~is_root_file=true,
-        ~hook=stop_after_parse,
-        loc.loc_start.pos_fname,
-      )
-    ) {
-    | exception exn => []
-    | {cstate_desc: Parsed(parsed_program)} => parsed_program.comments
-    | _ => failwith("Invalid compilation state")
-    };
+  let file = loc.loc_start.pos_fname;
 
-  Comments.to_ordered(comments);
+  switch (Hashtbl.find_opt(saved_comments, file)) {
+  | Some(comments) => comments
+  | None =>
+    let comments =
+      switch (compile_file(~is_root_file=true, ~hook=stop_after_parse, file)) {
+      | exception exn => []
+      | {cstate_desc: Parsed(parsed_program)} => parsed_program.comments
+      | _ => failwith("Invalid compilation state")
+      };
+
+    let ordered = Comments.to_ordered(comments);
+    Hashtbl.add(saved_comments, file, ordered);
+    ordered;
+  };
 };
 
 let attr_name = attr => {