tinymist_query/
semantic_tokens_full.rs

1use crate::prelude::*;
2
3/// The [`textDocument/semanticTokens/full`] request is sent from the client to
4/// the server to resolve the semantic tokens of a given file.
5///
6/// [`textDocument/semanticTokens/full`]: https://microsoft.github.io/language-server-protocol/specification#textDocument_semanticTokens
7///
8/// Semantic tokens are used to add additional color information to a file that
9/// depends on language specific symbol information. A semantic token request
10/// usually produces a large result. The protocol therefore supports encoding
11/// tokens with numbers. In addition, optional support for deltas is available,
12/// i.e. [`semantic_tokens_full_delta`].
13///
14/// [`semantic_tokens_full_delta`]: crate::SemanticTokensDeltaRequest
15///
16/// # Compatibility
17///
18/// This request was introduced in specification version 3.16.0.
19#[derive(Debug, Clone)]
20pub struct SemanticTokensFullRequest {
21    /// The path of the document to get semantic tokens for.
22    pub path: PathBuf,
23}
24
25impl SemanticRequest for SemanticTokensFullRequest {
26    type Response = SemanticTokensResult;
27
28    /// Handles the request to compute the semantic tokens for a given document.
29    fn request(self, ctx: &mut LocalContext) -> Option<Self::Response> {
30        let source = ctx.source_by_path(&self.path).ok()?;
31        let (tokens, result_id) = ctx.cached_tokens(&source);
32
33        Some(SemanticTokensResult::Tokens(SemanticTokens {
34            result_id,
35            data: tokens.as_ref().clone(),
36        }))
37    }
38}
39
40impl SemanticTokensFullRequest {
41    /// Computes the semantic tokens for a given source code.
42    pub fn compute(ctx: &mut LocalContext, source: &Source) -> crate::analysis::SemanticTokens {
43        crate::analysis::semantic_tokens::get_semantic_tokens(ctx, source)
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use crate::tests::*;
51
52    /// This is converted by Copilot from TypeScript `to_multiline_tokens2`.
53    /// <https://github.com/microsoft/vscode/blob/2acc0e52cbc7434c415f221d5c34ee1bbdd6cd71/src/vs/editor/common/services/semanticTokensProviderStyling.ts#L147>
54    fn check_tokens(tokens: &SemanticTokens) {
55        const DESIRED_TOKENS_PER_AREA: usize = 400;
56        const DESIRED_MAX_AREAS: usize = 1024;
57
58        let src_data = &tokens.data;
59        let token_count = src_data.len();
60        let tokens_per_area = std::cmp::max(
61            (token_count as f64 / DESIRED_MAX_AREAS as f64).ceil() as usize,
62            DESIRED_TOKENS_PER_AREA,
63        );
64
65        let mut token_index = 0;
66        let mut last_line_number = 1;
67        let mut last_start_character = 0;
68
69        while token_index < token_count {
70            let token_start_index = token_index;
71            let mut token_end_index =
72                std::cmp::min(token_start_index + tokens_per_area, token_count);
73
74            // Keep tokens on the same line in the same area...
75            if token_end_index < token_count {
76                let mut small_token_end_index = token_end_index;
77                while small_token_end_index - 1 > token_start_index
78                    && src_data[small_token_end_index].delta_line == 0
79                {
80                    small_token_end_index -= 1;
81                }
82
83                if small_token_end_index - 1 == token_start_index {
84                    // there are so many tokens on this line that our area would be empty, we must
85                    // now go right
86                    let mut big_token_end_index = token_end_index;
87                    while big_token_end_index + 1 < token_count
88                        && src_data[big_token_end_index].delta_line == 0
89                    {
90                        big_token_end_index += 1;
91                    }
92                    token_end_index = big_token_end_index;
93                } else {
94                    token_end_index = small_token_end_index;
95                }
96            }
97
98            let mut prev_line_number = 0;
99            let mut prev_end_character = 0;
100
101            while token_index < token_end_index {
102                let delta_line = src_data[token_index].delta_line;
103                let delta_character = src_data[token_index].delta_start;
104                let length = src_data[token_index].length;
105                let line_number = last_line_number + delta_line;
106                let start_character = if delta_line == 0 {
107                    last_start_character + delta_character
108                } else {
109                    delta_character
110                };
111                // delta_character
112                let end_character = start_character + length;
113
114                if end_character <= start_character {
115                    // this token is invalid (most likely a negative length casted to uint32)
116                    panic!(
117                        "Invalid length for semantic token at line {line_number}, character {start_character}, end: {end_character}"
118                    );
119                } else if prev_line_number == line_number && prev_end_character > start_character {
120                    // this token overlaps with the previous token
121                    panic!(
122                        "Overlapping semantic tokens at line {line_number}, character {start_character}, previous line {prev_line_number}, previous end {prev_end_character}"
123                    );
124                } else {
125                    prev_line_number = line_number;
126                    prev_end_character = end_character;
127                }
128
129                last_line_number = line_number;
130                last_start_character = start_character;
131                token_index += 1;
132            }
133        }
134    }
135
136    #[test]
137    fn test() {
138        snapshot_testing("semantic_tokens", &|ctx, path| {
139            let request = SemanticTokensFullRequest { path: path.clone() };
140
141            let mut result = request.request(ctx).unwrap();
142            if let SemanticTokensResult::Tokens(tokens) = &mut result {
143                tokens.result_id.take();
144            }
145
146            match &result {
147                SemanticTokensResult::Tokens(tokens) => {
148                    check_tokens(tokens);
149                }
150                SemanticTokensResult::Partial(_) => {
151                    panic!("Unexpected partial result");
152                }
153            }
154
155            assert_snapshot!(serde_json::to_string(&result).unwrap());
156        });
157    }
158}