sql_dialect_fmt_highlight/
semantic.rs

1//! LSP semantic-token mapping for the lexical highlighter.
2//!
3//! The [`crate::highlight`] pass produces [`HighlightKind`]s with byte ranges. Editors that speak
4//! the Language Server Protocol want *semantic tokens*: each significant token tagged with a
5//! standard token **type** (`keyword`, `string`, `number`, …) and a bitset of **modifiers**, then
6//! delta-encoded as `(deltaLine, deltaStartChar, length, tokenType, modifiers)` quintuples.
7//!
8//! This module is transport-free and depends on nothing but the lexer/syntax crates — the
9//! `sql-dialect-fmt-lsp` binary maps these onto `lsp_types`, but the mapping itself (and its tests) live
10//! here so the highlighter and the editor adapter can never disagree on what a kind *means*.
11//!
12//! It also understands [`Injection`]s: a `$$ … $$` body may embed JavaScript, Python, Java, Scala,
13//! or SQL (per a `LANGUAGE` clause). Tokens that fall inside an injection region are tagged so an
14//! editor can either re-highlight them with the embedded grammar or shade the whole body.
15
16use crate::{highlight, HighlightKind, HighlightToken};
17use sql_dialect_fmt_text::{utf16_len, LineIndex};
18
19/// The standard LSP semantic token types this highlighter emits.
20///
21/// The variant order is the *legend*: [`SemanticTokenType::index`] is the position an editor uses
22/// to decode the `tokenType` field, so adding a variant must only ever append.
23#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
24pub enum SemanticTokenType {
25    Keyword,
26    Type,
27    Variable,
28    String,
29    Number,
30    /// Bind / positional parameters (`$1`, `:name`, `?`) and session variables (`$name`).
31    Parameter,
32    Operator,
33    Comment,
34    /// Stage references (`@stage`, `@~`, `@%table`) — `namespace` is the closest standard type.
35    Namespace,
36    /// Built-in function names, currently used for Snowflake Cortex / AISQL functions.
37    Function,
38}
39
40impl SemanticTokenType {
41    /// The full legend, in `index` order. This is the contract an editor declares in its server
42    /// capabilities; the `tokenType` field of every emitted token indexes into this slice.
43    pub const LEGEND: &'static [SemanticTokenType] = &[
44        SemanticTokenType::Keyword,
45        SemanticTokenType::Type,
46        SemanticTokenType::Variable,
47        SemanticTokenType::String,
48        SemanticTokenType::Number,
49        SemanticTokenType::Parameter,
50        SemanticTokenType::Operator,
51        SemanticTokenType::Comment,
52        SemanticTokenType::Namespace,
53        SemanticTokenType::Function,
54    ];
55
56    /// The canonical LSP `SemanticTokenType` string (matches `lsp_types::SemanticTokenType`).
57    pub const fn name(self) -> &'static str {
58        match self {
59            SemanticTokenType::Keyword => "keyword",
60            SemanticTokenType::Type => "type",
61            SemanticTokenType::Variable => "variable",
62            SemanticTokenType::String => "string",
63            SemanticTokenType::Number => "number",
64            SemanticTokenType::Parameter => "parameter",
65            SemanticTokenType::Operator => "operator",
66            SemanticTokenType::Comment => "comment",
67            SemanticTokenType::Namespace => "namespace",
68            SemanticTokenType::Function => "function",
69        }
70    }
71
72    /// The legend index — the value carried in a token's `token_type` field.
73    pub const fn index(self) -> u32 {
74        match self {
75            SemanticTokenType::Keyword => 0,
76            SemanticTokenType::Type => 1,
77            SemanticTokenType::Variable => 2,
78            SemanticTokenType::String => 3,
79            SemanticTokenType::Number => 4,
80            SemanticTokenType::Parameter => 5,
81            SemanticTokenType::Operator => 6,
82            SemanticTokenType::Comment => 7,
83            SemanticTokenType::Namespace => 8,
84            SemanticTokenType::Function => 9,
85        }
86    }
87}
88
89/// LSP semantic token modifiers, as a bitset. Only the few the lexical layer can justify are
90/// modelled; `bits()` produces the `tokenModifiers` value an editor decodes against its legend.
91#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
92pub struct SemanticTokenModifiers(u32);
93
94impl SemanticTokenModifiers {
95    pub const NONE: SemanticTokenModifiers = SemanticTokenModifiers(0);
96    /// `documentation` — block/line comments.
97    pub const DOCUMENTATION: SemanticTokenModifiers = SemanticTokenModifiers(1 << 0);
98    /// `defaultLibrary` — built-in types and keywords are part of the language, not user code.
99    pub const DEFAULT_LIBRARY: SemanticTokenModifiers = SemanticTokenModifiers(1 << 1);
100
101    /// The modifier legend, in bit order; an editor declares this in its capabilities.
102    pub const LEGEND: &'static [&'static str] = &["documentation", "defaultLibrary"];
103
104    /// The raw bitset.
105    pub const fn bits(self) -> u32 {
106        self.0
107    }
108
109    /// Whether `other`'s bits are all set in `self`.
110    pub const fn contains(self, other: SemanticTokenModifiers) -> bool {
111        self.0 & other.0 == other.0
112    }
113}
114
115impl std::ops::BitOr for SemanticTokenModifiers {
116    type Output = SemanticTokenModifiers;
117    fn bitor(self, rhs: SemanticTokenModifiers) -> SemanticTokenModifiers {
118        SemanticTokenModifiers(self.0 | rhs.0)
119    }
120}
121
122/// The LSP token type + modifiers for a highlight kind, or `None` for kinds that carry no semantic
123/// token: whitespace (insignificant), punctuation (delimiters editors theme structurally), and lex
124/// errors (surfaced as diagnostics, not tokens).
125pub fn semantic_token(kind: HighlightKind) -> Option<(SemanticTokenType, SemanticTokenModifiers)> {
126    use HighlightKind::*;
127    use SemanticTokenModifiers as M;
128    Some(match kind {
129        Keyword => (SemanticTokenType::Keyword, M::DEFAULT_LIBRARY),
130        Type => (SemanticTokenType::Type, M::DEFAULT_LIBRARY),
131        Identifier | QuotedIdentifier => (SemanticTokenType::Variable, M::NONE),
132        String | DollarString => (SemanticTokenType::String, M::NONE),
133        Number => (SemanticTokenType::Number, M::NONE),
134        Variable => (SemanticTokenType::Parameter, M::NONE),
135        Operator => (SemanticTokenType::Operator, M::NONE),
136        Comment => (SemanticTokenType::Comment, M::DOCUMENTATION),
137        Whitespace | Punctuation | Error => return None,
138    })
139}
140
141/// An embedded-language region inside a `$$ … $$` body. Built by [`detect_injections`]; consumers
142/// can re-highlight the [`range`](Injection::range) with the embedded grammar named by
143/// [`language`](Injection::language). `#[non_exhaustive]` so fields can be added compatibly.
144#[derive(Clone, Debug, PartialEq, Eq)]
145#[non_exhaustive]
146pub struct Injection {
147    /// The embedded language declared (or inferred) for this body.
148    pub language: InjectedLanguage,
149    /// Byte range of the dollar-quoted body *including* its `$$` delimiters.
150    pub range: std::ops::Range<usize>,
151}
152
153/// A language that can be embedded in a Snowflake `$$ … $$` UDF / procedure body.
154#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
155pub enum InjectedLanguage {
156    Sql,
157    JavaScript,
158    Python,
159    Java,
160    Scala,
161}
162
163impl InjectedLanguage {
164    /// The Tree-sitter / TextMate-ish injection scope name editors use to pick a grammar.
165    pub const fn scope(self) -> &'static str {
166        match self {
167            InjectedLanguage::Sql => "source.snowflake-sql",
168            InjectedLanguage::JavaScript => "source.js",
169            InjectedLanguage::Python => "source.python",
170            InjectedLanguage::Java => "source.java",
171            InjectedLanguage::Scala => "source.scala",
172        }
173    }
174
175    /// Map a `LANGUAGE <name>` word (case-insensitive) to a language. SQL is the default body
176    /// language, so an unrecognized word falls back to [`InjectedLanguage::Sql`].
177    fn from_language_word(word: &str) -> InjectedLanguage {
178        if word.eq_ignore_ascii_case("javascript") {
179            InjectedLanguage::JavaScript
180        } else if word.eq_ignore_ascii_case("python") {
181            InjectedLanguage::Python
182        } else if word.eq_ignore_ascii_case("java") {
183            InjectedLanguage::Java
184        } else if word.eq_ignore_ascii_case("scala") {
185            InjectedLanguage::Scala
186        } else {
187            InjectedLanguage::Sql
188        }
189    }
190}
191
192/// A semantic token resolved to absolute byte coordinates (pre delta-encoding).
193/// `#[non_exhaustive]` so fields can be added without breaking downstream matches.
194#[derive(Clone, Debug, PartialEq, Eq)]
195#[non_exhaustive]
196pub struct ResolvedToken {
197    /// Byte range of the token in the source.
198    pub range: std::ops::Range<usize>,
199    /// The LSP semantic token type this token maps to.
200    pub token_type: SemanticTokenType,
201    /// The LSP semantic token modifiers (a bitset) for this token.
202    pub modifiers: SemanticTokenModifiers,
203}
204
205/// A semantic token as the LSP wire format wants it: zero-based line / UTF-16 char position with a
206/// UTF-16 length, plus the legend indices. These are *absolute* (not yet delta-encoded); call
207/// [`delta_encode`] for the on-wire `(deltaLine, deltaStartChar, …)` form. `#[non_exhaustive]` so
208/// fields can be added without breaking downstream matches.
209#[derive(Clone, Copy, Debug, PartialEq, Eq)]
210#[non_exhaustive]
211pub struct LineToken {
212    /// Zero-based line of the token's start.
213    pub line: u32,
214    /// Zero-based UTF-16 column of the token's start.
215    pub start_char: u32,
216    /// Length in UTF-16 code units (the LSP unit).
217    pub length: u32,
218    /// Legend index of the token type (see [`SemanticTokenType::index`]).
219    pub token_type: u32,
220    /// The token's modifier bitset (see [`SemanticTokenModifiers::bits`]).
221    pub modifiers: u32,
222}
223
224/// Scan `input` for embedded `$$ … $$` bodies and tag each with its language. Routine bodies use
225/// `LANGUAGE <name> ... AS $$...$$`; dynamic SQL uses `EXECUTE IMMEDIATE $$...$$`. Bodies after
226/// `AS $$...$$` with no `LANGUAGE` clause default to [`InjectedLanguage::Sql`]. Runs off the lexical
227/// highlighter, so it never parses or panics.
228pub fn detect_injections(input: &str) -> Vec<Injection> {
229    let highlighted = highlight(input);
230    let mut injections = Vec::new();
231    let mut language: Option<InjectedLanguage> = None;
232    let mut expect_language_name = false;
233    let mut saw_as_after_language = false;
234    let mut saw_as = false;
235    let mut saw_execute = false;
236    let mut saw_execute_immediate = false;
237
238    for token in &highlighted.tokens {
239        match token.kind {
240            HighlightKind::Whitespace | HighlightKind::Comment => {}
241            HighlightKind::DollarString => {
242                if saw_as_after_language || saw_as || saw_execute_immediate {
243                    injections.push(Injection {
244                        language: language.take().unwrap_or(InjectedLanguage::Sql),
245                        range: token.range.clone(),
246                    });
247                }
248                expect_language_name = false;
249                saw_as_after_language = false;
250                saw_as = false;
251                saw_execute = false;
252                saw_execute_immediate = false;
253            }
254            HighlightKind::Punctuation if token.text == ";" => {
255                // Statement boundary: a LANGUAGE clause does not carry across statements.
256                language = None;
257                expect_language_name = false;
258                saw_as_after_language = false;
259                saw_as = false;
260                saw_execute = false;
261                saw_execute_immediate = false;
262            }
263            HighlightKind::Keyword if token.text.eq_ignore_ascii_case("language") => {
264                expect_language_name = true;
265                saw_as_after_language = false;
266                saw_as = false;
267                saw_execute = false;
268                saw_execute_immediate = false;
269            }
270            HighlightKind::Keyword | HighlightKind::Identifier if expect_language_name => {
271                language = Some(InjectedLanguage::from_language_word(token.text));
272                expect_language_name = false;
273                saw_as_after_language = false;
274                saw_as = false;
275                saw_execute = false;
276                saw_execute_immediate = false;
277            }
278            HighlightKind::Keyword if token.text.eq_ignore_ascii_case("as") => {
279                saw_as_after_language = language.is_some();
280                saw_as = language.is_none();
281                expect_language_name = false;
282                saw_execute = false;
283                saw_execute_immediate = false;
284            }
285            HighlightKind::Keyword if token.text.eq_ignore_ascii_case("execute") => {
286                expect_language_name = false;
287                saw_as_after_language = false;
288                saw_as = false;
289                saw_execute = true;
290                saw_execute_immediate = false;
291            }
292            HighlightKind::Keyword
293                if token.text.eq_ignore_ascii_case("immediate") && saw_execute =>
294            {
295                expect_language_name = false;
296                saw_as_after_language = false;
297                saw_as = false;
298                saw_execute = false;
299                saw_execute_immediate = true;
300            }
301            _ => {
302                expect_language_name = false;
303                saw_as_after_language = false;
304                saw_as = false;
305                saw_execute = false;
306                saw_execute_immediate = false;
307            }
308        }
309    }
310    injections
311}
312
313/// Resolve `input` to absolute-coordinate semantic tokens (no delta-encoding). Trivia, punctuation
314/// and lex errors drop out; everything else becomes a [`ResolvedToken`] tagged per
315/// [`semantic_token`]. Injections do not change a token's *type* — the lexical kind already wins —
316/// but they are returned alongside by [`semantic_tokens`] so an editor can layer an embedded
317/// grammar over the body's range.
318pub fn resolve_tokens(input: &str) -> Vec<ResolvedToken> {
319    let highlighted = highlight(input);
320    highlighted
321        .tokens
322        .iter()
323        .enumerate()
324        .filter_map(|(index, token)| {
325            if is_cortex_or_aisql_function_token(&highlighted.tokens, index) {
326                return Some(ResolvedToken {
327                    range: token.range.clone(),
328                    token_type: SemanticTokenType::Function,
329                    modifiers: SemanticTokenModifiers::DEFAULT_LIBRARY,
330                });
331            }
332            let (token_type, modifiers) = semantic_token(token.kind)?;
333            Some(ResolvedToken {
334                range: token.range.clone(),
335                token_type,
336                modifiers,
337            })
338        })
339        .collect()
340}
341
342fn is_cortex_or_aisql_function_token(tokens: &[HighlightToken<'_>], index: usize) -> bool {
343    let token = &tokens[index];
344    if !matches!(
345        token.kind,
346        HighlightKind::Identifier | HighlightKind::Keyword
347    ) {
348        return false;
349    }
350    if next_significant(tokens, index).is_none_or(|next| tokens[next].text != "(") {
351        return false;
352    }
353    token.text.to_ascii_uppercase().starts_with("AI_")
354        || is_snowflake_cortex_qualified_leaf(tokens, index)
355}
356
357fn is_snowflake_cortex_qualified_leaf(tokens: &[HighlightToken<'_>], index: usize) -> bool {
358    let Some(dot_before_fn) = prev_significant(tokens, index) else {
359        return false;
360    };
361    if tokens[dot_before_fn].text != "." {
362        return false;
363    }
364    let Some(cortex) = prev_significant(tokens, dot_before_fn) else {
365        return false;
366    };
367    if !tokens[cortex].text.eq_ignore_ascii_case("cortex") {
368        return false;
369    }
370    let Some(dot_before_cortex) = prev_significant(tokens, cortex) else {
371        return false;
372    };
373    if tokens[dot_before_cortex].text != "." {
374        return false;
375    }
376    let Some(snowflake) = prev_significant(tokens, dot_before_cortex) else {
377        return false;
378    };
379    tokens[snowflake].text.eq_ignore_ascii_case("snowflake")
380}
381
382fn prev_significant(tokens: &[HighlightToken<'_>], index: usize) -> Option<usize> {
383    tokens[..index]
384        .iter()
385        .enumerate()
386        .rev()
387        .find(|(_, token)| {
388            !matches!(
389                token.kind,
390                HighlightKind::Whitespace | HighlightKind::Comment
391            )
392        })
393        .map(|(index, _)| index)
394}
395
396fn next_significant(tokens: &[HighlightToken<'_>], index: usize) -> Option<usize> {
397    tokens
398        .iter()
399        .enumerate()
400        .skip(index + 1)
401        .find(|(_, token)| {
402            !matches!(
403                token.kind,
404                HighlightKind::Whitespace | HighlightKind::Comment
405            )
406        })
407        .map(|(index, _)| index)
408}
409
410/// The full semantic-token result: the per-token tagging plus the embedded-language regions.
411/// `#[non_exhaustive]` so fields can be added without breaking downstream matches.
412#[derive(Clone, Debug, PartialEq, Eq)]
413#[non_exhaustive]
414pub struct SemanticTokens {
415    /// Every significant token, tagged with its LSP type + modifiers (trivia/punctuation dropped).
416    pub tokens: Vec<ResolvedToken>,
417    /// The `$$ … $$` embedded-language regions detected in the source.
418    pub injections: Vec<Injection>,
419}
420
421/// Resolve `input` to semantic tokens *and* its embedded-language injection regions in one pass.
422pub fn semantic_tokens(input: &str) -> SemanticTokens {
423    SemanticTokens {
424        tokens: resolve_tokens(input),
425        injections: detect_injections(input),
426    }
427}
428
429/// Lower resolved tokens to LSP [`LineToken`]s: split multi-line tokens (block comments,
430/// dollar-quoted bodies) into one token per line — the LSP encoding forbids a token spanning a
431/// newline — and compute UTF-16 positions and lengths. Output is sorted by `(line, start_char)`.
432pub fn line_tokens(input: &str) -> Vec<LineToken> {
433    let resolved = resolve_tokens(input);
434    let index = LineIndex::new(input);
435    let mut out = Vec::new();
436
437    for token in &resolved {
438        let mut piece_start = token.range.start;
439        for piece in input[token.range.clone()].split('\n') {
440            let length = utf16_len(piece);
441            if length > 0 {
442                let position = index.utf16_position(piece_start);
443                out.push(LineToken {
444                    line: position.line,
445                    start_char: position.character,
446                    length,
447                    token_type: token.token_type.index(),
448                    modifiers: token.modifiers.bits(),
449                });
450            }
451            piece_start += piece.len() + 1; // skip the piece and its trailing '\n'
452        }
453    }
454    out
455}
456
457/// Like [`line_tokens`], but columns and token lengths are measured in UTF-8 bytes.
458pub fn line_tokens_utf8(input: &str) -> Vec<LineToken> {
459    let resolved = resolve_tokens(input);
460    let index = LineIndex::new(input);
461    let mut out = Vec::new();
462
463    for token in &resolved {
464        let mut piece_start = token.range.start;
465        for piece in input[token.range.clone()].split('\n') {
466            let length = piece.len() as u32;
467            if length > 0 {
468                let position = index.utf8_position(piece_start);
469                out.push(LineToken {
470                    line: position.line,
471                    start_char: position.character,
472                    length,
473                    token_type: token.token_type.index(),
474                    modifiers: token.modifiers.bits(),
475                });
476            }
477            piece_start += piece.len() + 1; // skip the piece and its trailing '\n'
478        }
479    }
480    out
481}
482
483/// Delta-encode absolute [`LineToken`]s into the LSP wire form: each quintuple is
484/// `(deltaLine, deltaStartChar, length, tokenType, tokenModifiers)`, where deltas are relative to
485/// the previous token (and `deltaStartChar` resets to the absolute column when the line advances).
486/// Tokens must be sorted by `(line, start_char)`; [`line_tokens`] already produces them that way.
487pub fn delta_encode(tokens: &[LineToken]) -> Vec<[u32; 5]> {
488    let mut out = Vec::with_capacity(tokens.len());
489    let (mut prev_line, mut prev_col) = (0u32, 0u32);
490    for token in tokens {
491        let delta_line = token.line - prev_line;
492        let delta_start = if delta_line == 0 {
493            token.start_char - prev_col
494        } else {
495            token.start_char
496        };
497        out.push([
498            delta_line,
499            delta_start,
500            token.length,
501            token.token_type,
502            token.modifiers,
503        ]);
504        (prev_line, prev_col) = (token.line, token.start_char);
505    }
506    out
507}
508
509/// One-shot: resolve, lower to lines, and delta-encode. The vector an LSP server hands back for
510/// `textDocument/semanticTokens/full` (modulo the server's wrapper type).
511pub fn semantic_tokens_lsp(input: &str) -> Vec<[u32; 5]> {
512    delta_encode(&line_tokens(input))
513}
514
515/// One-shot semantic tokens using UTF-8 columns/lengths.
516pub fn semantic_tokens_lsp_utf8(input: &str) -> Vec<[u32; 5]> {
517    delta_encode(&line_tokens_utf8(input))
518}
519
520#[cfg(test)]
521mod tests {
522    use super::*;
523
524    #[test]
525    fn legend_indices_match_their_position() {
526        for (i, ty) in SemanticTokenType::LEGEND.iter().enumerate() {
527            assert_eq!(ty.index() as usize, i, "legend index drift for {ty:?}");
528        }
529        // Names are the canonical LSP spellings.
530        assert_eq!(SemanticTokenType::Keyword.name(), "keyword");
531        assert_eq!(SemanticTokenType::Parameter.name(), "parameter");
532        assert_eq!(SemanticTokenType::Namespace.name(), "namespace");
533    }
534
535    #[test]
536    fn modifier_bitset_is_a_real_bitset() {
537        let both = SemanticTokenModifiers::DOCUMENTATION | SemanticTokenModifiers::DEFAULT_LIBRARY;
538        assert_eq!(both.bits(), 0b11);
539        assert!(both.contains(SemanticTokenModifiers::DOCUMENTATION));
540        assert!(both.contains(SemanticTokenModifiers::DEFAULT_LIBRARY));
541        assert!(!SemanticTokenModifiers::NONE.contains(SemanticTokenModifiers::DOCUMENTATION));
542        assert_eq!(SemanticTokenModifiers::LEGEND.len(), 2);
543    }
544
545    #[test]
546    fn every_highlight_kind_maps_consistently() {
547        use HighlightKind::*;
548        // Significant kinds carry a token; insignificant ones do not.
549        for kind in [
550            Keyword,
551            Type,
552            Identifier,
553            QuotedIdentifier,
554            String,
555            DollarString,
556            Number,
557            Variable,
558            Operator,
559            Comment,
560        ] {
561            assert!(semantic_token(kind).is_some(), "{kind:?} should map");
562        }
563        for kind in [Whitespace, Punctuation, Error] {
564            assert!(semantic_token(kind).is_none(), "{kind:?} should not map");
565        }
566        // Keywords and types are part of the language.
567        assert_eq!(
568            semantic_token(Keyword).unwrap().1,
569            SemanticTokenModifiers::DEFAULT_LIBRARY
570        );
571        assert_eq!(
572            semantic_token(Type).unwrap().1,
573            SemanticTokenModifiers::DEFAULT_LIBRARY
574        );
575        // Comments are documentation.
576        assert_eq!(
577            semantic_token(Comment).unwrap().1,
578            SemanticTokenModifiers::DOCUMENTATION
579        );
580    }
581
582    #[test]
583    fn resolve_drops_trivia_and_punctuation() {
584        let toks = resolve_tokens("SELECT a, 1 -- c\n");
585        let types: Vec<_> = toks.iter().map(|t| t.token_type).collect();
586        // SELECT(kw) a(var) 1(num) -- c(comment); the comma is punctuation and drops.
587        assert_eq!(
588            types,
589            vec![
590                SemanticTokenType::Keyword,
591                SemanticTokenType::Variable,
592                SemanticTokenType::Number,
593                SemanticTokenType::Comment,
594            ]
595        );
596    }
597
598    #[test]
599    fn resolved_ranges_match_source_text() {
600        let sql = "SELECT $1::NUMBER FROM t";
601        for tok in resolve_tokens(sql) {
602            assert!(tok.range.end <= sql.len());
603        }
604        // $1 is a parameter, NUMBER a type, :: an operator.
605        let sem = semantic_tokens(sql);
606        let by_text: Vec<_> = sem
607            .tokens
608            .iter()
609            .map(|t| (&sql[t.range.clone()], t.token_type))
610            .collect();
611        assert!(by_text.contains(&("$1", SemanticTokenType::Parameter)));
612        assert!(by_text.contains(&("NUMBER", SemanticTokenType::Type)));
613        assert!(by_text.contains(&("::", SemanticTokenType::Operator)));
614    }
615
616    #[test]
617    fn line_tokens_are_utf16_and_split_multiline() {
618        // 芋 is 3 bytes but 1 UTF-16 unit; the block comment spans two lines.
619        let sql = "SELECT '芋' /* a\nb */ x";
620        let lines = line_tokens(sql);
621        // The string token length is in UTF-16 units: '芋' => quote+芋+quote = 3 units.
622        let string_tok = lines
623            .iter()
624            .find(|t| t.token_type == SemanticTokenType::String.index())
625            .unwrap();
626        assert_eq!(string_tok.length, 3);
627        // The block comment is split into two LineTokens on consecutive lines.
628        let comment_lines: Vec<_> = lines
629            .iter()
630            .filter(|t| t.token_type == SemanticTokenType::Comment.index())
631            .map(|t| t.line)
632            .collect();
633        assert_eq!(comment_lines, vec![0, 1]);
634    }
635
636    #[test]
637    fn delta_encode_resets_column_on_newline() {
638        let sql = "SELECT a\nFROM t";
639        let encoded = semantic_tokens_lsp(sql);
640        // SELECT@(0,0) a@(0,7) FROM@(1,0) t@(1,5)
641        assert_eq!(
642            encoded[0],
643            [0, 0, 6, 0, SemanticTokenModifiers::DEFAULT_LIBRARY.bits()]
644        );
645        // a: same line, +7 from SELECT start; variable type=2, no modifiers.
646        assert_eq!(encoded[1], [0, 7, 1, 2, 0]);
647        // FROM: next line, absolute column 0; keyword.
648        assert_eq!(
649            encoded[2],
650            [1, 0, 4, 0, SemanticTokenModifiers::DEFAULT_LIBRARY.bits()]
651        );
652        // t: same line, +5.
653        assert_eq!(encoded[3], [0, 5, 1, 2, 0]);
654    }
655
656    #[test]
657    fn empty_input_yields_no_tokens() {
658        assert!(line_tokens("").is_empty());
659        assert!(semantic_tokens_lsp("").is_empty());
660        assert!(detect_injections("").is_empty());
661        assert!(semantic_tokens("").tokens.is_empty());
662    }
663
664    #[test]
665    fn detects_javascript_injection_from_language_clause() {
666        let sql = "CREATE FUNCTION f() RETURNS STRING LANGUAGE JAVASCRIPT AS $$ return 1; $$;";
667        let injections = detect_injections(sql);
668        assert_eq!(injections.len(), 1);
669        assert_eq!(injections[0].language, InjectedLanguage::JavaScript);
670        // The range covers the whole $$ body including delimiters.
671        let body = &sql[injections[0].range.clone()];
672        assert!(body.starts_with("$$"));
673        assert!(body.ends_with("$$"));
674        assert!(body.contains("return 1;"));
675    }
676
677    #[test]
678    fn detects_python_and_scala_and_java_injections() {
679        for (word, expected) in [
680            ("PYTHON", InjectedLanguage::Python),
681            ("Java", InjectedLanguage::Java),
682            ("scala", InjectedLanguage::Scala),
683            ("SQL", InjectedLanguage::Sql),
684        ] {
685            let sql = format!("CREATE FUNCTION f() RETURNS INT LANGUAGE {word} AS $$x$$;");
686            let injections = detect_injections(&sql);
687            assert_eq!(injections.len(), 1, "for {word}");
688            assert_eq!(injections[0].language, expected, "for {word}");
689        }
690    }
691
692    #[test]
693    fn body_without_language_clause_defaults_to_sql() {
694        let sql = "CREATE PROCEDURE p() RETURNS STRING AS $$ BEGIN RETURN 'ok'; END $$; \
695                   EXECUTE IMMEDIATE $$ SELECT 1 $$;";
696        let injections = detect_injections(sql);
697        assert_eq!(injections.len(), 2);
698        assert_eq!(injections[0].language, InjectedLanguage::Sql);
699        assert_eq!(injections[1].language, InjectedLanguage::Sql);
700        assert_eq!(InjectedLanguage::Sql.scope(), "source.snowflake-sql");
701        assert_eq!(InjectedLanguage::JavaScript.scope(), "source.js");
702    }
703
704    #[test]
705    fn non_body_dollar_strings_do_not_get_injected() {
706        let sql = "SELECT $$plain text$$ AS value; \
707                   CREATE PROCEDURE p() LANGUAGE PYTHON IMPORTS = ($$stage/file.py$$) AS $$body$$;";
708        let injections = detect_injections(sql);
709        assert_eq!(injections.len(), 1);
710        assert_eq!(injections[0].language, InjectedLanguage::Python);
711        assert_eq!(&sql[injections[0].range.clone()], "$$body$$");
712    }
713
714    #[test]
715    fn language_clause_does_not_leak_across_statements() {
716        // The first statement is JS; the second has a bare body and must default to SQL.
717        let sql = "CREATE FUNCTION a() RETURNS INT LANGUAGE JAVASCRIPT AS $$1$$; \
718                   EXECUTE IMMEDIATE $$ SELECT 2 $$;";
719        let injections = detect_injections(sql);
720        assert_eq!(injections.len(), 2);
721        assert_eq!(injections[0].language, InjectedLanguage::JavaScript);
722        assert_eq!(injections[1].language, InjectedLanguage::Sql);
723    }
724
725    #[test]
726    fn multiple_bodies_each_get_their_language() {
727        let sql = "CREATE FUNCTION a() LANGUAGE PYTHON AS $$py$$; \
728                   CREATE FUNCTION b() LANGUAGE JAVASCRIPT AS $$js$$;";
729        let injections = detect_injections(sql);
730        assert_eq!(injections.len(), 2);
731        assert_eq!(injections[0].language, InjectedLanguage::Python);
732        assert_eq!(injections[1].language, InjectedLanguage::JavaScript);
733    }
734
735    #[test]
736    fn semantic_tokens_bundles_tokens_and_injections() {
737        let sql = "CREATE FUNCTION f() LANGUAGE JAVASCRIPT AS $$ return 1; $$";
738        let sem = semantic_tokens(sql);
739        assert_eq!(sem.injections.len(), 1);
740        assert_eq!(sem.injections[0].language, InjectedLanguage::JavaScript);
741        // The dollar body is itself a String semantic token covering the injection range.
742        let dollar = sem
743            .tokens
744            .iter()
745            .find(|t| {
746                t.token_type == SemanticTokenType::String && sql[t.range.clone()].contains("$$")
747            })
748            .expect("dollar string token");
749        assert_eq!(dollar.range, sem.injections[0].range);
750    }
751
752    #[test]
753    fn never_panics_on_adversarial_input() {
754        for sql in [
755            "$$",
756            "$$ unterminated",
757            "LANGUAGE",
758            "LANGUAGE $$x$$",
759            ";;;;",
760            "@~/stage/ $1 :name ? ->> |> => :: ->",
761            "'長芋' \"畑\" -- 芋\n/* 芋 */ $$芋$$",
762            "\r\n\r\n",
763        ] {
764            // Each entry point must complete without panicking and stay self-consistent.
765            let lines = line_tokens(sql);
766            let encoded = delta_encode(&lines);
767            assert_eq!(lines.len(), encoded.len());
768            let _ = semantic_tokens(sql);
769            let _ = semantic_tokens_lsp(sql);
770        }
771    }
772}
sql_dialect_fmt_highlight/semantic.rs

sql_dialect_fmt_highlight/
semantic.rs
