feat: add command to profile the entire server (#1438)

* feat: declare and bind tinymist.profileServer command

* feat: editor bridge with the frontend

* feat: start and stop server profiling

* feat: add profile-server prototype (#1440)

* Add profile-server prototype

* fix: use branch

---------

Co-authored-by: Myriad-Dreamin <camiyoru@gmail.com>

* feat: make it good

* build: update cargo.lock

* dev: ls profile impl and hook

* test: update snapshot

---------

Co-authored-by: Derived Cat <hooyuser@outlook.com>

Author: Myriad-Dreamin

Cargo.lock

@@ -137,6 +137,7 @@ reflexo-vec2svg = { version = "=0.6.0-rc1" }
 typst = "0.13.1"
 typst-html = "0.13.1"
 typst-library = "0.13.1"
+typst-macros = "0.13.1"
 typst-timing = "0.13.1"
 typst-svg = "0.13.1"
 typst-render = "0.13.1"
@@ -266,6 +267,7 @@ extend-exclude = ["/.git", "fixtures"]
 #
 # A regular build MUST use `tag` or `rev` to specify the version of the patched crate to ensure stability.
 typst = { git = "https://github.com/Myriad-Dreamin/typst.git", tag = "tinymist/v0.13.10" }
+typst-macros = { git = "https://github.com/Myriad-Dreamin/typst.git", tag = "tinymist/v0.13.10" }
 typst-library = { git = "https://github.com/Myriad-Dreamin/typst.git", tag = "tinymist/v0.13.10" }
 typst-html = { git = "https://github.com/Myriad-Dreamin/typst.git", tag = "tinymist/v0.13.10" }
 typst-timing = { git = "https://github.com/Myriad-Dreamin/typst.git", tag = "tinymist/v0.13.10" }

Cargo.toml

@@ -199,10 +199,18 @@ impl LspClientRoot {
             msg_kind: M::get_message_kind(),
             sender: Arc::downgrade(&_strong),
             req_queue: Arc::new(Mutex::new(ReqQueue::default())),
+
+            hook: Arc::new(()),
         };
         Self { weak, _strong }
     }
 
+    /// Sets the hook for the language server host.
+    pub fn with_hook(mut self, hook: Arc<dyn LsHook>) -> Self {
+        self.weak.hook = hook;
+        self
+    }
+
     /// Returns the weak reference to the language server host.
     pub fn weak(&self) -> LspClient {
         self.weak.clone()
@@ -221,6 +229,8 @@ pub struct LspClient {
     pub(crate) msg_kind: MessageKind,
     pub(crate) sender: Weak<ConnectionTx>,
     pub(crate) req_queue: Arc<Mutex<ReqQueue>>,
+
+    pub(crate) hook: Arc<dyn LsHook>,
 }
 
 impl LspClient {
@@ -290,7 +300,7 @@ impl LspClient {
     /// Registers an client2server request in the request queue.
     pub fn register_request(&self, method: &str, id: &RequestId, received_at: Instant) {
         let mut req_queue = self.req_queue.lock();
-        self.start_request(id, method);
+        self.hook.start_request(id, method);
         req_queue
             .incoming
             .register(id.clone(), (method.to_owned(), received_at));
@@ -327,7 +337,7 @@ impl LspClient {
             return;
         };
 
-        self.stop_request(&id, &method, received_at);
+        self.hook.stop_request(&id, &method, received_at);
 
         let Some(sender) = self.sender.upgrade() else {
             log::warn!("failed to send response ({method}, {id}): connection closed");
@@ -380,7 +390,19 @@ impl LspClient {
     }
 }
 
-impl LspClient {
+/// A trait that defines the hook for the language server.
+pub trait LsHook: fmt::Debug + Send + Sync {
+    /// Starts a request.
+    fn start_request(&self, req_id: &RequestId, method: &str);
+    /// Stops a request.
+    fn stop_request(&self, req_id: &RequestId, method: &str, received_at: Instant);
+    /// Starts a notification.
+    fn start_notification(&self, method: &str);
+    /// Stops a notification.
+    fn stop_notification(&self, method: &str, received_at: Instant, result: LspResult<()>);
+}
+
+impl LsHook for () {
     fn start_request(&self, req_id: &RequestId, method: &str) {
         log::info!("handling {method} - ({req_id})");
     }

crates/sync-lsp/src/server.rs

@@ -260,7 +260,7 @@ where
 
     /// Handles an incoming event.
     fn on_event(&mut self, received_at: Instant, not: dap::Event) -> anyhow::Result<()> {
-        self.client.start_notification(&not.event);
+        self.client.hook.start_notification(&not.event);
         let handle = |s,
                       dap::Event {
                           seq: _,
@@ -273,7 +273,9 @@ where
             };
 
             let result = handler(s, body);
-            self.client.stop_notification(&event, received_at, result);
+            self.client
+                .hook
+                .stop_notification(&event, received_at, result);
 
             Ok(())
         };

crates/sync-lsp/src/server/dap_srv.rs

@@ -354,15 +354,17 @@ where
 
     /// Handles an incoming notification.
     fn on_notification(&mut self, received_at: Instant, not: Notification) -> anyhow::Result<()> {
-        self.client.start_notification(&not.method);
+        self.client.hook.start_notification(&not.method);
         let handle = |s, Notification { method, params }: Notification| {
             let Some(handler) = self.notifications.get(method.as_str()) else {
                 log::warn!("unhandled notification: {method}");
                 return Ok(());
             };
 
             let result = handler(s, params);
-            self.client.stop_notification(&method, received_at, result);
+            self.client
+                .hook
+                .stop_notification(&method, received_at, result);
 
             Ok(())
         };

crates/sync-lsp/src/server/lsp_srv.rs

@@ -35,7 +35,9 @@ tinymist-world.workspace = true
 toml.workspace = true
 triomphe.workspace = true
 typst.workspace = true
+typst-macros.workspace = true
 typst-shim.workspace = true
+typst-timing.workspace = true
 unscanny.workspace = true
 
 [dev-dependencies]

crates/tinymist-analysis/Cargo.toml

@@ -23,6 +23,7 @@ pub fn analyze_expr(world: &dyn World, node: &LinkedNode) -> EcoVec<(Value, Opti
 }
 
 /// Try to determine a set of possible values for an expression.
+#[typst_macros::time(span = node.span())]
 pub fn analyze_expr_(world: &dyn World, node: &SyntaxNode) -> EcoVec<(Value, Option<Styles>)> {
     let Some(expr) = node.cast::<ast::Expr>() else {
         return eco_vec![];
@@ -51,6 +52,7 @@ pub fn analyze_expr_(world: &dyn World, node: &SyntaxNode) -> EcoVec<(Value, Opt
 }
 
 /// Try to load a module from the current source file.
+#[typst_macros::time(span = source.span())]
 pub fn analyze_import_(world: &dyn World, source: &SyntaxNode) -> (Option<Value>, Option<Value>) {
     let source_span = source.span();
     let Some((source, _)) = analyze_expr_(world, source).into_iter().next() else {
@@ -109,6 +111,7 @@ pub struct DynLabel {
 /// - All labels and descriptions for them, if available
 /// - A split offset: All labels before this offset belong to nodes, all after
 ///   belong to a bibliography.
+#[typst_macros::time]
 pub fn analyze_labels(document: &TypstDocument) -> (Vec<DynLabel>, usize) {
     let mut output = vec![];
 

crates/tinymist-analysis/src/track_values.rs

@@ -51,7 +51,9 @@ tinymist-std.workspace = true
 tinymist-l10n.workspace = true
 tinymist-lint.workspace = true
 typst.workspace = true
+typst-macros.workspace = true
 typst-shim.workspace = true
+typst-timing.workspace = true
 unscanny.workspace = true
 walkdir.workspace = true
 yaml-rust2.workspace = true

crates/tinymist-query/Cargo.toml

@@ -38,6 +38,7 @@ pub struct CallInfo {
 
 // todo: cache call
 /// Analyzes a function call.
+#[typst_macros::time(span = node.span())]
 pub fn analyze_call(
     ctx: &mut LocalContext,
     source: Source,

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

@@ -50,6 +50,7 @@ impl CompletionPair<'_, '_, '_> {
         self.def_completions(defines, parens);
     }
 
+    #[typst_macros::time]
     pub fn scope_defs(&mut self) -> Option<Defines> {
         let mut defines = Defines {
             types: self.worker.ctx.type_check(&self.cursor.source),

crates/tinymist-query/src/analysis/completion/scope.rs

@@ -81,6 +81,7 @@ impl HasNameRange for Decl {
 
 // todo: field definition
 /// Finds the definition of a symbol.
+#[typst_macros::time(span = syntax.node().span())]
 pub fn definition(
     ctx: &Arc<SharedContext>,
     source: &Source,

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

@@ -779,6 +779,7 @@ impl SharedContext {
     }
 
     /// Get the lint result of a source file.
+    #[typst_macros::time(span = source.root().span())]
     pub(crate) fn lint(self: &Arc<Self>, source: &Source) -> LintInfo {
         let ei = self.expr_stage(source);
         let ti = self.type_check(source);
@@ -1296,6 +1297,7 @@ fn ceil_char_boundary(text: &str, mut cursor: usize) -> usize {
     cursor.min(text.len())
 }
 
+#[typst_macros::time]
 #[comemo::memoize]
 fn analyze_bib(
     world: Tracked<dyn World + '_>,

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

@@ -8,6 +8,7 @@ use tinymist_world::package::PackageSpec;
 use super::prelude::*;
 
 /// Get link expressions from a source.
+#[typst_macros::time(span = src.root().span())]
 #[comemo::memoize]
 pub fn get_link_exprs(src: &Source) -> Arc<LinkInfo> {
     let root = LinkedNode::new(src.root());

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

@@ -13,6 +13,7 @@ use crate::ty::BuiltinTy;
 
 /// With given type information, check the type of a literal expression again by
 /// touching the possible related nodes.
+#[typst_macros::time(span = node.span())]
 pub(crate) fn post_type_check(
     ctx: Arc<SharedContext>,
     ti: &TypeInfo,

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

@@ -26,6 +26,7 @@ use crate::{
 pub type SemanticTokens = Arc<Vec<SemanticToken>>;
 
 /// Get the semantic tokens for a source.
+#[typst_macros::time(span = source.root().span())]
 pub(crate) fn get_semantic_tokens(ctx: &mut LocalContext, source: &Source) -> SemanticTokens {
     let mut tokenizer = Tokenizer::new(
         source.clone(),

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

@@ -30,6 +30,18 @@ pub enum SignatureTarget {
     Convert(Func),
 }
 
+impl SignatureTarget {
+    /// Returns the span of the callee node.
+    pub fn span(&self) -> Span {
+        match self {
+            SignatureTarget::Def(_, def) => def.decl.span(),
+            SignatureTarget::SyntaxFast(_, span) | SignatureTarget::Syntax(_, span) => *span,
+            SignatureTarget::Runtime(func) | SignatureTarget::Convert(func) => func.span(),
+        }
+    }
+}
+
+#[typst_macros::time(span = callee_node.span())]
 pub(crate) fn analyze_signature(
     ctx: &Arc<SharedContext>,
     callee_node: SignatureTarget,
@@ -41,6 +53,7 @@ pub(crate) fn analyze_signature(
     })
 }
 
+#[typst_macros::time(span = callee_node.span())]
 fn analyze_type_signature(
     ctx: &Arc<SharedContext>,
     callee_node: &SignatureTarget,
@@ -315,6 +328,7 @@ impl BoundChecker for AliasStackChecker<'_, '_> {
     }
 }
 
+#[typst_macros::time(span = callee_node.span())]
 fn analyze_dyn_signature(
     ctx: &Arc<SharedContext>,
     callee_node: &SignatureTarget,

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

@@ -29,6 +29,7 @@ pub struct TypeEnv {
 }
 
 /// Type checking at the source unit level.
+#[typst_macros::time(span = ei.source.root().span())]
 pub(crate) fn type_check(
     ctx: Arc<SharedContext>,
     ei: ExprInfo,
@@ -228,14 +229,14 @@ impl TypeChecker<'_> {
         &mut self,
         sig: &Interned<SigTy>,
         args: &Interned<SigTy>,
-        withs: Option<&Vec<Interned<SigTy>>>,
+        with: Option<&Vec<Interned<SigTy>>>,
     ) {
-        let call_desc = (sig.clone(), args.clone(), withs.cloned());
+        let call_desc = (sig.clone(), args.clone(), with.cloned());
         if !self.call_cache.insert(call_desc) {
             return;
         }
 
-        for (arg_recv, arg_ins) in sig.matches(args, withs) {
+        for (arg_recv, arg_ins) in sig.matches(args, with) {
             self.constrain(arg_ins, arg_recv);
         }
     }

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

@@ -12,6 +12,7 @@ static EMPTY_DOCSTRING: LazyLock<DocString> = LazyLock::new(DocString::default);
 static EMPTY_VAR_DOC: LazyLock<VarDoc> = LazyLock::new(VarDoc::default);
 
 impl TypeChecker<'_> {
+    #[typst_macros::time(span = expr.span())]
     pub(crate) fn check_syntax(&mut self, expr: &Expr) -> Option<Ty> {
         Some(match expr {
             Expr::Block(exprs) => self.check_block(exprs),

crates/tinymist-query/src/analysis/tyck/syntax.rs

@@ -28,6 +28,7 @@ use super::{compute_docstring, def::*, DocCommentMatcher, InterpretMode};
 
 pub type ExprRoute = FxHashMap<TypstFileId, Option<Arc<LazyHash<LexicalScope>>>>;
 
+#[typst_macros::time(span = source.root().span())]
 pub(crate) fn expr_of(
     ctx: Arc<SharedContext>,
     source: Source,

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

@@ -12,6 +12,7 @@ pub struct IndexInfo {
     pub(crate) identifiers: FxHashSet<Interned<str>>,
 }
 
+#[typst_macros::time(span = src.root().span())]
 #[comemo::memoize]
 pub fn get_index_info(src: &Source) -> Arc<IndexInfo> {
     let root = src.root();

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

@@ -11,6 +11,7 @@ use typst_shim::utils::LazyHash;
 
 use super::{is_mark, CommentGroupMatcher};
 
+#[typst_macros::time(span = source.root().span())]
 pub(crate) fn get_lexical_hierarchy(
     source: &Source,
     scope_kind: LexicalScopeKind,

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

@@ -18,6 +18,7 @@ pub struct ModuleDependency {
 /// It will scan all the files in the context, using
 /// [`LocalContext::source_files`], and find the dependencies and dependents
 /// of each file.
+#[typst_macros::time]
 pub fn construct_module_dependencies(
     ctx: &mut LocalContext,
 ) -> HashMap<TypstFileId, ModuleDependency> {

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

@@ -318,6 +318,22 @@ impl<T, E: std::fmt::Display> IgnoreLogging<T> for Result<T, E> {
     }
 }
 
+impl<T> IgnoreLogging<T> for Option<T> {
+    fn log_error(self, msg: &str) -> Option<T> {
+        self.or_else(|| {
+            log::error!("{msg}");
+            None
+        })
+    }
+
+    fn log_error_with(self, f: impl FnOnce() -> String) -> Option<T> {
+        self.or_else(|| {
+            log::error!("{}", f());
+            None
+        })
+    }
+}
+
 /// A trait to add context to a result.
 pub trait WithContext<T>: Sized {
     /// Add a context to the result.

crates/tinymist-std/src/error.rs

@@ -15,8 +15,7 @@ use rkyv::{Archive, Deserialize as rDeser, Serialize as rSer};
 
 use crate::error::prelude::Result;
 
-pub(crate) type FxBuildHasher = std::hash::BuildHasherDefault<FxHasher>;
-pub use rustc_hash::{FxHashMap, FxHashSet, FxHasher};
+pub use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet, FxHasher};
 // pub type FxIndexSet<K> = indexmap::IndexSet<K, FxHasher>;
 // pub type FxIndexMap<K, V> = indexmap::IndexMap<K, V, FxHasher>;
 /// A dashmap that uses the FxHasher as the underlying hasher.
@@ -34,7 +33,8 @@ pub type FxDashMap<K, V> = dashmap::DashMap<K, V, FxBuildHasher>;
 /// > is a probability of 1 in 36,890,000,000,000 of a `StableCrateId`
 /// > collision.
 ///
-/// This stores the 16-bytes data in a pair of `u64` instead of single `u128` to avoid heavily affecting `align` of the embedding structs.
+/// This stores the 16-bytes data in a pair of `u64` instead of single `u128` to
+/// avoid heavily affecting `align` of the embedding structs.
 #[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
 #[cfg_attr(feature = "rkyv", derive(Archive, rDeser, rSer))]
 #[cfg_attr(feature = "rkyv-validation", archive(check_bytes))]