perf: add cache for docstrings (#1877)

crates/tinymist-query/src/analysis/global.rs

@@ -8,6 +8,7 @@ use comemo::{Track, Tracked};
 use lsp_types::Url;
 use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
+use tinymist_analysis::docs::DocString;
 use tinymist_analysis::stats::AllocStats;
 use tinymist_analysis::ty::term_value;
 use tinymist_analysis::{analyze_expr_, analyze_import_};
@@ -163,6 +164,7 @@ impl Analysis {
     /// Clear all cached resources.
     pub fn clear_cache(&self) {
         self.caches.signatures.clear();
+        self.caches.docstrings.clear();
         self.caches.def_signatures.clear();
         self.caches.static_signatures.clear();
         self.caches.terms.clear();
@@ -289,6 +291,7 @@ impl LocalContextGuard {
         caches.static_signatures.retain(|(l, _)| retainer(*l));
         caches.terms.retain(|(l, _)| retainer(*l));
         caches.signatures.retain(|(l, _)| retainer(*l));
+        caches.docstrings.retain(|(l, _)| retainer(*l));
     }
 }
 
@@ -969,6 +972,23 @@ impl SharedContext {
         res.get_or_init(|| compute(self)).clone()
     }
 
+    pub(crate) fn compute_docstring(
+        self: &Arc<Self>,
+        fid: TypstFileId,
+        docs: String,
+        kind: DefKind,
+    ) -> Option<Arc<DocString>> {
+        let res = self
+            .analysis
+            .caches
+            .docstrings
+            .entry(hash128(&(fid, &docs, kind)), self.lifetime);
+        res.get_or_init(|| {
+            crate::syntax::docs::do_compute_docstring(self, fid, docs, kind).map(Arc::new)
+        })
+        .clone()
+    }
+
     /// Remove html tags from markup content if necessary.
     pub fn remove_html(&self, markup: EcoString) -> EcoString {
         if !self.analysis.remove_html {
@@ -1163,6 +1183,7 @@ pub struct AnalysisGlobalCaches {
     def_signatures: CacheMap<DeferredCompute<Option<Signature>>>,
     static_signatures: CacheMap<DeferredCompute<Option<Signature>>>,
     signatures: CacheMap<DeferredCompute<Option<Signature>>>,
+    docstrings: CacheMap<DeferredCompute<Option<Arc<DocString>>>>,
     terms: CacheMap<(Value, Ty)>,
 }
 

crates/tinymist-query/src/syntax/docs.rs

@@ -14,7 +14,7 @@ use crate::{
 
 use super::DeclExpr;
 
-pub(crate) fn compute_docstring(
+pub(crate) fn do_compute_docstring(
     ctx: &Arc<SharedContext>,
     fid: TypstFileId,
     docs: String,

crates/tinymist-query/src/syntax/expr.rs

@@ -24,7 +24,7 @@ use crate::{
     ty::{BuiltinTy, InsTy, Ty},
 };
 
-use super::{compute_docstring, def::*, DocCommentMatcher, InterpretMode};
+use super::{def::*, DocCommentMatcher, InterpretMode};
 
 pub type ExprRoute = FxHashMap<TypstFileId, Option<Arc<LazyHash<LexicalScope>>>>;
 
@@ -84,11 +84,9 @@ pub(crate) fn expr_of(
     let imports_base = Arc::new(Mutex::new(FxHashMap::default()));
     let imports = imports_base.clone();
 
-    let module_docstring = Arc::new(
-        find_module_level_docs(&source)
-            .and_then(|docs| compute_docstring(&ctx, source.id(), docs, DefKind::Module))
-            .unwrap_or_default(),
-    );
+    let module_docstring = find_module_level_docs(&source)
+        .and_then(|docs| ctx.compute_docstring(source.id(), docs, DefKind::Module))
+        .unwrap_or_default();
 
     let (exports, root) = {
         let mut w = ExprWorker {
@@ -220,11 +218,9 @@ impl ExprWorker<'_> {
 
     fn check_docstring(&mut self, decl: &DeclExpr, docs: Option<String>, kind: DefKind) {
         if let Some(docs) = docs {
-            let docstring = compute_docstring(&self.ctx, self.fid, docs, kind);
+            let docstring = self.ctx.compute_docstring(self.fid, docs, kind);
             if let Some(docstring) = docstring {
-                self.docstrings
-                    .lock()
-                    .insert(decl.clone(), Arc::new(docstring));
+                self.docstrings.lock().insert(decl.clone(), docstring);
             }
         }
     }

crates/tinymist-query/src/syntax/mod.rs

@@ -10,7 +10,6 @@ pub(crate) mod index;
 pub(crate) mod lexical_hierarchy;
 pub(crate) mod module;
 
-pub(crate) use docs::*;
 pub use expr::*;
 pub use index::*;
 pub use lexical_hierarchy::*;