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

sql_dialect_fmt_highlight/
semantic.rs
