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

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,
-    Hover, InlayHint, LanguageString, Location as LspLocation, LocationLink, MarkedString,
-    MarkupContent, MarkupKind, Position as LspPosition, PrepareRenameResponse, SelectionRange,
-    SemanticTokens, SemanticTokensDelta, SemanticTokensFullDeltaResult, SemanticTokensResult,
-    SignatureHelp, SignatureInformation, SymbolInformation, TextEdit, Url, WorkspaceEdit,
+    DocumentHighlight, DocumentSymbol, DocumentSymbolResponse, Documentation, FoldingRange,
+    GotoDefinitionResponse, Hover, InlayHint, LanguageString, Location as LspLocation,
+    LocationLink, MarkedString, MarkupContent, MarkupKind, Position as LspPosition,
+    PrepareRenameResponse, SelectionRange, SemanticTokens, SemanticTokensDelta,
+    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(||