feat: highlights all break points for that loop context (#317)

* feat: highlights all break points for that loop or block context

* docs: update readme

* dev: remove a debug logging

* QAQ

README.md

@@ -16,6 +16,11 @@ Language service (LSP) features:
   - Also known as "syntax highlighting".
 - [Diagnostics](https://code.visualstudio.com/api/language-extensions/programmatic-language-features#provide-diagnostics)
   - Also known as "error checking" or "error reporting".
+- [Document highlight](https://code.visualstudio.com/api/language-extensions/programmatic-language-features#highlight-all-occurrences-of-a-symbol-in-a-document)
+  - Highlight all break points in a loop context.
+  - (Todo) Highlight all exit points in a function context.
+  - (Todo) Highlight all captures in a closure context.
+  - (Todo) Highlight all occurrences of a symbol in a document.
 - [Document symbols](https://code.visualstudio.com/docs/getstarted/userinterface#_outline-view)
   - Also known as "document outline" or "table of contents" **in Typst**.
 - [Folding ranges](https://burkeholland.gitbook.io/vs-code-can-do-that/exercise-3-navigation-and-refactoring/folding-sections)

crates/tinymist-query/src/document_highlight.rs

@@ -0,0 +1,170 @@
+use crate::{prelude::*, SemanticRequest};
+
+/// The [`textDocument/documentHighlight`] request
+///
+/// [`textDocument/documentHighlight`]: https://microsoft.github.io/language-server-protocol/specification#textDocument_documentHighlight
+#[derive(Debug, Clone)]
+pub struct DocumentHighlightRequest {
+    /// The path of the document to request highlight for.
+    pub path: PathBuf,
+    /// The position of the document to request highlight for.
+    pub position: LspPosition,
+}
+
+impl SemanticRequest for DocumentHighlightRequest {
+    type Response = Vec<DocumentHighlight>;
+
+    fn request(self, ctx: &mut AnalysisContext) -> Option<Self::Response> {
+        let source = ctx.source_by_path(&self.path).ok()?;
+        let cursor = ctx.to_typst_pos(self.position, &source)?;
+
+        let root = LinkedNode::new(source.root());
+        let mut node = &root.leaf_at(cursor)?;
+
+        loop {
+            match node.kind() {
+                SyntaxKind::For
+                | SyntaxKind::While
+                | SyntaxKind::Break
+                | SyntaxKind::Continue
+                | SyntaxKind::LoopBreak
+                | SyntaxKind::LoopContinue => {
+                    return DocumentHighlightWorker::new(ctx, &source).highlight_loop_of(node)
+                }
+                SyntaxKind::Arrow
+                | SyntaxKind::Params
+                | SyntaxKind::Return
+                | SyntaxKind::FuncReturn => return highlight_func_returns(ctx, node),
+                _ => {}
+            }
+            node = node.parent()?;
+        }
+    }
+}
+
+struct DocumentHighlightWorker<'a, 'w> {
+    ctx: &'a mut AnalysisContext<'w>,
+    current: &'a Source,
+    highlights: Vec<DocumentHighlight>,
+    worklist: Vec<LinkedNode<'a>>,
+}
+
+impl<'a, 'w> DocumentHighlightWorker<'a, 'w> {
+    fn new(ctx: &'a mut AnalysisContext<'w>, current: &'a Source) -> Self {
+        Self {
+            ctx,
+            current,
+            highlights: Vec::new(),
+            worklist: Vec::new(),
+        }
+    }
+
+    fn finish(self) -> Option<Vec<DocumentHighlight>> {
+        (!self.highlights.is_empty()).then_some(self.highlights)
+    }
+
+    fn annotate(&mut self, node: &LinkedNode) {
+        let mut rng = node.range();
+
+        // if previous node is hash
+        if rng.start > 0 && self.current.text().as_bytes()[rng.start - 1] == b'#' {
+            rng.start -= 1;
+        }
+
+        self.highlights.push(DocumentHighlight {
+            range: self.ctx.to_lsp_range(rng, self.current),
+            kind: None,
+        });
+    }
+
+    fn check<F>(&mut self, check: F)
+    where
+        F: Fn(&mut Self, LinkedNode<'a>),
+    {
+        while let Some(node) = self.worklist.pop() {
+            check(self, node);
+        }
+    }
+
+    fn check_children(&mut self, node: &LinkedNode<'a>) {
+        if node.get().children().len() == 0 {
+            return;
+        }
+
+        for child in node.children() {
+            self.worklist.push(child.clone());
+        }
+    }
+
+    fn check_loop(&mut self, node: LinkedNode<'a>) {
+        match node.kind() {
+            SyntaxKind::ForLoop
+            | SyntaxKind::WhileLoop
+            | SyntaxKind::Closure
+            | SyntaxKind::Contextual => {
+                return;
+            }
+            SyntaxKind::LoopBreak | SyntaxKind::LoopContinue => {
+                self.annotate(&node);
+                return;
+            }
+            _ => {}
+        }
+
+        self.check_children(&node);
+    }
+
+    fn highlight_loop_of(mut self, node: &'a LinkedNode<'a>) -> Option<Vec<DocumentHighlight>> {
+        let _ = self.ctx;
+
+        // find the nearest loop node
+        let loop_node = ancestors(node)
+            .find(|node| matches!(node.kind(), SyntaxKind::ForLoop | SyntaxKind::WhileLoop))?;
+
+        // find the first key word of the loop node
+        let keyword = loop_node.children().find(|node| node.kind().is_keyword());
+        if let Some(keyword) = keyword {
+            self.annotate(&keyword);
+        }
+
+        self.check_children(loop_node);
+        self.check(Self::check_loop);
+
+        log::debug!("highlights: {:?}", self.highlights);
+        self.finish()
+    }
+}
+
+fn highlight_func_returns(
+    ctx: &mut AnalysisContext,
+    node: &LinkedNode,
+) -> Option<Vec<DocumentHighlight>> {
+    let _ = ctx;
+    let _ = node;
+    None
+}
+
+fn ancestors<'a>(node: &'a LinkedNode<'a>) -> impl Iterator<Item = &'a LinkedNode<'a>> {
+    std::iter::successors(Some(node), |node| node.parent())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::*;
+
+    #[test]
+    fn test() {
+        snapshot_testing("document_highlight", &|ctx, path| {
+            let source = ctx.source_by_path(&path).unwrap();
+
+            let request = DocumentHighlightRequest {
+                path: path.clone(),
+                position: find_test_position_after(&source),
+            };
+
+            let result = request.request(ctx);
+            assert_snapshot!(JsonRepr::new_redacted(result, &REDACT_LOC));
+        });
+    }
+}

crates/tinymist-query/src/fixtures/document_highlight/base.typ

@@ -0,0 +1,3 @@
+#for i in range(0) {
+  (/* position after */break)
+}

crates/tinymist-query/src/fixtures/document_highlight/nest.typ

@@ -0,0 +1,6 @@
+#for i in range(0) {
+  (/* position after */break)
+  for j in range(0) {
+    break
+  }
+}

crates/tinymist-query/src/fixtures/document_highlight/nest2.typ

@@ -0,0 +1,6 @@
+#for i in range(0) {
+  (/* position after */break)
+  context {
+    break
+  }
+}

crates/tinymist-query/src/fixtures/document_highlight/snaps/test@base.typ.snap

@@ -0,0 +1,13 @@
+---
+source: crates/tinymist-query/src/document_highlight.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/document_highlight/base.typ
+---
+[
+ {
+  "range": "0:0:0:4"
+ },
+ {
+  "range": "1:23:1:28"
+ }
+]

crates/tinymist-query/src/fixtures/document_highlight/snaps/test@nest.typ.snap

@@ -0,0 +1,13 @@
+---
+source: crates/tinymist-query/src/document_highlight.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/document_highlight/nest.typ
+---
+[
+ {
+  "range": "0:0:0:4"
+ },
+ {
+  "range": "1:23:1:28"
+ }
+]

crates/tinymist-query/src/fixtures/document_highlight/snaps/test@nest2.typ.snap

@@ -0,0 +1,13 @@
+---
+source: crates/tinymist-query/src/document_highlight.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/document_highlight/nest2.typ
+---
+[
+ {
+  "range": "0:0:0:4"
+ },
+ {
+  "range": "1:23:1:28"
+ }
+]

crates/tinymist-query/src/fixtures/document_highlight/snaps/test@syntax_error.typ.snap

@@ -0,0 +1,13 @@
+---
+source: crates/tinymist-query/src/document_highlight.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/document_highlight/syntax_error.typ
+---
+[
+ {
+  "range": "0:0:0:4"
+ },
+ {
+  "range": "1:2:1:7"
+ }
+]

crates/tinymist-query/src/fixtures/document_highlight/syntax_error.typ

@@ -0,0 +1,6 @@
+#for i in range(0) {
+  break
+  context {
+    (/* position after */break)
+  }
+}

crates/tinymist-query/src/lib.rs

@@ -33,6 +33,8 @@ pub(crate) mod color_presentation;
 pub use color_presentation::*;
 pub(crate) mod document_color;
 pub use document_color::*;
+pub(crate) mod document_highlight;
+pub use document_highlight::*;
 pub(crate) mod document_symbol;
 pub use document_symbol::*;
 pub(crate) mod document_metrics;
@@ -203,6 +205,7 @@ mod polymorphic {
         References(ReferencesRequest),
         InlayHint(InlayHintRequest),
         DocumentColor(DocumentColorRequest),
+        DocumentHighlight(DocumentHighlightRequest),
         ColorPresentation(ColorPresentationRequest),
         CodeAction(CodeActionRequest),
         CodeLens(CodeLensRequest),
@@ -235,6 +238,7 @@ mod polymorphic {
                 CompilerQueryRequest::References(..) => PinnedFirst,
                 CompilerQueryRequest::InlayHint(..) => Unique,
                 CompilerQueryRequest::DocumentColor(..) => PinnedFirst,
+                CompilerQueryRequest::DocumentHighlight(..) => PinnedFirst,
                 CompilerQueryRequest::ColorPresentation(..) => ContextFreeUnique,
                 CompilerQueryRequest::CodeAction(..) => Unique,
                 CompilerQueryRequest::CodeLens(..) => Unique,
@@ -266,6 +270,7 @@ mod polymorphic {
                 CompilerQueryRequest::References(req) => &req.path,
                 CompilerQueryRequest::InlayHint(req) => &req.path,
                 CompilerQueryRequest::DocumentColor(req) => &req.path,
+                CompilerQueryRequest::DocumentHighlight(req) => &req.path,
                 CompilerQueryRequest::ColorPresentation(req) => &req.path,
                 CompilerQueryRequest::CodeAction(req) => &req.path,
                 CompilerQueryRequest::CodeLens(req) => &req.path,
@@ -298,6 +303,7 @@ mod polymorphic {
         References(Option<Vec<LspLocation>>),
         InlayHint(Option<Vec<InlayHint>>),
         DocumentColor(Option<Vec<ColorInformation>>),
+        DocumentHighlight(Option<Vec<DocumentHighlight>>),
         ColorPresentation(Option<Vec<ColorPresentation>>),
         CodeAction(Option<Vec<CodeActionOrCommand>>),
         CodeLens(Option<Vec<CodeLens>>),

crates/tinymist-query/src/prelude.rs

@@ -11,11 +11,12 @@ pub use log::{error, trace};
 pub use lsp_types::{
     request::GotoDeclarationResponse, CodeAction, CodeActionKind, CodeActionOrCommand, CodeLens,
     ColorInformation, ColorPresentation, CompletionResponse, DiagnosticRelatedInformation,
-    DocumentSymbol, DocumentSymbolResponse, Documentation, FoldingRange, GotoDefinitionResponse,
+    DocumentHighlight, DocumentSymbol, DocumentSymbolResponse, Documentation, FoldingRange,
-    Hover, InlayHint, LanguageString, Location as LspLocation, LocationLink, MarkedString,
+    GotoDefinitionResponse, Hover, InlayHint, LanguageString, Location as LspLocation,
-    MarkupContent, MarkupKind, Position as LspPosition, PrepareRenameResponse, SelectionRange,
+    LocationLink, MarkedString, MarkupContent, MarkupKind, Position as LspPosition,
-    SemanticTokens, SemanticTokensDelta, SemanticTokensFullDeltaResult, SemanticTokensResult,
+    PrepareRenameResponse, SelectionRange, SemanticTokens, SemanticTokensDelta,
-    SignatureHelp, SignatureInformation, SymbolInformation, TextEdit, Url, WorkspaceEdit,
+    SemanticTokensFullDeltaResult, SemanticTokensResult, SignatureHelp, SignatureInformation,
+    SymbolInformation, TextEdit, Url, WorkspaceEdit,
 };
 pub use reflexo::vector::ir::DefId;
 pub use serde_json::Value as JsonValue;

crates/tinymist-query/src/tests.rs

@@ -189,7 +189,15 @@ pub fn find_test_range(s: &Source) -> Range<usize> {
     start as usize..end as usize
 }
 
+pub fn find_test_position_after(s: &Source) -> LspPosition {
+    find_test_position_(s, 1)
+}
+
 pub fn find_test_position(s: &Source) -> LspPosition {
+    find_test_position_(s, 0)
+}
+
+pub fn find_test_position_(s: &Source, offset: usize) -> LspPosition {
     enum AstMatcher {
         MatchAny { prev: bool },
         MatchIdent { prev: bool },
@@ -271,7 +279,7 @@ pub fn find_test_position(s: &Source) -> LspPosition {
         break;
     }
 
-    typst_to_lsp::offset_to_position(n.offset(), PositionEncoding::Utf16, s)
+    typst_to_lsp::offset_to_position(n.offset() + offset, PositionEncoding::Utf16, s)
 }
 
 // pub static REDACT_URI: Lazy<RedactFields> = Lazy::new(||

crates/tinymist/src/server/lsp.rs

@@ -257,6 +257,7 @@ impl TypstLanguageServer {
             request_fn!(Completion, Self::completion),
             request_fn!(SemanticTokensFullRequest, Self::semantic_tokens_full),
             request_fn!(SemanticTokensFullDeltaRequest, Self::semantic_tokens_full_delta),
+            request_fn!(DocumentHighlightRequest, Self::document_highlight),
             request_fn!(DocumentSymbolRequest, Self::document_symbol),
             // Sync for low latency
             request_fn_!(Formatting, Self::formatting),
@@ -1050,6 +1051,14 @@ impl TypstLanguageServer {
         run_query!(self.SelectionRange(path, positions))
     }
 
+    fn document_highlight(
+        &mut self,
+        params: DocumentHighlightParams,
+    ) -> LspResult<Option<Vec<DocumentHighlight>>> {
+        let (path, position) = as_path_pos(params.text_document_position_params);
+        run_query!(self.DocumentHighlight(path, position))
+    }
+
     fn document_symbol(
         &mut self,
         params: DocumentSymbolParams,

crates/tinymist/src/server/lsp_init.rs

@@ -373,6 +373,7 @@ impl Init {
                     },
                 }),
                 color_provider: Some(ColorProviderCapability::Simple(true)),
+                document_highlight_provider: Some(OneOf::Left(true)),
                 document_symbol_provider: Some(OneOf::Left(true)),
                 workspace_symbol_provider: Some(OneOf::Left(true)),
                 selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),

crates/tinymist/src/state.rs

@@ -291,9 +291,9 @@ impl TypstLanguageServer {
         assert!(query.fold_feature() != FoldRequestFeature::ContextFreeUnique);
 
         match query {
-            OnExport(OnExportRequest { kind, path }) => Ok(
+            OnExport(OnExportRequest { kind, path }) => Ok(CompilerQueryResponse::OnExport(
-                CompilerQueryResponse::OnExport(client.on_export(kind, path)?),
+                client.on_export(kind, path)?,
-            ),
+            )),
             OnSaveExport(OnSaveExportRequest { path }) => {
                 client.on_save_export(path)?;
                 Ok(CompilerQueryResponse::OnSaveExport(()))
@@ -303,6 +303,7 @@ impl TypstLanguageServer {
             GotoDeclaration(req) => query_world!(client, GotoDeclaration, req),
             References(req) => query_world!(client, References, req),
             InlayHint(req) => query_world!(client, InlayHint, req),
+            DocumentHighlight(req) => query_world!(client, DocumentHighlight, req),
             DocumentColor(req) => query_world!(client, DocumentColor, req),
             CodeAction(req) => query_world!(client, CodeAction, req),
             CodeLens(req) => query_world!(client, CodeLens, req),

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:dcea50b5683ee4fcb2ed84578c6a0250");
+        insta::assert_snapshot!(hash, @"siphash128_13:3d34c7ce3c351dda8198c2c78f7c643a");
     }
 
     {
@@ -385,7 +385,7 @@ fn e2e() {
         });
 
         let hash = replay_log(&tinymist_binary, &root.join("vscode"));
-        insta::assert_snapshot!(hash, @"siphash128_13:d8da3e31bea7c9cf4743c8f80c497242");
+        insta::assert_snapshot!(hash, @"siphash128_13:74ed8df5e8b474e3565659a09a1d7f38");
     }
 }