docs: split and documenting document highlight worker (#989)

crates/tinymist-query/src/analysis.rs

@@ -10,6 +10,8 @@ pub mod code_action;
 pub use code_action::*;
 pub mod color_expr;
 pub use color_expr::*;
+pub mod doc_highlight;
+pub use doc_highlight::*;
 pub mod link_exprs;
 pub use link_exprs::*;
 pub mod stats;

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

@@ -0,0 +1,126 @@
+//! Analyze related expressions to highlight in a source file.
+
+use crate::{prelude::*, syntax::node_ancestors};
+
+/// Analyzes the document and provides related expression information to
+/// highlight.
+pub struct DocumentHighlightWorker<'a> {
+    /// The local analysis context to work with.
+    ctx: &'a mut LocalContext,
+    /// The source document to analyze.
+    source: &'a Source,
+    /// The related expressions to provide.
+    pub annotated: Vec<DocumentHighlight>,
+    /// The worklist to check the nodes.
+    worklist: Vec<LinkedNode<'a>>,
+}
+
+impl<'a> DocumentHighlightWorker<'a> {
+    /// Creates a new worker
+    pub fn new(ctx: &'a mut LocalContext, source: &'a Source) -> Self {
+        Self {
+            ctx,
+            source,
+            annotated: Vec::new(),
+            worklist: Vec::new(),
+        }
+    }
+
+    /// Starts to work
+    pub fn work(&mut self, mut node: &'a LinkedNode<'a>) -> Option<()> {
+        loop {
+            match node.kind() {
+                SyntaxKind::For
+                | SyntaxKind::While
+                | SyntaxKind::Break
+                | SyntaxKind::Continue
+                | SyntaxKind::LoopBreak
+                | SyntaxKind::LoopContinue => return self.work_loop(node),
+                SyntaxKind::Arrow
+                | SyntaxKind::Params
+                | SyntaxKind::Return
+                | SyntaxKind::FuncReturn => return self.work_func(node),
+                _ => {}
+            }
+            node = node.parent()?;
+        }
+    }
+
+    fn work_loop(&mut self, node: &'a LinkedNode<'a>) -> Option<()> {
+        let _ = self.ctx;
+
+        // find the nearest loop node
+        let loop_node = 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);
+
+        crate::log_debug_ct!("highlights: {:?}", self.annotated);
+        Some(())
+    }
+
+    fn work_func(&mut self, _node: &'a LinkedNode<'a>) -> Option<()> {
+        None
+    }
+
+    /// Annotate the node for highlight
+    fn annotate(&mut self, node: &LinkedNode) {
+        let mut rng = node.range();
+
+        // if previous node is hash
+        if rng.start > 0 && self.source.text().as_bytes()[rng.start - 1] == b'#' {
+            rng.start -= 1;
+        }
+
+        self.annotated.push(DocumentHighlight {
+            range: self.ctx.to_lsp_range(rng, self.source),
+            kind: None,
+        });
+    }
+
+    /// Consumes the worklist and checks the nodes
+    fn check<F>(&mut self, check: F)
+    where
+        F: Fn(&mut Self, LinkedNode<'a>),
+    {
+        while let Some(node) = self.worklist.pop() {
+            check(self, node);
+        }
+    }
+
+    /// Pushes the children of the node to check
+    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);
+    }
+}

crates/tinymist-query/src/document_highlight.rs

@@ -1,6 +1,6 @@
 use typst_shim::syntax::LinkedNodeExt;
 
-use crate::{prelude::*, syntax::node_ancestors, SemanticRequest};
+use crate::{analysis::doc_highlight::DocumentHighlightWorker, prelude::*, SemanticRequest};
 
 /// The [`textDocument/documentHighlight`] request
 ///
@@ -21,131 +21,14 @@ impl SemanticRequest for DocumentHighlightRequest {
         let cursor = ctx.to_typst_pos(self.position, &source)?;
 
         let root = LinkedNode::new(source.root());
-        let mut node = &root.leaf_at_compat(cursor)?;
+        let node = root.leaf_at_compat(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()?;
-        }
+        let mut worker = DocumentHighlightWorker::new(ctx, &source);
+        worker.work(&node)?;
+        (!worker.annotated.is_empty()).then_some(worker.annotated)
     }
 }
 
-struct DocumentHighlightWorker<'a> {
-    ctx: &'a mut LocalContext,
-    current: &'a Source,
-    highlights: Vec<DocumentHighlight>,
-    worklist: Vec<LinkedNode<'a>>,
-}
-
-impl<'a> DocumentHighlightWorker<'a> {
-    fn new(ctx: &'a mut LocalContext, 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 = 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);
-
-        crate::log_debug_ct!("highlights: {:?}", self.highlights);
-        self.finish()
-    }
-}
-
-fn highlight_func_returns(
-    ctx: &mut LocalContext,
-    node: &LinkedNode,
-) -> Option<Vec<DocumentHighlight>> {
-    let _ = ctx;
-    let _ = node;
-    None
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;