tinymist_query/analysis/
doc_highlight.rs

1//! Analyze related expressions to highlight in a source file.
2
3use crate::{prelude::*, syntax::node_ancestors};
4
5/// Analyzes the document and provides related expression information to
6/// highlight.
7pub struct DocumentHighlightWorker<'a> {
8    /// The local analysis context to work with.
9    ctx: &'a mut LocalContext,
10    /// The source document to analyze.
11    source: &'a Source,
12    /// The related expressions to provide.
13    pub annotated: Vec<DocumentHighlight>,
14    /// The worklist to check the nodes.
15    worklist: Vec<LinkedNode<'a>>,
16}
17
18impl<'a> DocumentHighlightWorker<'a> {
19    /// Creates a new worker
20    pub fn new(ctx: &'a mut LocalContext, source: &'a Source) -> Self {
21        Self {
22            ctx,
23            source,
24            annotated: Vec::new(),
25            worklist: Vec::new(),
26        }
27    }
28
29    /// Starts to work
30    pub fn work(&mut self, mut node: &'a LinkedNode<'a>) -> Option<()> {
31        loop {
32            match node.kind() {
33                SyntaxKind::For
34                | SyntaxKind::While
35                | SyntaxKind::Break
36                | SyntaxKind::Continue
37                | SyntaxKind::LoopBreak
38                | SyntaxKind::LoopContinue => return self.work_loop(node),
39                SyntaxKind::Arrow
40                | SyntaxKind::Params
41                | SyntaxKind::Return
42                | SyntaxKind::FuncReturn
43                | SyntaxKind::Contextual => return self.work_func(node),
44                _ => {}
45            }
46            node = node.parent()?;
47        }
48    }
49
50    fn work_loop(&mut self, node: &'a LinkedNode<'a>) -> Option<()> {
51        let _ = self.ctx;
52
53        // find the nearest loop node
54        let loop_node = 'find_loop: {
55            for anc in node_ancestors(node) {
56                if matches!(anc.kind(), SyntaxKind::Contextual | SyntaxKind::Closure) {
57                    return None;
58                }
59                if matches!(anc.kind(), SyntaxKind::ForLoop | SyntaxKind::WhileLoop) {
60                    break 'find_loop anc;
61                }
62            }
63            return None;
64        };
65
66        // find the first key word of the loop node
67        let keyword = loop_node.children().find(|node| node.kind().is_keyword());
68        if let Some(keyword) = keyword {
69            self.annotate(&keyword);
70        }
71
72        self.check_children(loop_node);
73        self.check(Self::check_loop);
74
75        crate::log_debug_ct!("highlights: {:?}", self.annotated);
76        Some(())
77    }
78
79    fn work_func(&mut self, _node: &'a LinkedNode<'a>) -> Option<()> {
80        None
81    }
82
83    /// Annotate the node for highlight
84    fn annotate(&mut self, node: &LinkedNode) {
85        let mut rng = node.range();
86
87        // if previous node is hash
88        if rng.start > 0 && self.source.text().as_bytes()[rng.start - 1] == b'#' {
89            rng.start -= 1;
90        }
91
92        self.annotated.push(DocumentHighlight {
93            range: self.ctx.to_lsp_range(rng, self.source),
94            kind: None,
95        });
96    }
97
98    /// Consumes the worklist and checks the nodes
99    fn check<F>(&mut self, check: F)
100    where
101        F: Fn(&mut Self, LinkedNode<'a>),
102    {
103        while let Some(node) = self.worklist.pop() {
104            check(self, node);
105        }
106    }
107
108    /// Pushes the children of the node to check
109    fn check_children(&mut self, node: &LinkedNode<'a>) {
110        if node.get().children().len() == 0 {
111            return;
112        }
113
114        for child in node.children() {
115            self.worklist.push(child.clone());
116        }
117    }
118
119    fn check_loop(&mut self, node: LinkedNode<'a>) {
120        match node.kind() {
121            SyntaxKind::ForLoop
122            | SyntaxKind::WhileLoop
123            | SyntaxKind::Closure
124            | SyntaxKind::Contextual => {
125                return;
126            }
127            SyntaxKind::LoopBreak | SyntaxKind::LoopContinue => {
128                self.annotate(&node);
129                return;
130            }
131            _ => {}
132        }
133
134        self.check_children(&node);
135    }
136}