tinymist_query/
folding_range.rs

1use std::collections::HashSet;
2
3use crate::{
4    SyntaxRequest,
5    prelude::*,
6    syntax::{LexicalHierarchy, LexicalKind, LexicalScopeKind, get_lexical_hierarchy},
7};
8
9/// The [`textDocument/foldingRange`] request is sent from the client to the
10/// server to return all folding ranges found in a given text document.
11///
12/// [`textDocument/foldingRange`]: https://microsoft.github.io/language-server-protocol/specification#textDocument_foldingRange
13///
14/// # Compatibility
15///
16/// This request was introduced in specification version 3.10.0.
17#[derive(Debug, Clone)]
18pub struct FoldingRangeRequest {
19    /// The path of the document to get folding ranges for.
20    pub path: PathBuf,
21    /// If set, the client can only provide folding ranges that consist of whole
22    /// lines.
23    pub line_folding_only: bool,
24}
25
26impl SyntaxRequest for FoldingRangeRequest {
27    type Response = Vec<FoldingRange>;
28
29    fn request(
30        self,
31        source: &Source,
32        position_encoding: PositionEncoding,
33    ) -> Option<Self::Response> {
34        let line_folding_only = self.line_folding_only;
35
36        let hierarchy = get_lexical_hierarchy(source, LexicalScopeKind::Braced)?;
37
38        let mut results = vec![];
39        let LspPosition { line, character } =
40            to_lsp_position(source.text().len(), position_encoding, source);
41        let loc = (line, Some(character));
42
43        calc_folding_range(
44            &hierarchy,
45            source,
46            position_encoding,
47            loc,
48            loc,
49            true,
50            &mut results,
51        );
52
53        // Generally process of folding ranges with line_folding_only
54        if line_folding_only {
55            let mut max_line = 0;
56            for r in &mut results {
57                r.start_character = None;
58                r.end_character = None;
59                max_line = max_line.max(r.end_line);
60            }
61            let mut line_coverage = vec![false; max_line as usize + 1];
62            let mut pair_coverage = HashSet::new();
63            results.reverse();
64            results.retain_mut(|r| {
65                if pair_coverage.contains(&(r.start_line, r.end_line)) {
66                    return false;
67                }
68
69                if line_coverage[r.start_line as usize] {
70                    r.start_line += 1;
71                }
72                if line_coverage[r.end_line as usize] {
73                    r.end_line = r.end_line.saturating_sub(1);
74                }
75                if r.start_line >= r.end_line {
76                    return false;
77                }
78
79                line_coverage[r.start_line as usize] = true;
80                pair_coverage.insert((r.start_line, r.end_line));
81                true
82            });
83            results.reverse();
84        }
85
86        crate::log_debug_ct!(
87            "FoldingRangeRequest(line_folding_only={line_folding_only}) symbols: {hierarchy:#?} results: {results:#?}"
88        );
89
90        Some(results)
91    }
92}
93
94type LoC = (u32, Option<u32>);
95
96fn calc_folding_range(
97    hierarchy: &[LexicalHierarchy],
98    source: &Source,
99    position_encoding: PositionEncoding,
100    parent_last_loc: LoC,
101    last_loc: LoC,
102    is_last_range: bool,
103    folding_ranges: &mut Vec<FoldingRange>,
104) {
105    for (idx, child) in hierarchy.iter().enumerate() {
106        let range = to_lsp_range(child.info.range.clone(), source, position_encoding);
107        let is_not_last_range = idx + 1 < hierarchy.len();
108        let is_not_final_last_range = !is_last_range || is_not_last_range;
109
110        let mut folding_range = FoldingRange {
111            start_line: range.start.line,
112            start_character: Some(range.start.character),
113            end_line: range.end.line,
114            end_character: Some(range.end.character),
115            kind: None,
116            collapsed_text: Some(child.info.name.to_string()),
117        };
118
119        let next_start = if is_not_last_range {
120            let next = &hierarchy[idx + 1];
121            let next_rng = to_lsp_range(next.info.range.clone(), source, position_encoding);
122            (next_rng.start.line, Some(next_rng.start.character))
123        } else if is_not_final_last_range {
124            parent_last_loc
125        } else {
126            last_loc
127        };
128
129        if matches!(child.info.kind, LexicalKind::Heading(..)) {
130            folding_range.end_line = folding_range.end_line.max(if is_not_last_range {
131                next_start.0.saturating_sub(1)
132            } else {
133                next_start.0
134            });
135        }
136
137        if matches!(child.info.kind, LexicalKind::CommentGroup) {
138            folding_range.kind = Some(lsp_types::FoldingRangeKind::Comment);
139        }
140
141        if let Some(ch) = &child.children {
142            let parent_last_loc = if is_not_last_range {
143                (range.end.line, Some(range.end.character))
144            } else {
145                parent_last_loc
146            };
147
148            calc_folding_range(
149                ch,
150                source,
151                position_encoding,
152                parent_last_loc,
153                last_loc,
154                !is_not_final_last_range,
155                folding_ranges,
156            );
157        }
158        folding_ranges.push(folding_range);
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use crate::tests::*;
166
167    #[test]
168    fn test() {
169        snapshot_testing("folding_range", &|world, path| {
170            let r = |line_folding_only| {
171                let request = FoldingRangeRequest {
172                    path: path.clone(),
173                    line_folding_only,
174                };
175
176                let source = world.source_by_path(&path).unwrap();
177
178                request.request(&source, PositionEncoding::Utf16)
179            };
180
181            let result_false = r(false);
182            let result_true = r(true);
183            assert_snapshot!(JsonRepr::new_pure(json!({
184                "false": result_false,
185                "true": result_true,
186            })));
187        });
188    }
189}