feat: support code context queries (#217)

* feat: support code context queries

* dev: improve switch

* dev: update snapshot

crates/tinymist-query/src/code_context.rs

@@ -0,0 +1,121 @@
+use serde::{Deserialize, Serialize};
+
+use crate::{prelude::*, SyntaxRequest};
+
+/// A mode in which a text document is interpreted.
+#[derive(Debug, Clone, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub enum InterpretMode {
+    /// The position is in a comment.
+    Comment,
+    /// The position is in a string.
+    String,
+    /// The position is in a raw.
+    Raw,
+    /// The position is in a markup block.
+    Markup,
+    /// The position is in a code block.
+    Code,
+    /// The position is in a math equation.
+    Math,
+}
+
+/// A query to get the mode at a specific position in a text document.
+#[derive(Debug, Clone, Deserialize)]
+#[serde(tag = "kind", rename_all = "camelCase")]
+pub enum InteractCodeContextQuery {
+    /// Get the mode at a specific position in a text document.
+    ModeAt {
+        /// The position inside the text document.
+        position: LspPosition,
+    },
+}
+
+/// A response to a `InteractCodeContextQuery`.
+#[derive(Debug, Clone, Serialize)]
+#[serde(tag = "kind", rename_all = "camelCase")]
+pub enum InteractCodeContextResponse {
+    /// The mode at the requested position.
+    ModeAt {
+        /// The mode at the requested position.
+        mode: InterpretMode,
+    },
+}
+
+/// A request to get the mode at a specific position in a text document.
+#[derive(Debug, Clone, Deserialize)]
+#[serde(tag = "kind")]
+pub struct InteractCodeContextRequest {
+    /// The path to the text document.
+    pub path: PathBuf,
+    /// The queries to execute.
+    pub query: Vec<InteractCodeContextQuery>,
+}
+
+impl SyntaxRequest for InteractCodeContextRequest {
+    type Response = Vec<InteractCodeContextResponse>;
+
+    fn request(
+        self,
+        source: &Source,
+        positing_encoding: PositionEncoding,
+    ) -> Option<Self::Response> {
+        let mut responses = Vec::new();
+
+        for query in self.query {
+            match query {
+                InteractCodeContextQuery::ModeAt { position } => {
+                    let pos = lsp_to_typst::position(position, positing_encoding, source)?;
+
+                    // get mode
+                    let root = LinkedNode::new(source.root());
+                    let leaf = root.leaf_at(pos);
+                    let mut leaf = leaf.as_ref();
+                    let mode = loop {
+                        log::info!("leaf for context: {:?}", leaf);
+                        use SyntaxKind::*;
+                        if let Some(t) = leaf {
+                            match t.kind() {
+                                LineComment | BlockComment => break InterpretMode::Comment,
+                                Raw => break InterpretMode::Raw,
+                                Str => break InterpretMode::String,
+                                CodeBlock | Code => break InterpretMode::Code,
+                                ContentBlock | Markup => break InterpretMode::Markup,
+                                Equation | Math => break InterpretMode::Math,
+                                Space | Linebreak | Parbreak | Escape | Shorthand | SmartQuote
+                                | RawLang | RawDelim | RawTrimmed | Hash | LeftBrace
+                                | RightBrace | LeftBracket | RightBracket | LeftParen
+                                | RightParen | Comma | Semicolon | Colon | Star | Underscore
+                                | Dollar | Plus | Minus | Slash | Hat | Prime | Dot | Eq | EqEq
+                                | ExclEq | Lt | LtEq | Gt | GtEq | PlusEq | HyphEq | StarEq
+                                | SlashEq | Dots | Arrow | Root | Not | And | Or | None | Auto
+                                | As | Named | Keyed | Error | Eof => {}
+                                Text | Strong | Emph | Link | Label | Ref | RefMarker | Heading
+                                | HeadingMarker | ListItem | ListMarker | EnumItem | EnumMarker
+                                | TermItem | TermMarker => break InterpretMode::Markup,
+                                MathIdent | MathAlignPoint | MathDelimited | MathAttach
+                                | MathPrimes | MathFrac | MathRoot => break InterpretMode::Math,
+                                Let | Set | Show | Context | If | Else | For | In | While
+                                | Break | Continue | Return | Import | Include | Ident | Bool
+                                | Int | Float | Numeric | FieldAccess | Args | Spread | Closure
+                                | Params | LetBinding | SetRule | ShowRule | Contextual
+                                | Conditional | WhileLoop | ForLoop | ModuleImport
+                                | ImportItems | RenamedImportItem | ModuleInclude | LoopBreak
+                                | LoopContinue | FuncReturn | FuncCall | Unary | Binary
+                                | Parenthesized | Dict | Array | Destructuring
+                                | DestructAssignment => break InterpretMode::Code,
+                            }
+                            leaf = t.parent();
+                        } else {
+                            break InterpretMode::Markup;
+                        }
+                    };
+
+                    responses.push(InteractCodeContextResponse::ModeAt { mode });
+                }
+            }
+        }
+
+        Some(responses)
+    }
+}

crates/tinymist-query/src/lib.rs

@@ -20,6 +20,8 @@ pub use analysis::AnalysisContext;
 use typst::{model::Document as TypstDocument, syntax::Source};
 
 pub use diagnostics::*;
+pub(crate) mod code_context;
+pub use code_context::*;
 pub(crate) mod code_lens;
 pub use code_lens::*;
 pub(crate) mod completion;
@@ -215,6 +217,7 @@ mod polymorphic {
         Formatting(FormattingRequest),
         FoldingRange(FoldingRangeRequest),
         SelectionRange(SelectionRangeRequest),
+        InteractCodeContext(InteractCodeContextRequest),
 
         DocumentMetrics(DocumentMetricsRequest),
         ServerInfo(ServerInfoRequest),
@@ -245,6 +248,7 @@ mod polymorphic {
                 CompilerQueryRequest::Formatting(..) => ContextFreeUnique,
                 CompilerQueryRequest::FoldingRange(..) => ContextFreeUnique,
                 CompilerQueryRequest::SelectionRange(..) => ContextFreeUnique,
+                CompilerQueryRequest::InteractCodeContext(..) => PinnedFirst,
 
                 CompilerQueryRequest::DocumentMetrics(..) => PinnedFirst,
                 CompilerQueryRequest::ServerInfo(..) => Mergeable,
@@ -274,6 +278,7 @@ mod polymorphic {
                 CompilerQueryRequest::Formatting(req) => &req.path,
                 CompilerQueryRequest::FoldingRange(req) => &req.path,
                 CompilerQueryRequest::SelectionRange(req) => &req.path,
+                CompilerQueryRequest::InteractCodeContext(req) => &req.path,
 
                 CompilerQueryRequest::DocumentMetrics(req) => &req.path,
                 CompilerQueryRequest::ServerInfo(..) => return None,
@@ -304,6 +309,7 @@ mod polymorphic {
         Formatting(Option<Vec<TextEdit>>),
         FoldingRange(Option<Vec<FoldingRange>>),
         SelectionRange(Option<Vec<SelectionRange>>),
+        InteractCodeContext(Option<Vec<InteractCodeContextResponse>>),
 
         DocumentMetrics(Option<DocumentMetricsResponse>),
         ServerInfo(Option<HashMap<String, ServerInfoResponse>>),

crates/tinymist/src/server/lsp.rs

@@ -650,6 +650,7 @@ impl TypstLanguageServer {
             exec_fn!("tinymist.focusMain", Self::focus_document),
             exec_fn!("tinymist.doInitTemplate", Self::init_template),
             exec_fn!("tinymist.doGetTemplateEntry", Self::do_get_template_entry),
+            exec_fn!("tinymist.interactCodeContext", Self::interact_code_context),
             exec_fn_!("tinymist.getDocumentTrace", Self::get_document_trace),
             exec_fn!("tinymist.getDocumentMetrics", Self::get_document_metrics),
             exec_fn!("tinymist.getServerInfo", Self::get_server_info),
@@ -686,6 +687,31 @@ impl TypstLanguageServer {
         Ok(res)
     }
 
+    /// Interact with the code context at the source file.
+    pub fn interact_code_context(&mut self, _arguments: Vec<JsonValue>) -> LspResult<JsonValue> {
+        let queries = _arguments.into_iter().next().ok_or_else(|| {
+            invalid_params("The first parameter is not a valid code context query array")
+        })?;
+
+        #[derive(Debug, Clone, Deserialize)]
+        #[serde(rename_all = "camelCase")]
+        pub struct InteractCodeContextParams {
+            pub text_document: TextDocumentIdentifier,
+            pub query: Vec<tinymist_query::InteractCodeContextQuery>,
+        }
+
+        let params: InteractCodeContextParams = serde_json::from_value(queries)
+            .map_err(|e| invalid_params(format!("Cannot parse code context queries: {e}")))?;
+        let path = as_path(params.text_document);
+        let query = params.query;
+
+        let res = run_query!(self.InteractCodeContext(path, query))?;
+        let res =
+            serde_json::to_value(res).map_err(|_| internal_error("Cannot serialize responses"))?;
+
+        Ok(res)
+    }
+
     /// Get the trace data of the document.
     pub fn get_document_trace(
         &mut self,

crates/tinymist/src/state.rs

@@ -222,6 +222,7 @@ impl TypstLanguageServer {
         use CompilerQueryRequest::*;
 
         match query {
+            InteractCodeContext(req) => query_source!(self, InteractCodeContext, req),
             SemanticTokensFull(req) => query_tokens_cache!(self, SemanticTokensFull, req),
             SemanticTokensDelta(req) => query_tokens_cache!(self, SemanticTokensDelta, req),
             FoldingRange(req) => query_source!(self, FoldingRange, req),
@@ -275,7 +276,8 @@ impl TypstLanguageServer {
                 Ok(CompilerQueryResponse::ServerInfo(Some(res)))
             }
 
-            FoldingRange(..)
+            InteractCodeContext(..)
+            | FoldingRange(..)
             | SelectionRange(..)
             | SemanticTokensDelta(..)
             | Formatting(..)

tests/e2e/main.rs

@@ -374,7 +374,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("neovim"));
-        insta::assert_snapshot!(hash, @"siphash128_13:9480f1a9fe4bb5166bf916d610af1e37");
+        insta::assert_snapshot!(hash, @"siphash128_13:35e217e61e97a024b17ba8020374a65f");
     }
 
     {
@@ -385,7 +385,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("vscode"));
-        insta::assert_snapshot!(hash, @"siphash128_13:1cc4396062d1b450002eb8d89b668705");
+        insta::assert_snapshot!(hash, @"siphash128_13:e40d55d49e7012058d1623c8ea1a1573");
     }
 }