dev: perform simple rate limit on heavy dynamic analysis (#532)

* dev: perform simple rate limit on heavy dynamic analysis

* chore: wording

* dev: remove a todo

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

@@ -13,7 +13,7 @@ use once_cell::sync::OnceCell;
 use reflexo::hash::{hash128, FxDashMap};
 use reflexo::{debug_loc::DataSource, ImmutPath};
 use typst::eval::Eval;
-use typst::foundations::{self, Func};
+use typst::foundations::{self, Func, Styles};
 use typst::syntax::{FileId, LinkedNode, SyntaxNode};
 use typst::{
     diag::{eco_format, FileError, FileResult, PackageError},
@@ -21,17 +21,19 @@ use typst::{
     syntax::{package::PackageSpec, Source, Span, VirtualPath},
     World,
 };
-use typst::{foundations::Value, syntax::ast, text::Font};
+use typst::{foundations::Value, model::Document, syntax::ast, text::Font};
 use typst::{layout::Position, syntax::FileId as TypstFileId};
 
 use super::{
-    analyze_bib, post_type_check, BibInfo, DefUseInfo, DefinitionLink, IdentRef, ImportInfo,
-    PathPreference, SigTy, Signature, SignatureTarget, Ty, TypeScheme,
+    analyze_bib, analyze_expr_, analyze_import_, post_type_check, BibInfo, DefUseInfo,
+    DefinitionLink, IdentRef, ImportInfo, PathPreference, SigTy, Signature, SignatureTarget, Ty,
+    TypeScheme,
 };
 use crate::adt::interner::Interned;
 use crate::analysis::analyze_dyn_signature;
 use crate::path_to_url;
 use crate::syntax::{get_deref_target, resolve_id_by_path, DerefTarget};
+use crate::upstream::{tooltip_, Tooltip};
 use crate::{
     lsp_to_typst,
     syntax::{
@@ -40,58 +42,8 @@ use crate::{
     typst_to_lsp, LspPosition, LspRange, PositionEncoding, TypstRange, VersionedDocument,
 };
 
-/// A cache for module-level analysis results of a module.
-///
-/// You should not holds across requests, because source code may change.
-#[derive(Default)]
-pub struct ModuleAnalysisCache {
-    file: OnceCell<FileResult<Bytes>>,
-    source: OnceCell<FileResult<Source>>,
-    def_use: OnceCell<Option<Arc<DefUseInfo>>>,
-    type_check: OnceCell<Option<Arc<TypeScheme>>>,
-}
-
-impl ModuleAnalysisCache {
-    /// Get the bytes content of a file.
-    pub fn file(&self, ctx: &AnalysisContext, file_id: TypstFileId) -> FileResult<Bytes> {
-        self.file.get_or_init(|| ctx.world().file(file_id)).clone()
-    }
-
-    /// Get the source of a file.
-    pub fn source(&self, ctx: &AnalysisContext, file_id: TypstFileId) -> FileResult<Source> {
-        self.source
-            .get_or_init(|| ctx.world().source(file_id))
-            .clone()
-    }
-
-    /// Try to get the def-use information of a file.
-    pub fn def_use(&self) -> Option<Arc<DefUseInfo>> {
-        self.def_use.get().cloned().flatten()
-    }
-
-    /// Compute the def-use information of a file.
-    pub(crate) fn compute_def_use(
-        &self,
-        f: impl FnOnce() -> Option<Arc<DefUseInfo>>,
-    ) -> Option<Arc<DefUseInfo>> {
-        self.def_use.get_or_init(f).clone()
-    }
-
-    /// Try to get the type check information of a file.
-    pub(crate) fn type_check(&self) -> Option<Arc<TypeScheme>> {
-        self.type_check.get().cloned().flatten()
-    }
-
-    /// Compute the type check information of a file.
-    pub(crate) fn compute_type_check(
-        &self,
-        f: impl FnOnce() -> Option<Arc<TypeScheme>>,
-    ) -> Option<Arc<TypeScheme>> {
-        self.type_check.get_or_init(f).clone()
-    }
-}
-
 /// The analysis data holds globally.
+#[derive(Default)]
 pub struct Analysis {
     /// The position encoding for the workspace.
     pub position_encoding: PositionEncoding,
@@ -99,6 +51,8 @@ pub struct Analysis {
     pub enable_periscope: bool,
     /// The global caches for analysis.
     pub caches: AnalysisGlobalCaches,
+    /// The global caches for analysis.
+    pub workers: AnalysisGlobalWorkers,
 }
 
 impl Analysis {
@@ -152,6 +106,57 @@ pub struct AnalysisCaches {
     module_deps: OnceCell<HashMap<TypstFileId, ModuleDependency>>,
 }
 
+/// A cache for module-level analysis results of a module.
+///
+/// You should not holds across requests, because source code may change.
+#[derive(Default)]
+pub struct ModuleAnalysisCache {
+    file: OnceCell<FileResult<Bytes>>,
+    source: OnceCell<FileResult<Source>>,
+    def_use: OnceCell<Option<Arc<DefUseInfo>>>,
+    type_check: OnceCell<Option<Arc<TypeScheme>>>,
+}
+
+impl ModuleAnalysisCache {
+    /// Get the bytes content of a file.
+    pub fn file(&self, ctx: &AnalysisContext, file_id: TypstFileId) -> FileResult<Bytes> {
+        self.file.get_or_init(|| ctx.world().file(file_id)).clone()
+    }
+
+    /// Get the source of a file.
+    pub fn source(&self, ctx: &AnalysisContext, file_id: TypstFileId) -> FileResult<Source> {
+        self.source
+            .get_or_init(|| ctx.world().source(file_id))
+            .clone()
+    }
+
+    /// Try to get the def-use information of a file.
+    pub fn def_use(&self) -> Option<Arc<DefUseInfo>> {
+        self.def_use.get().cloned().flatten()
+    }
+
+    /// Compute the def-use information of a file.
+    pub(crate) fn compute_def_use(
+        &self,
+        f: impl FnOnce() -> Option<Arc<DefUseInfo>>,
+    ) -> Option<Arc<DefUseInfo>> {
+        self.def_use.get_or_init(f).clone()
+    }
+
+    /// Try to get the type check information of a file.
+    pub(crate) fn type_check(&self) -> Option<Arc<TypeScheme>> {
+        self.type_check.get().cloned().flatten()
+    }
+
+    /// Compute the type check information of a file.
+    pub(crate) fn compute_type_check(
+        &self,
+        f: impl FnOnce() -> Option<Arc<TypeScheme>>,
+    ) -> Option<Arc<TypeScheme>> {
+        self.type_check.get_or_init(f).clone()
+    }
+}
+
 /// The resources for analysis.
 pub trait AnalysisResources {
     /// Get the world surface for Typst compiler.
@@ -179,6 +184,17 @@ pub trait AnalysisResources {
     }
 }
 
+/// Shared workers to limit resource usage
+#[derive(Default)]
+pub struct AnalysisGlobalWorkers {
+    /// A possible long running import dynamic analysis task
+    import: RateLimiter,
+    /// A possible long running expression dynamic analysis task
+    expression: RateLimiter,
+    /// A possible long running tooltip dynamic analysis task
+    tooltip: RateLimiter,
+}
+
 /// The context for analyzers.
 pub struct AnalysisContext<'a> {
     /// The root of the workspace.
@@ -510,7 +526,8 @@ impl<'w> AnalysisContext<'w> {
         let w = self.resources.world();
         let w = w.track();
 
-        import_info(w, source)
+        let token = &self.analysis.workers.import;
+        token.enter(|| import_info(w, source))
     }
 
     /// Get the def-use information of a source file.
@@ -661,6 +678,33 @@ impl<'w> AnalysisContext<'w> {
         post_type_check(self, &ty_chk, k.clone()).or_else(|| ty_chk.type_of_span(k.span()))
     }
 
+    /// Try to load a module from the current source file.
+    pub fn analyze_import(&mut self, source: &LinkedNode) -> Option<Value> {
+        let token = &self.analysis.workers.import;
+        token.enter(|| analyze_import_(self.world(), source))
+    }
+
+    /// Try to determine a set of possible values for an expression.
+    pub fn analyze_expr(&mut self, node: &LinkedNode) -> EcoVec<(Value, Option<Styles>)> {
+        let token = &self.analysis.workers.expression;
+        token.enter(|| analyze_expr_(self.world(), node))
+    }
+
+    /// Describe the item under the cursor.
+    ///
+    /// Passing a `document` (from a previous compilation) is optional, but
+    /// enhances the autocompletions. Label completions, for instance, are
+    /// only generated when the document is available.
+    pub fn tooltip(
+        &mut self,
+        document: Option<&Document>,
+        source: &Source,
+        cursor: usize,
+    ) -> Option<Tooltip> {
+        let token = &self.analysis.workers.tooltip;
+        token.enter(|| tooltip_(self.world(), document, source, cursor))
+    }
+
     fn gc(&self) {
         let lifetime = self.lifetime;
         loop {
@@ -814,3 +858,18 @@ impl SearchCtx<'_, '_> {
         }
     }
 }
+
+/// A rate limiter on some (cpu-heavy) action
+#[derive(Default)]
+pub struct RateLimiter {
+    token: std::sync::Mutex<()>,
+}
+
+impl RateLimiter {
+    /// Executes some (cpu-heavy) action with rate limit
+    #[must_use]
+    pub fn enter<T>(&self, f: impl FnOnce() -> T) -> T {
+        let _c = self.token.lock().unwrap();
+        f()
+    }
+}

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

@@ -7,9 +7,8 @@ use typst::{
     World,
 };
 
-use crate::syntax::resolve_id_by_path;
+use crate::{analysis::analyze_import_, syntax::resolve_id_by_path};
 
-use super::analyze_import;
 pub use super::prelude::*;
 
 /// The import information of a source file.
@@ -95,7 +94,7 @@ impl<'a, 'w> ImportCollector<'a, 'w> {
                             let exp = find_import_expr(self.root.leaf_at(exp.range.end));
                             let val = exp
                                 .as_ref()
-                                .and_then(|exp| analyze_import(self.ctx.deref(), exp));
+                                .and_then(|exp| analyze_import_(self.ctx.deref(), exp));
 
                             match val {
                                 Some(Value::Module(m)) => {

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

@@ -210,7 +210,7 @@ pub fn find_definition(
             let root = LinkedNode::new(def_source.root());
             let def_name = root.leaf_at(def.range.start + 1)?;
             log::info!("def_name for function: {def_name:?}", def_name = def_name);
-            let values = analyze_expr(ctx.world(), &def_name);
+            let values = ctx.analyze_expr(&def_name);
             let func = values.into_iter().find(|v| matches!(v.0, Value::Func(..)));
             log::info!("okay for function: {func:?}");
 
@@ -378,7 +378,7 @@ fn resolve_callee_(
         this: None,
     })
     .or_else(|| {
-        let values = analyze_expr(ctx.world(), callee);
+        let values = ctx.analyze_expr(callee);
 
         if let Some(func) = values.into_iter().find_map(|v| match v.0 {
             Value::Func(f) => Some(f),
@@ -397,7 +397,7 @@ fn resolve_callee_(
             } {
                 let target = access.target();
                 let field = access.field().get();
-                let values = analyze_expr(ctx.world(), &callee.find(target.span())?);
+                let values = ctx.analyze_expr(&callee.find(target.span())?);
                 if let Some((this, func_ptr)) = values.into_iter().find_map(|(this, _styles)| {
                     if let Some(Value::Func(f)) = this.ty().scope().get(field) {
                         return Some((this, f.clone()));

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

@@ -11,7 +11,7 @@ use typst::syntax::{ast, LinkedNode, Span, SyntaxKind};
 use typst::World;
 
 /// Try to determine a set of possible values for an expression.
-pub fn analyze_expr(world: &dyn World, node: &LinkedNode) -> EcoVec<(Value, Option<Styles>)> {
+pub fn analyze_expr_(world: &dyn World, node: &LinkedNode) -> EcoVec<(Value, Option<Styles>)> {
     let Some(expr) = node.cast::<ast::Expr>() else {
         return eco_vec![];
     };
@@ -27,13 +27,13 @@ pub fn analyze_expr(world: &dyn World, node: &LinkedNode) -> EcoVec<(Value, Opti
         _ => {
             if node.kind() == SyntaxKind::Contextual {
                 if let Some(child) = node.children().last() {
-                    return analyze_expr(world, &child);
+                    return analyze_expr_(world, &child);
                 }
             }
 
             if let Some(parent) = node.parent() {
                 if parent.kind() == SyntaxKind::FieldAccess && node.index() > 0 {
-                    return analyze_expr(world, parent);
+                    return analyze_expr_(world, parent);
                 }
             }
 
@@ -48,9 +48,9 @@ pub fn analyze_expr(world: &dyn World, node: &LinkedNode) -> EcoVec<(Value, Opti
 }
 
 /// Try to load a module from the current source file.
-pub fn analyze_import(world: &dyn World, source: &LinkedNode) -> Option<Value> {
+pub fn analyze_import_(world: &dyn World, source: &LinkedNode) -> Option<Value> {
     let source_span = source.span();
-    let (source, _) = analyze_expr(world, source).into_iter().next()?;
+    let (source, _) = analyze_expr_(world, source).into_iter().next()?;
     if source.scope().is_some() {
         return Some(source);
     }
@@ -86,7 +86,8 @@ pub struct DynLabel {
     pub label_desc: Option<EcoString>,
     /// Additional details about the label.
     pub detail: Option<EcoString>,
-    /// The title of the bibliography entry. Not present for non-bibliography labels.
+    /// The title of the bibliography entry. Not present for non-bibliography
+    /// labels.
     pub bib_title: Option<EcoString>,
 }
 

crates/tinymist-query/src/hover.rs

@@ -5,9 +5,7 @@ use crate::{
     jump_from_cursor,
     prelude::*,
     syntax::{find_docs_before, get_deref_target, LexicalKind, LexicalVarKind},
-    upstream::{
-        expr_tooltip, plain_docs_sentence, route_of_value, tooltip, truncated_repr, Tooltip,
-    },
+    upstream::{expr_tooltip, plain_docs_sentence, route_of_value, truncated_repr, Tooltip},
     LspHoverContents, StatefulRequest,
 };
 
@@ -42,12 +40,9 @@ impl StatefulRequest for HoverRequest {
         let cursor = offset + 1;
 
         let contents = def_tooltip(ctx, &source, doc.as_ref(), cursor).or_else(|| {
-            Some(typst_to_lsp::tooltip(&tooltip(
-                ctx.world(),
-                doc_ref,
-                &source,
-                cursor,
-            )?))
+            Some(typst_to_lsp::tooltip(
+                &ctx.tooltip(doc_ref, &source, cursor)?,
+            ))
         })?;
 
         let ast_node = LinkedNode::new(source.root()).leaf_at(cursor)?;

crates/tinymist-query/src/prelude.rs

@@ -30,7 +30,7 @@ pub use typst::syntax::{
 };
 pub use typst::World;
 
-pub use crate::analysis::{analyze_expr, AnalysisContext};
+pub use crate::analysis::AnalysisContext;
 pub use crate::lsp_typst_boundary::{
     lsp_to_typst, path_to_url, typst_to_lsp, LspDiagnostic, LspRange, LspSeverity,
     PositionEncoding, TypstDiagnostic, TypstSeverity, TypstSpan,

crates/tinymist-query/src/tests.rs

@@ -71,11 +71,7 @@ pub fn snapshot_testing(name: &str, f: &impl Fn(&mut AnalysisContext, PathBuf))
                     .collect::<Vec<_>>();
                 let mut w = w.snapshot();
                 let w = WrapWorld(&mut w);
-                let a = Analysis {
-                    position_encoding: PositionEncoding::Utf16,
-                    enable_periscope: false,
-                    caches: Default::default(),
-                };
+                let a = Analysis::default();
                 let mut ctx = AnalysisContext::new(root, &w, &a);
                 ctx.test_completion_files(Vec::new);
                 ctx.test_files(|| paths);

crates/tinymist-query/src/upstream/complete.rs

@@ -14,7 +14,7 @@ use unscanny::Scanner;
 
 use super::{plain_docs_sentence, summarize_font_family};
 use crate::adt::interner::Interned;
-use crate::analysis::{analyze_expr, analyze_import, analyze_labels, DynLabel, Ty};
+use crate::analysis::{analyze_labels, DynLabel, Ty};
 use crate::AnalysisContext;
 
 mod ext;
@@ -366,7 +366,7 @@ fn complete_field_accesses(ctx: &mut CompletionContext) -> bool {
         if prev.is::<ast::Expr>();
         if prev.parent_kind() != Some(SyntaxKind::Markup) ||
            prev.prev_sibling_kind() == Some(SyntaxKind::Hash);
-        if let Some((value, styles)) = analyze_expr(ctx.world(), &prev).into_iter().next();
+        if let Some((value, styles)) = ctx.ctx.analyze_expr(&prev).into_iter().next();
         then {
             ctx.from = ctx.cursor;
             field_access_completions(ctx, &value, &styles);
@@ -381,7 +381,7 @@ fn complete_field_accesses(ctx: &mut CompletionContext) -> bool {
         if prev.kind() == SyntaxKind::Dot;
         if let Some(prev_prev) = prev.prev_sibling();
         if prev_prev.is::<ast::Expr>();
-        if let Some((value, styles)) = analyze_expr(ctx.world(), &prev_prev).into_iter().next();
+        if let Some((value, styles)) = ctx.ctx.analyze_expr(&prev_prev).into_iter().next();
         then {
             ctx.from = ctx.leaf.offset();
             field_access_completions(ctx, &value, &styles);
@@ -542,7 +542,7 @@ fn import_item_completions<'a>(
     existing: ast::ImportItems<'a>,
     source: &LinkedNode,
 ) {
-    let Some(value) = analyze_import(ctx.world(), source) else {
+    let Some(value) = ctx.ctx.analyze_import(source) else {
         return;
     };
     let Some(scope) = value.scope() else { return };

crates/tinymist-query/src/upstream/complete/ext.rs

@@ -12,9 +12,7 @@ use typst::visualize::Color;
 
 use super::{Completion, CompletionContext, CompletionKind};
 use crate::adt::interner::Interned;
-use crate::analysis::{
-    analyze_dyn_signature, analyze_import, resolve_call_target, BuiltinTy, PathPreference, Ty,
-};
+use crate::analysis::{analyze_dyn_signature, resolve_call_target, BuiltinTy, PathPreference, Ty};
 use crate::syntax::{param_index_at_leaf, CheckTarget};
 use crate::upstream::complete::complete_code;
 use crate::upstream::plain_docs_sentence;
@@ -90,7 +88,7 @@ impl<'a, 'w> CompletionContext<'a, 'w> {
                     let anaylyze = node.children().find(|child| child.is::<ast::Expr>());
                     let analyzed = anaylyze
                         .as_ref()
-                        .and_then(|source| analyze_import(self.world(), source));
+                        .and_then(|source| self.ctx.analyze_import(source));
                     if analyzed.is_none() {
                         log::debug!("failed to analyze import: {:?}", anaylyze);
                     }

crates/tinymist-query/src/upstream/tooltip.rs

@@ -11,14 +11,14 @@ use typst::util::{round_2, Numeric};
 use typst::World;
 
 use super::{plain_docs_sentence, summarize_font_family, truncated_repr};
-use crate::analysis::{analyze_expr, analyze_labels, DynLabel};
+use crate::analysis::{analyze_expr_, analyze_labels, DynLabel};
 
 /// Describe the item under the cursor.
 ///
 /// Passing a `document` (from a previous compilation) is optional, but enhances
 /// the autocompletions. Label completions, for instance, are only generated
 /// when the document is available.
-pub fn tooltip(
+pub fn tooltip_(
     world: &dyn World,
     document: Option<&Document>,
     source: &Source,
@@ -57,7 +57,7 @@ pub fn expr_tooltip(world: &dyn World, leaf: &LinkedNode) -> Option<Tooltip> {
         return None;
     }
 
-    let values = analyze_expr(world, ancestor);
+    let values = analyze_expr_(world, ancestor);
 
     if let [(value, _)] = values.as_slice() {
         if let Some(docs) = value.docs() {