Skip to main content

sql_dialect_fmt_lsp/
lib.rs

1//! LSP feature logic for Snowflake SQL: formatting edits, parse diagnostics, and semantic tokens.
2//!
3//! This module is deliberately transport-free (it never touches stdio or an `lsp-server`
4//! connection), so every feature is a pure `&str -> data` function that is unit-testable. The
5//! binary ([`crate::main`]) is the thin adapter that wires these into a language server.
6//!
7//! Positions follow the LSP convention: zero-based lines and **UTF-16** column offsets.
8
9use lsp_types::{
10    Diagnostic, DiagnosticSeverity, FoldingRange, FoldingRangeKind, Hover, HoverContents,
11    MarkupContent, MarkupKind, Position, Range, SemanticToken, SemanticTokenModifier,
12    SemanticTokenType, TextEdit,
13};
14use sql_dialect_fmt_formatter::{format, FormatOptions};
15use sql_dialect_fmt_highlight::{semantic, HighlightKind};
16
17/// The semantic-token type legend, mirrored from the single source of truth in
18/// `sql-dialect-fmt-highlight` ([`semantic::SemanticTokenType::LEGEND`]). A token's `token_type` field is
19/// an index into this slice, so the order here is the contract with the editor (declared in the
20/// server's capabilities) and *must* equal the highlighter's legend — otherwise an editor would
21/// decode our token types against the wrong names.
22pub fn token_types() -> Vec<SemanticTokenType> {
23    semantic::SemanticTokenType::LEGEND
24        .iter()
25        .map(|ty| SemanticTokenType::new(ty.name()))
26        .collect()
27}
28
29/// The semantic-token modifier legend, mirrored from
30/// [`semantic::SemanticTokenModifiers::LEGEND`]. A token's `token_modifiers_bitset` is decoded
31/// bit-by-bit against this slice, so it too must match the highlighter.
32pub fn token_modifiers() -> Vec<SemanticTokenModifier> {
33    semantic::SemanticTokenModifiers::LEGEND
34        .iter()
35        .map(|&name| SemanticTokenModifier::new(name))
36        .collect()
37}
38
39/// Maps byte offsets into a document to LSP [`Position`]s (UTF-16 columns).
40pub struct LineIndex<'a> {
41    text: &'a str,
42    /// Byte offset of the start of each line.
43    line_starts: Vec<usize>,
44}
45
46impl<'a> LineIndex<'a> {
47    pub fn new(text: &'a str) -> Self {
48        let mut line_starts = vec![0];
49        line_starts.extend(
50            text.bytes()
51                .enumerate()
52                .filter(|&(_, b)| b == b'\n')
53                .map(|(i, _)| i + 1),
54        );
55        LineIndex { text, line_starts }
56    }
57
58    /// The LSP position of a byte `offset` (clamped to the document end).
59    pub fn position(&self, offset: usize) -> Position {
60        let offset = offset.min(self.text.len());
61        let line = match self.line_starts.binary_search(&offset) {
62            Ok(line) => line,
63            Err(next) => next - 1,
64        };
65        let line_start = self.line_starts[line];
66        let col: usize = self.text[line_start..offset]
67            .chars()
68            .map(char::len_utf16)
69            .sum();
70        Position::new(line as u32, col as u32)
71    }
72
73    /// The position one past the last character — the end of the document.
74    pub fn end(&self) -> Position {
75        self.position(self.text.len())
76    }
77
78    /// The byte offset of an LSP [`Position`] (the inverse of [`Self::position`]). Out-of-range
79    /// lines/columns clamp to the line or document end.
80    pub fn offset(&self, position: Position) -> usize {
81        let line = position.line as usize;
82        let Some(&line_start) = self.line_starts.get(line) else {
83            return self.text.len();
84        };
85        let mut remaining = position.character as usize; // UTF-16 units to consume
86        let mut offset = line_start;
87        for ch in self.text[line_start..].chars() {
88            let width = ch.len_utf16();
89            if remaining < width || ch == '\n' {
90                break;
91            }
92            remaining -= width;
93            offset += ch.len_utf8();
94        }
95        offset
96    }
97}
98
99/// The edits to apply for `textDocument/formatting`: a single whole-document replacement, or an
100/// empty list when the input is already formatted (so the editor records no change).
101pub fn format_edits(text: &str, options: &FormatOptions) -> Vec<TextEdit> {
102    let formatted = format(text, options);
103    if formatted == text {
104        return Vec::new();
105    }
106    let index = LineIndex::new(text);
107    vec![TextEdit {
108        range: Range::new(Position::new(0, 0), index.end()),
109        new_text: formatted,
110    }]
111}
112
113/// Diagnostics for `textDocument/publishDiagnostics`: both lexer and parser errors. Neither stage
114/// ever fails, so this is the set of *recovered* errors (empty for clean input).
115///
116/// Lexer errors (unterminated literals/comments, stray characters) are surfaced too — they are
117/// reported by the tokenizer before the parser runs, so the parser-only error list would otherwise
118/// miss them. We read them through the highlighter, which already re-exposes the lexer's errors.
119/// Each diagnostic's range covers the whole offending token (via its byte `range()`), not a single
120/// character, so editors underline the real span.
121pub fn diagnostics(text: &str) -> Vec<Diagnostic> {
122    let index = LineIndex::new(text);
123    let to_range = |span: std::ops::Range<usize>| {
124        Range::new(index.position(span.start), index.position(span.end))
125    };
126    let make = |range: Range, message: String| Diagnostic {
127        range,
128        severity: Some(DiagnosticSeverity::ERROR),
129        source: Some("sql-dialect-fmt".to_string()),
130        message,
131        ..Default::default()
132    };
133
134    let lex_errors = sql_dialect_fmt_highlight::highlight(text).errors;
135    let parse = sql_dialect_fmt_parser::parse(text);
136    let mut diagnostics: Vec<_> = lex_errors
137        .into_iter()
138        .map(|err| make(to_range(err.range()), err.message))
139        .chain(
140            parse
141                .errors()
142                .iter()
143                .map(|err| make(to_range(err.range()), err.message.clone())),
144        )
145        .collect();
146    diagnostics.extend(embedded_language_diagnostics(text, &index));
147    diagnostics
148}
149
150fn embedded_language_diagnostics(text: &str, index: &LineIndex<'_>) -> Vec<Diagnostic> {
151    let mut diagnostics = Vec::new();
152    let mut expect_language_name = false;
153    let mut language_name: Option<(&str, std::ops::Range<usize>)> = None;
154    let mut saw_as_after_language = false;
155
156    for token in sql_dialect_fmt_highlight::highlight(text).tokens {
157        match token.kind {
158            HighlightKind::Whitespace | HighlightKind::Comment => {}
159            HighlightKind::Punctuation if token.text == ";" => {
160                expect_language_name = false;
161                language_name = None;
162                saw_as_after_language = false;
163            }
164            HighlightKind::DollarString => {
165                if saw_as_after_language {
166                    if let Some((word, range)) = language_name.take() {
167                        if !is_supported_embedded_language(word) {
168                            diagnostics.push(Diagnostic {
169                                range: Range::new(index.position(range.start), index.position(range.end)),
170                                severity: Some(DiagnosticSeverity::WARNING),
171                                source: Some("sql-dialect-fmt".to_string()),
172                                message: format!(
173                                    "unsupported embedded language {word}; expected SQL, JAVASCRIPT, PYTHON, JAVA, or SCALA"
174                                ),
175                                ..Default::default()
176                            });
177                        }
178                    }
179                }
180                expect_language_name = false;
181                saw_as_after_language = false;
182            }
183            HighlightKind::Keyword if token.text.eq_ignore_ascii_case("language") => {
184                expect_language_name = true;
185                language_name = None;
186                saw_as_after_language = false;
187            }
188            HighlightKind::Keyword | HighlightKind::Identifier | HighlightKind::Type
189                if expect_language_name =>
190            {
191                language_name = Some((token.text, token.range));
192                expect_language_name = false;
193            }
194            HighlightKind::Keyword
195                if language_name.is_some() && token.text.eq_ignore_ascii_case("as") =>
196            {
197                saw_as_after_language = true;
198            }
199            _ => {
200                expect_language_name = false;
201            }
202        }
203    }
204
205    diagnostics
206}
207
208fn is_supported_embedded_language(word: &str) -> bool {
209    ["SQL", "JAVASCRIPT", "PYTHON", "JAVA", "SCALA"]
210        .iter()
211        .any(|candidate| candidate.eq_ignore_ascii_case(word))
212}
213
214/// Hover information for `textDocument/hover`: the keyword/type/symbol description at `position`,
215/// rendered as Markdown with an optional docs link, scoped to the hovered token's range.
216pub fn hover(text: &str, position: Position) -> Option<Hover> {
217    let index = LineIndex::new(text);
218    let info = sql_dialect_fmt_hover::hover_at(text, index.offset(position))?;
219    let mut value = format!("**{}**\n\n{}", info.title, info.body);
220    if let Some(url) = info.docs_url {
221        value.push_str(&format!("\n\n[Snowflake docs]({url})"));
222    }
223    Some(Hover {
224        contents: HoverContents::Markup(MarkupContent {
225            kind: MarkupKind::Markdown,
226            value,
227        }),
228        range: Some(Range::new(
229            index.position(info.range.start),
230            index.position(info.range.end),
231        )),
232    })
233}
234
235/// Folding ranges for `textDocument/foldingRange`: one region per multi-line top-level statement,
236/// so an editor can collapse each statement in a script. The CST's root children are the statements.
237pub fn folding_ranges(text: &str) -> Vec<FoldingRange> {
238    let index = LineIndex::new(text);
239    let root = sql_dialect_fmt_parser::parse(text).syntax();
240    root.children()
241        .filter_map(|stmt| {
242            // Use the span of the statement's significant tokens, so a leading/trailing blank line
243            // (attached as trivia) doesn't inflate a single-line statement into a foldable region.
244            let mut tokens = stmt
245                .descendants_with_tokens()
246                .filter_map(|el| el.into_token())
247                .filter(|t| !t.kind().is_trivia());
248            let first = tokens.next()?;
249            let last = tokens.last().unwrap_or_else(|| first.clone());
250            let start = index.position(first.text_range().start().into()).line;
251            let end = index.position(last.text_range().end().into()).line;
252            (end > start).then_some(FoldingRange {
253                start_line: start,
254                end_line: end,
255                kind: Some(FoldingRangeKind::Region),
256                ..FoldingRange::default()
257            })
258        })
259        .collect()
260}
261
262/// Apply one `textDocument/didChange` content change to `text`, returning the new document.
263///
264/// Supports incremental sync: a `Some(range)` splices `new_text` over the byte span the range
265/// covers; a `None` range is a whole-document replacement (the editor sent the full text). The
266/// range is treated as ordered and clamped to the document, so a malformed event can't panic.
267pub fn apply_change(text: &str, range: Option<Range>, new_text: &str) -> String {
268    let Some(range) = range else {
269        return new_text.to_string();
270    };
271    let index = LineIndex::new(text);
272    let a = index.offset(range.start);
273    let b = index.offset(range.end);
274    let (start, end) = (a.min(b), a.max(b));
275    let mut out = String::with_capacity(text.len() - (end - start) + new_text.len());
276    out.push_str(&text[..start]);
277    out.push_str(new_text);
278    out.push_str(&text[end..]);
279    out
280}
281
282/// The delta-encoded semantic tokens for `textDocument/semanticTokens/full`.
283///
284/// This is a thin adapter over `sql-dialect-fmt-highlight`'s [`semantic::semantic_tokens_lsp`]: the
285/// highlighter already splits multi-line tokens (block comments, dollar-quoted strings) into one
286/// token per line — the LSP encoding requires each token to stay on a single line — computes UTF-16
287/// columns/lengths, and delta-encodes into `(deltaLine, deltaStartChar, length, tokenType,
288/// tokenModifiers)` quintuples against the same legend the server advertises. We only reshape each
289/// quintuple into an `lsp_types::SemanticToken`, **preserving** the modifier bitset (rather than
290/// hardcoding 0) so `defaultLibrary` / `documentation` modifiers reach the editor.
291pub fn semantic_tokens(text: &str) -> Vec<SemanticToken> {
292    semantic::semantic_tokens_lsp(text)
293        .into_iter()
294        .map(
295            |[delta_line, delta_start, length, token_type, token_modifiers_bitset]| SemanticToken {
296                delta_line,
297                delta_start,
298                length,
299                token_type,
300                token_modifiers_bitset,
301            },
302        )
303        .collect()
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309
310    #[test]
311    fn line_index_maps_offsets_to_utf16_positions() {
312        let text = "SELECT a\nFROM 芋;\n"; // 芋 is one UTF-16 unit but 3 bytes
313        let index = LineIndex::new(text);
314        assert_eq!(index.position(0), Position::new(0, 0));
315        assert_eq!(index.position(7), Position::new(0, 7)); // the `a`
316        let from = text.find("FROM").unwrap();
317        assert_eq!(index.position(from), Position::new(1, 0));
318        let semicolon = text.find(';').unwrap();
319        assert_eq!(index.position(semicolon), Position::new(1, 6)); // FROM<sp>芋 = 6 utf16 units
320    }
321
322    #[test]
323    fn formatting_replaces_the_whole_document() {
324        let edits = format_edits("select a,b from t", &FormatOptions::default());
325        assert_eq!(edits.len(), 1);
326        assert_eq!(edits[0].new_text, "SELECT a, b\nFROM t;\n");
327        assert_eq!(edits[0].range.start, Position::new(0, 0));
328    }
329
330    #[test]
331    fn already_formatted_input_yields_no_edits() {
332        let formatted = "SELECT a, b\nFROM t;\n";
333        assert!(format_edits(formatted, &FormatOptions::default()).is_empty());
334    }
335
336    #[test]
337    fn clean_sql_has_no_diagnostics() {
338        assert!(diagnostics("select 1").is_empty());
339    }
340
341    #[test]
342    fn broken_sql_reports_a_diagnostic() {
343        let diags = diagnostics("select from where");
344        assert!(!diags.is_empty());
345        assert_eq!(diags[0].severity, Some(DiagnosticSeverity::ERROR));
346    }
347
348    #[test]
349    fn lexer_errors_reach_diagnostics() {
350        // An unterminated string is a *lexer* error (the parser never sees it as a token boundary
351        // problem). It must still surface as a diagnostic, and its range must cover the whole
352        // unterminated literal, not a single character.
353        let text = "SELECT 'oops";
354        let diags = diagnostics(text);
355        let lex_diag = diags
356            .iter()
357            .find(|d| d.message.contains("unterminated string"))
358            .expect("an unterminated-string diagnostic");
359        assert_eq!(lex_diag.severity, Some(DiagnosticSeverity::ERROR));
360        let quote = text.find('\'').unwrap() as u32;
361        assert_eq!(lex_diag.range.start, Position::new(0, quote));
362        // Spans to end of the line (the whole literal), so end column > start column.
363        assert!(lex_diag.range.end.character > lex_diag.range.start.character);
364    }
365
366    #[test]
367    fn parser_diagnostic_range_covers_the_token() {
368        // `MERGE tgt ...` (no INTO) reports "expected INTO" at the `tgt` token; the LSP range must
369        // span the whole 3-character identifier, not one character.
370        let text = "MERGE tgt USING src ON a = b";
371        let diags = diagnostics(text);
372        let into = diags
373            .iter()
374            .find(|d| d.message == "expected INTO")
375            .expect("an INTO diagnostic");
376        let col = text.find("tgt").unwrap() as u32;
377        assert_eq!(into.range.start, Position::new(0, col));
378        assert_eq!(into.range.end, Position::new(0, col + 3));
379    }
380
381    #[test]
382    fn clean_sql_still_has_no_lexer_or_parser_diagnostics() {
383        assert!(diagnostics("SELECT a FROM t").is_empty());
384    }
385
386    #[test]
387    fn unsupported_embedded_language_is_a_warning() {
388        let text = "CREATE FUNCTION f() RETURNS STRING LANGUAGE RUBY AS $$x$$;";
389        let diags = diagnostics(text);
390        let language = diags
391            .iter()
392            .find(|d| d.message.contains("unsupported embedded language RUBY"))
393            .expect("unsupported-language diagnostic");
394        assert_eq!(language.severity, Some(DiagnosticSeverity::WARNING));
395        let col = text.find("RUBY").unwrap() as u32;
396        assert_eq!(language.range.start, Position::new(0, col));
397        assert_eq!(language.range.end, Position::new(0, col + 4));
398    }
399
400    #[test]
401    fn embedded_language_warning_does_not_fire_for_plain_columns_or_dynamic_sql() {
402        for text in [
403            "SELECT language FROM t;",
404            "EXECUTE IMMEDIATE $$ SELECT 1 $$;",
405        ] {
406            assert!(
407                diagnostics(text)
408                    .iter()
409                    .all(|diag| !diag.message.contains("unsupported embedded language")),
410                "{text}"
411            );
412        }
413    }
414
415    #[test]
416    fn offset_is_the_inverse_of_position() {
417        let text = "SELECT a\nFROM 芋;\n";
418        let index = LineIndex::new(text);
419        for offset in [
420            0usize,
421            7,
422            text.find("FROM").unwrap(),
423            text.find(';').unwrap(),
424        ] {
425            assert_eq!(index.offset(index.position(offset)), offset);
426        }
427    }
428
429    #[test]
430    fn hover_describes_a_type() {
431        // Hover over the `varchar` cast target should return a Snowflake type description.
432        let src = "select x::varchar from t";
433        let col = src.find("varchar").unwrap() as u32;
434        let hover = hover(src, Position::new(0, col)).expect("hover");
435        assert!(hover.range.is_some());
436        match hover.contents {
437            HoverContents::Markup(m) => assert!(m.value.to_lowercase().contains("varchar")),
438            _ => panic!("expected markup"),
439        }
440    }
441
442    #[test]
443    fn apply_change_splices_an_incremental_edit() {
444        // Replace "world" (line 1, cols 0..5) with "snow".
445        let text = "hello\nworld\n";
446        let range = Range::new(Position::new(1, 0), Position::new(1, 5));
447        assert_eq!(apply_change(text, Some(range), "snow"), "hello\nsnow\n");
448    }
449
450    #[test]
451    fn apply_change_with_no_range_replaces_whole_document() {
452        assert_eq!(apply_change("old", None, "new text"), "new text");
453    }
454
455    #[test]
456    fn folding_ranges_cover_multiline_statements() {
457        let ranges = folding_ranges("select a,\nb\nfrom t;\n\nselect 1;");
458        assert_eq!(ranges.len(), 1); // only the first (multi-line) statement folds
459        assert_eq!(ranges[0].start_line, 0);
460        assert_eq!(ranges[0].end_line, 2);
461    }
462
463    #[test]
464    fn semantic_tokens_tag_keywords() {
465        let tokens = semantic_tokens("select a from t");
466        assert!(!tokens.is_empty());
467        // The first token is `select`, a keyword (legend index 0), at line 0 column 0.
468        assert_eq!(tokens[0].delta_line, 0);
469        assert_eq!(tokens[0].delta_start, 0);
470        assert_eq!(tokens[0].length, 6);
471        assert_eq!(tokens[0].token_type, 0);
472        // The keyword carries the `defaultLibrary` modifier — it must not be hardcoded to 0.
473        assert_eq!(
474            tokens[0].token_modifiers_bitset,
475            semantic::SemanticTokenModifiers::DEFAULT_LIBRARY.bits()
476        );
477    }
478
479    #[test]
480    fn server_legend_equals_the_highlighter_legend() {
481        // The advertised type legend must be exactly the highlighter's LEGEND, in order — this is
482        // the contract an editor decodes `token_type` against.
483        let advertised: Vec<String> = token_types()
484            .iter()
485            .map(|t| t.as_str().to_string())
486            .collect();
487        let expected: Vec<String> = semantic::SemanticTokenType::LEGEND
488            .iter()
489            .map(|t| t.name().to_string())
490            .collect();
491        assert_eq!(advertised, expected);
492        // It includes `function` (currently the last appended type) for Cortex/AISQL recognition.
493        assert_eq!(advertised.last().map(String::as_str), Some("function"));
494
495        // Likewise the modifier legend mirrors the highlighter's.
496        let mods: Vec<String> = token_modifiers()
497            .iter()
498            .map(|m| m.as_str().to_string())
499            .collect();
500        assert_eq!(mods, vec!["documentation", "defaultLibrary"]);
501    }
502
503    #[test]
504    fn semantic_tokens_are_monotonic_and_never_panic_on_multiline() {
505        // A multi-line block comment must split into per-line tokens without panicking.
506        let tokens = semantic_tokens("select 1 /* a\nb */ from t");
507        // Deltas must be non-negative by construction (u32) and the stream stays consistent.
508        assert!(tokens.iter().all(|t| t.length > 0));
509    }
510}