Skip to main content

rudb_parse/
tokenize.rs

1//! The tokenizer, matched to DuckDB's by behaviour.
2//!
3//! This is the one part of the front end with no declarative artifact behind it. The grammar is
4//! vendored and generated from, per `spec/20-the-grammar.md`, and it says nothing about string
5//! literals, dollar quoting, numeric literal forms, comments or the operator rules. All of that
6//! is 613 lines of hand written C++ upstream, in `src/parser/peg/tokenizer/base_tokenizer.cpp`
7//! and `parser_tokenizer.cpp`, and this is a port of it rather than an interpretation.
8//!
9//! Which is worth saying plainly: everything here is a fact about DuckDB and not a design
10//! decision of ours. Where the behaviour looks wrong, it is still the behaviour, because a query
11//! that returns a different answer is worse than a query that returns a surprising one. Section
12//! 20.7 enumerates the ten that bite. The reason a differential fuzzer against a real DuckDB is
13//! scheduled from week one rather than at M5 is this file.
14//!
15//! One deliberate divergence, and it is representational. Upstream types a quoted identifier and
16//! a bare one both as `IDENTIFIER` and recovers the difference from the first byte where it
17//! matters. We give them separate kinds. Nothing downstream may treat them differently in a place
18//! upstream does not.
19//!
20//! Nothing is decoded. A string keeps its quotes and its escapes, a number keeps its underscores,
21//! an identifier keeps its case. `spec/20-the-grammar.md` section 7 for why case in particular:
22//! DuckDB is case insensitive and case preserving, including for quoted identifiers, which is not
23//! what PostgreSQL does and not what folding here would give.
24
25use rudb_common::{Error, Result, Span};
26
27use crate::generated::keywords::{KEYWORDS, LONGEST};
28use crate::token::{Flags, Kind, NOT_A_KEYWORD, Token};
29
30/// Split `query` into tokens, ending with exactly one [`Kind::EndOfInput`].
31///
32/// Fails only where DuckDB's parser tokenizer throws, which is four cases: an unterminated block
33/// comment, an unterminated string literal, an unterminated quoted identifier, and the empty
34/// quoted identifier `""`. An unterminated dollar quoted string is not one of them and comes back
35/// as a token with [`Flags::UNTERMINATED`] set.
36pub fn tokenize(query: &str) -> Result<Vec<Token>> {
37    Tokenizer::new(query).run()
38}
39
40/// Where the scan is.
41///
42/// Named for what upstream calls them so that reading the two side by side stays possible, which
43/// matters more here than anywhere else in the crate.
44#[derive(Clone, Copy, PartialEq, Eq)]
45enum State {
46    Standard,
47    LineComment,
48    BlockComment,
49    QuotedIdentifier,
50    StringLiteral,
51    Word,
52    Numeric,
53    Operator,
54    DollarQuoted,
55}
56
57struct Tokenizer<'a> {
58    query: &'a str,
59    bytes: &'a [u8],
60    tokens: Vec<Token>,
61    /// Where the token being built starts, and after a token is pushed, where the gap starts.
62    last: usize,
63    /// A block comment ended here. Used to set [`Flags::BLOCK_COMMENT`] on the next token, which
64    /// is the whole reason comments are tracked rather than skipped.
65    block_comment_at: Option<usize>,
66    /// Set by an `E` prefix, which is the only one of the four that changes how the body is read.
67    escape_string: bool,
68    /// The tag between the dollars, as a span, so that closing it is a slice comparison.
69    dollar_tag: Span,
70    depth: u32,
71}
72
73impl<'a> Tokenizer<'a> {
74    fn new(query: &'a str) -> Self {
75        Tokenizer {
76            query,
77            bytes: query.as_bytes(),
78            // Two tokens every seven bytes is what the ClickBench and TPC-H queries come out at.
79            // Being wrong here costs a realloc, being absent costs six.
80            tokens: Vec::with_capacity(query.len() / 4 + 4),
81            last: 0,
82            block_comment_at: None,
83            escape_string: false,
84            dollar_tag: Span::new(0, 0),
85            depth: 0,
86        }
87    }
88
89    fn run(mut self) -> Result<Vec<Token>> {
90        let mut state = State::Standard;
91        let mut i = 0;
92        while i < self.bytes.len() {
93            let c = self.bytes[i];
94            match state {
95                State::Standard => {
96                    if let Some(next) = self.standard(&mut i, c)? {
97                        state = next;
98                    }
99                }
100                State::Numeric => self.numeric(&mut state, &mut i, c),
101                State::Operator => self.operator(&mut state, &mut i, c),
102                State::Word => self.word(&mut state, &mut i, c),
103                State::StringLiteral => self.string_literal(&mut state, &mut i, c),
104                State::QuotedIdentifier => self.quoted_identifier(&mut state, &mut i, c)?,
105                State::LineComment => {
106                    if c == b'\n' || c == b'\r' {
107                        self.comment(self.last, i + 1);
108                        self.last = i + 1;
109                        state = State::Standard;
110                    }
111                }
112                State::BlockComment => self.block_comment(&mut state, &mut i, c),
113                State::DollarQuoted => self.dollar_quoted(&mut state, &mut i),
114            }
115            i += 1;
116        }
117        self.finish(state)
118    }
119
120    /// The end of the input, which is a different decision per state and not a loop exit.
121    fn finish(mut self, state: State) -> Result<Vec<Token>> {
122        let end = self.bytes.len();
123        match state {
124            State::LineComment => {
125                self.comment(self.last, end);
126            }
127            State::BlockComment => {
128                return Err(self.error(
129                    format!(
130                        "unterminated /* comment at or near \"{}\"",
131                        &self.query[self.last..end]
132                    ),
133                    self.last,
134                ));
135            }
136            State::Operator => self.push_operator(self.last, end),
137            State::DollarQuoted => {
138                // Not an error upstream, which is worth noticing rather than tidying up. It comes
139                // back as a token that ran off the end and the grammar gets to decide.
140                self.push_flagged(self.last, end, Kind::String, Flags::UNTERMINATED);
141            }
142            State::StringLiteral => {
143                return Err(self.error("unterminated string literal", self.last));
144            }
145            State::QuotedIdentifier => {
146                return Err(self.error("unterminated quoted identifier", self.last));
147            }
148            State::Numeric => self.push(self.last, end, Kind::Number),
149            State::Word => self.push_word(self.last, end),
150            // A `$` with nothing after it lands here, and upstream calls it an identifier. It is
151            // reproduced rather than corrected, because a tokenizer that disagrees with DuckDB
152            // about a one byte query disagrees with it about something.
153            State::Standard => self.push(self.last, end, Kind::Identifier),
154        }
155        self.tokens.push(Token {
156            kind: Kind::EndOfInput,
157            flags: Flags::default(),
158            keyword: NOT_A_KEYWORD,
159            start: end as u32,
160            end: end as u32,
161        });
162        Ok(self.tokens)
163    }
164
165    /// The dispatch at the top of a token. Returns the state to move to, if it changes.
166    fn standard(&mut self, i: &mut usize, c: u8) -> Result<Option<State>> {
167        match c {
168            b'\'' => {
169                self.last = *i;
170                self.escape_string = false;
171                return Ok(Some(State::StringLiteral));
172            }
173            b'"' => {
174                self.last = *i;
175                return Ok(Some(State::QuotedIdentifier));
176            }
177            b';' => {
178                // The base tokenizer emits nothing here and lets its caller decide. The parser's
179                // caller emits `;`, because `Program <- TopLevelStatement*` consumes it. The
180                // autocomplete caller does something else, which is why the hook exists.
181                self.tokens.push(Token {
182                    kind: Kind::Terminator,
183                    flags: self.gap_flags(*i),
184                    keyword: NOT_A_KEYWORD,
185                    start: *i as u32,
186                    end: *i as u32 + 1,
187                });
188                self.last = *i + 1;
189                return Ok(None);
190            }
191            b'$' => return Ok(self.dollar(i)),
192            b'-' if self.bytes.get(*i + 1) == Some(&b'-') => {
193                *i += 1;
194                return Ok(Some(State::LineComment));
195            }
196            b'/' if self.bytes.get(*i + 1) == Some(&b'*') => {
197                *i += 1;
198                self.depth = 1;
199                return Ok(Some(State::BlockComment));
200            }
201            _ => {}
202        }
203
204        if is_space(c) {
205            self.last = *i + 1;
206            return Ok(None);
207        }
208
209        if let Some(len) = special_operator(self.bytes, *i) {
210            // `::=` is an operator run and not `::` followed by `=`. The three character check
211            // above is why `->` survives the rule that a `-` never joins anything.
212            if self.bytes.get(*i + len).is_some_and(|&next| is_operator_char_in_run(next)) {
213                self.last = *i;
214                return Ok(Some(State::Operator));
215            }
216            self.push(*i, *i + len, Kind::Operator);
217            *i += len - 1;
218            self.last = *i + 1;
219            return Ok(None);
220        }
221
222        if is_single_byte_operator(c) {
223            self.push(*i, *i + 1, Kind::Operator);
224            self.last = *i + 1;
225            return Ok(None);
226        }
227
228        if is_initial_number(c) {
229            self.last = *i;
230            return Ok(Some(State::Numeric));
231        }
232
233        // `E`, `X`, `B` and `N`, in either case, and only when the quote is the very next byte.
234        // `SELECT e 'a'` is an identifier and then a string. Only `E` changes how the body is
235        // read, and it is the only one of the four this tokenizer does anything else with.
236        if is_string_prefix(c) && self.bytes.get(*i + 1) == Some(&b'\'') {
237            self.last = *i;
238            self.escape_string = c == b'E' || c == b'e';
239            *i += 1;
240            return Ok(Some(State::StringLiteral));
241        }
242
243        if is_operator_char(c) {
244            self.last = *i;
245            return Ok(Some(State::Operator));
246        }
247
248        self.last = *i;
249        Ok(Some(State::Word))
250    }
251
252    /// `$` is three things and which one it is depends on what follows.
253    fn dollar(&mut self, i: &mut usize) -> Option<State> {
254        let Some(&next) = self.bytes.get(*i + 1) else {
255            // Nothing after it, and upstream leaves `last` where it is, so this byte ends up in
256            // whatever the final token turns out to be.
257            return None;
258        };
259        if next.is_ascii_digit() {
260            // `$1` is a parameter, and it is two tokens rather than one. The grammar spells it,
261            // which is why the tokenizer does not have to.
262            self.push(*i, *i + 1, Kind::Operator);
263            return None;
264        }
265
266        // A tag runs to the next `$` and may contain only tag characters. Anything else and this
267        // was `$name`, which is also two tokens.
268        let mut close = None;
269        for at in *i + 1..self.bytes.len() {
270            if self.bytes[at] == b'$' {
271                close = Some(at);
272                break;
273            }
274            if !is_dollar_tag_char(self.bytes[at]) {
275                break;
276            }
277        }
278        let Some(close) = close else {
279            self.push(*i, *i + 1, Kind::Operator);
280            return None;
281        };
282
283        self.last = *i;
284        self.dollar_tag = Span::new(*i as u32 + 1, close as u32);
285        *i = close;
286        Some(State::DollarQuoted)
287    }
288
289    fn numeric(&mut self, state: &mut State, i: &mut usize, c: u8) {
290        if is_initial_number(c) {
291            return;
292        }
293        // Only between two digits, which is why `1_000` is a thousand and `SELECT 1_` is `1`
294        // aliased `_`.
295        if c == b'_' && self.bytes.get(*i + 1).is_some_and(|&n| is_initial_number(n)) {
296            return;
297        }
298        if is_scientific(c) && !is_scientific(self.bytes[*i - 1]) {
299            // A digit has to be in there somewhere, which rules out `.e100` while allowing both
300            // `1e5` and `.1e5`.
301            if self.bytes[self.last].is_ascii_digit() || self.bytes[*i - 1].is_ascii_digit() {
302                return;
303            }
304        }
305        if (c == b'+' || c == b'-') && is_scientific(self.bytes[*i - 1]) {
306            return;
307        }
308
309        // Give back anything on the end that is not a digit or a dot, which is what turns the `e`
310        // of `1e+` back into something the next pass has to deal with.
311        while !is_initial_number(self.bytes[*i - 1]) {
312            *i -= 1;
313        }
314        self.push(self.last, *i, Kind::Number);
315        *state = State::Standard;
316        self.last = *i;
317        *i -= 1;
318    }
319
320    fn operator(&mut self, state: &mut State, i: &mut usize, c: u8) {
321        if c == b'/' && self.bytes.get(*i + 1) == Some(&b'*') {
322            self.push_operator(self.last, *i);
323            *state = State::Standard;
324            self.last = *i;
325            *i -= 1;
326            return;
327        }
328        if !is_operator_char_in_run(c) {
329            self.push_operator(self.last, *i);
330            *state = State::Standard;
331            self.last = *i;
332            *i -= 1;
333        }
334    }
335
336    fn word(&mut self, state: &mut State, i: &mut usize, c: u8) {
337        // `$` is a legal non-initial identifier character, which is PostgreSQL's rule and is why
338        // this one test is not part of `is_word_char`.
339        if c == b'$' || is_word_char(c) {
340            return;
341        }
342        self.push_word(self.last, *i);
343        *state = State::Standard;
344        self.last = *i;
345        *i -= 1;
346    }
347
348    fn string_literal(&mut self, state: &mut State, i: &mut usize, c: u8) {
349        if self.escape_string && c == b'\\' && *i + 1 < self.bytes.len() {
350            *i += 1;
351            return;
352        }
353        if c != b'\'' {
354            return;
355        }
356        if self.bytes.get(*i + 1) == Some(&b'\'') {
357            *i += 1;
358            return;
359        }
360        self.push(self.last, *i + 1, Kind::String);
361        self.last = *i + 1;
362        self.escape_string = false;
363        *state = State::Standard;
364    }
365
366    fn quoted_identifier(&mut self, state: &mut State, i: &mut usize, c: u8) -> Result<()> {
367        if c != b'"' {
368            return Ok(());
369        }
370        if self.bytes.get(*i + 1) == Some(&b'"') {
371            *i += 1;
372            return Ok(());
373        }
374        if *i + 1 == self.last + 2 {
375            return Err(self.error("zero-length delimited identifier", self.last));
376        }
377        self.push(self.last, *i + 1, Kind::QuotedIdentifier);
378        self.last = *i + 1;
379        *state = State::Standard;
380        Ok(())
381    }
382
383    fn block_comment(&mut self, state: &mut State, i: &mut usize, c: u8) {
384        // Nested, which is the difference between commenting out a block that contains a comment
385        // and getting a syntax error halfway down the file.
386        if c == b'/' && self.bytes.get(*i + 1) == Some(&b'*') {
387            *i += 1;
388            self.depth += 1;
389        } else if c == b'*' && self.bytes.get(*i + 1) == Some(&b'/') {
390            *i += 1;
391            self.depth -= 1;
392            if self.depth == 0 {
393                self.comment(self.last, *i + 1);
394                self.last = *i + 1;
395                *state = State::Standard;
396            }
397        }
398    }
399
400    fn dollar_quoted(&mut self, state: &mut State, i: &mut usize) {
401        if self.bytes[*i] != b'$' || *i + 1 >= self.bytes.len() {
402            return;
403        }
404        let start = *i + 1;
405        let mut end = start;
406        while end < self.bytes.len() && self.bytes[end] != b'$' {
407            end += 1;
408        }
409        if end >= self.bytes.len() {
410            return;
411        }
412        let tag = &self.bytes[self.dollar_tag.start as usize..self.dollar_tag.end as usize];
413        if end - start != tag.len() || &self.bytes[start..end] != tag {
414            return;
415        }
416        self.push(self.last, end + 1, Kind::String);
417        *state = State::Standard;
418        *i = end;
419        self.last = *i + 1;
420    }
421
422    /// Push a bare word, having decided whether it is a keyword.
423    ///
424    /// A word is a keyword when it is in at least one class, which is not the same as being in
425    /// the table. The 15 soft words are in the table with a mask of zero and are identifiers here,
426    /// exactly as `PEGKeywordHelper::IsKeyword` has it, because their lists are the five class
427    /// lists and a soft word is in none of them. They still keep their index, so `ORDER BY x
428    /// ASCENDING` can match the literal without `SELECT ascending FROM t` becoming a syntax error.
429    /// `spec/20-the-grammar.md` section 5.
430    fn push_word(&mut self, start: usize, end: usize) {
431        if start >= end {
432            return;
433        }
434        let keyword = lookup(&self.query[start..end]);
435        let kind = if classes(keyword) == 0 { Kind::Identifier } else { Kind::Keyword };
436        let flags = self.gap_flags(start);
437        self.tokens.push(Token { kind, flags, keyword, start: start as u32, end: end as u32 });
438    }
439
440    /// An operator run, minus a trailing `+` where PostgreSQL says to give it back.
441    ///
442    /// `SELECT 1 =+ 1` is `1 = +1`, and `SELECT 1 !=+ 1` goes looking for an operator named `!=+`,
443    /// because the run in the second case contains a character from the special set. The rule is
444    /// PostgreSQL's and it exists so that a user defined operator can be told apart from an
445    /// operator followed by a signed number.
446    fn push_operator(&mut self, start: usize, end: usize) {
447        let special = self.bytes[start..end].iter().any(|&b| {
448            matches!(b, b'~' | b'!' | b'@' | b'#' | b'%' | b'^' | b'&' | b'|' | b'`' | b'?')
449        });
450        let mut cut = end;
451        if !special {
452            while cut > start && self.bytes[cut - 1] == b'+' {
453                cut -= 1;
454            }
455        }
456        self.push(start, cut, Kind::Operator);
457        for at in cut..end {
458            self.push(at, at + 1, Kind::Operator);
459        }
460    }
461
462    /// Record a comment. Nothing is pushed, because a comment is not a token anywhere the parser
463    /// can see, but where the block ones were has to be remembered so the next token can say it
464    /// was preceded by one.
465    fn comment(&mut self, start: usize, end: usize) {
466        if end >= start + 2 && &self.bytes[start..start + 2] == b"/*" {
467            self.block_comment_at = Some(start);
468        }
469    }
470
471    /// Push a token that is not a bare word, unless it is empty.
472    ///
473    /// The empty check is upstream's and it is what lets several states push unconditionally at a
474    /// boundary without first asking whether they have anything.
475    ///
476    /// One divergence, and it is in a field rather than in a token. Upstream reaches around
477    /// `PushToken` for single byte operators, special operators, a trimmed `+`, a `$` and a `;`,
478    /// so those five arrive with both gap flags clear no matter what was in the gap. We set them
479    /// on everything. Nothing reads a gap flag on an operator today, and a rule that holds for
480    /// every token is one fewer thing to remember when something starts to.
481    fn push(&mut self, start: usize, end: usize, kind: Kind) {
482        if start >= end {
483            return;
484        }
485        let flags = self.gap_flags(start);
486        self.tokens.push(Token {
487            kind,
488            flags,
489            keyword: NOT_A_KEYWORD,
490            start: start as u32,
491            end: end as u32,
492        });
493    }
494
495    fn push_flagged(&mut self, start: usize, end: usize, kind: Kind, extra: Flags) {
496        self.push(start, end, kind);
497        if let Some(token) = self.tokens.last_mut() {
498            token.flags = token.flags.with(extra);
499        }
500    }
501
502    /// What was in the gap between the previous token and `start`.
503    ///
504    /// Two literals separated by whitespace containing a newline are one literal, a line comment
505    /// between them keeps the join, and a block comment breaks it. That rule is PostgreSQL's,
506    /// DuckDB kept it, and it is the only reason either flag exists.
507    fn gap_flags(&self, start: usize) -> Flags {
508        let Some(previous) = self.tokens.last() else { return Flags::default() };
509        let from = previous.end as usize;
510        let mut flags = Flags::default();
511        if self.block_comment_at.is_some_and(|at| at >= from && at < start) {
512            flags = flags.with(Flags::BLOCK_COMMENT);
513        }
514        if self.bytes[from..start.min(self.bytes.len())].iter().any(|&b| b == b'\n' || b == b'\r') {
515            flags = flags.with(Flags::NEWLINE);
516        }
517        flags
518    }
519
520    fn error(&self, message: impl Into<String>, at: usize) -> Error {
521        Error::parser(message).with_span(Span::new(at as u32, self.bytes.len() as u32))
522    }
523}
524
525/// The index of `word` in the generated keyword table, or [`NOT_A_KEYWORD`].
526///
527/// ASCII folded into a fixed buffer, because every keyword is a plain ASCII word and the longest
528/// is 15 bytes, so a longer candidate cannot be one and never touches the table.
529pub fn lookup(word: &str) -> u16 {
530    if word.len() > LONGEST {
531        return NOT_A_KEYWORD;
532    }
533    let mut folded = [0u8; LONGEST];
534    for (slot, byte) in folded.iter_mut().zip(word.bytes()) {
535        *slot = byte.to_ascii_lowercase();
536    }
537    let folded = &folded[..word.len()];
538    // The table is sorted, so this is one binary search over 514 entries, which is nine
539    // comparisons and fits in a handful of cache lines.
540    match KEYWORDS.binary_search_by(|(candidate, _)| candidate.as_bytes().cmp(folded)) {
541        Ok(at) => at as u16,
542        Err(_) => NOT_A_KEYWORD,
543    }
544}
545
546/// The classes `word` belongs to, or zero.
547///
548/// Zero for a word that is in the table with no class, which is one of the 15 soft words, and
549/// zero for a word that is not in the table at all. Those two are the same answer to the only
550/// question this function is asked, which is whether the word blocks an identifier here.
551pub fn classes(keyword: u16) -> u8 {
552    if keyword == NOT_A_KEYWORD { 0 } else { KEYWORDS[keyword as usize].1 }
553}
554
555const fn is_space(c: u8) -> bool {
556    matches!(c, b' ' | b'\t' | b'\n' | 0x0b | 0x0c | b'\r')
557}
558
559/// The dozen characters that are always their own token and never join a run.
560///
561/// `-` and `#` being in here is the surprise. It is why `SELECT 1 =- 1` is `1 = -1` and why `--`
562/// can be a comment without ever having to be told it is not an operator.
563const fn is_single_byte_operator(c: u8) -> bool {
564    matches!(c, b'(' | b')' | b'{' | b'}' | b'[' | b']' | b',' | b'?' | b'$' | b'-' | b'#')
565}
566
567/// PostgreSQL's operator character set, which is every ASCII punctuation character except `_`.
568const fn is_operator_char(c: u8) -> bool {
569    if c == b'_' {
570        return false;
571    }
572    matches!(c, b'!'..=b'/' | b':'..=b'@' | b'['..=b'`' | b'{'..=b'~')
573}
574
575/// Whether `c` can continue an operator run.
576///
577/// The single byte operators are excluded, and so are the five characters that always end one:
578/// `'`, `-`, `;`, `"` and `.`. `-` is in both lists and that is not redundant, since the first
579/// list is also consulted where the second is not.
580const fn is_operator_char_in_run(c: u8) -> bool {
581    if is_single_byte_operator(c) || is_control_flow(c) {
582        return false;
583    }
584    is_operator_char(c)
585}
586
587const fn is_control_flow(c: u8) -> bool {
588    matches!(c, b'\'' | b'-' | b';' | b'"' | b'.')
589}
590
591/// Whether `c` can continue a bare word.
592///
593/// Anything that is not punctuation, whitespace or a delimiter, which by the arithmetic above
594/// includes every byte at or above 0x80. So an identifier may be any UTF-8 the user likes, and
595/// the tokenizer never has to decode it to find that out.
596const fn is_word_char(c: u8) -> bool {
597    if is_single_byte_operator(c) || is_operator_char(c) || is_space(c) || is_control_flow(c) {
598        return false;
599    }
600    true
601}
602
603/// A digit or a dot, which is both what starts a number and what can appear anywhere in one.
604///
605/// A dot being unconditional here is why `SELECT 1.2.3` is a single number token rather than
606/// three things. What it means is somebody else's problem, which is exactly how upstream has it.
607const fn is_initial_number(c: u8) -> bool {
608    c.is_ascii_digit() || c == b'.'
609}
610
611const fn is_scientific(c: u8) -> bool {
612    c == b'e' || c == b'E'
613}
614
615const fn is_string_prefix(c: u8) -> bool {
616    matches!(c, b'N' | b'n' | b'X' | b'x' | b'E' | b'e' | b'B' | b'b')
617}
618
619/// A-Z, a-z, 0-9, `_`, and anything at or above 0x80. Digits are only legal after the first byte
620/// and the caller checks that.
621const fn is_dollar_tag_char(c: u8) -> bool {
622    c.is_ascii_alphanumeric() || c == b'_' || c >= 0x80
623}
624
625/// The six sequences checked before the maximal run, longest first.
626///
627/// `->` is here because `-` is a single byte operator and would otherwise never join anything,
628/// which would leave the JSON arrow unspellable.
629fn special_operator(bytes: &[u8], at: usize) -> Option<usize> {
630    if bytes[at..].starts_with(b"->>") {
631        return Some(3);
632    }
633    for candidate in [b"::".as_slice(), b":=", b"->", b"**", b"//"] {
634        if bytes[at..].starts_with(candidate) {
635            return Some(2);
636        }
637    }
638    None
639}
640
641#[cfg(test)]
642mod tests {
643    use super::{classes, lookup, tokenize};
644    use crate::generated::keywords::{RESERVED, UNRESERVED};
645    use crate::token::{Flags, Kind, NOT_A_KEYWORD, Token};
646
647    /// Every token but the sentinel, as the pair worth asserting on.
648    fn scan(query: &str) -> Vec<(Kind, &str)> {
649        let tokens = tokenize(query).expect("tokenizes");
650        assert_eq!(tokens.last().map(|t| t.kind), Some(Kind::EndOfInput));
651        assert_eq!(tokens.iter().filter(|t| t.kind == Kind::EndOfInput).count(), 1);
652        tokens[..tokens.len() - 1].iter().map(|t| (t.kind, t.text(query))).collect()
653    }
654
655    fn texts(query: &str) -> Vec<&str> {
656        scan(query).into_iter().map(|(_, text)| text).collect()
657    }
658
659    fn all(query: &str) -> Vec<Token> {
660        tokenize(query).expect("tokenizes")
661    }
662
663    fn message(query: &str) -> String {
664        tokenize(query).expect_err("fails").message().to_string()
665    }
666
667    #[test]
668    fn the_empty_query_is_one_sentinel() {
669        let tokens = tokenize("").expect("tokenizes");
670        assert_eq!(tokens.len(), 1);
671        assert_eq!(tokens[0].kind, Kind::EndOfInput);
672        assert_eq!(tokens[0].span(), rudb_common::Span::new(0, 0));
673        assert!(scan("   \t\n  ").is_empty());
674    }
675
676    #[test]
677    fn a_word_in_a_class_is_a_keyword_and_one_in_none_is_not() {
678        assert_eq!(scan("SELECT"), [(Kind::Keyword, "SELECT")]);
679        assert_eq!(scan("banana"), [(Kind::Identifier, "banana")]);
680    }
681
682    #[test]
683    fn case_is_matched_but_not_folded() {
684        // DuckDB is case insensitive and case preserving, and unlike PostgreSQL that holds for
685        // quoted identifiers too. So the keyword is recognized whatever its case and the bytes
686        // come back exactly as written. Folding here would be the wrong answer twice over.
687        for spelling in ["select", "SELECT", "SeLeCt"] {
688            let tokens = all(spelling);
689            assert_eq!(tokens[0].kind, Kind::Keyword);
690            assert_eq!(tokens[0].text(spelling), spelling);
691            assert_eq!(tokens[0].keyword, lookup("select"));
692        }
693        assert_eq!(scan(r#""Foo""#), [(Kind::QuotedIdentifier, r#""Foo""#)]);
694    }
695
696    #[test]
697    fn a_soft_word_keeps_its_index_and_stays_an_identifier() {
698        // ASCENDING is spelled by a rule and is in none of the five lists. If it came back as a
699        // keyword then `SELECT ascending FROM t` would stop working, and if it came back without
700        // an index then `ORDER BY x ASCENDING` could not be filtered on the literal.
701        let tokens = all("ascending");
702        assert_eq!(tokens[0].kind, Kind::Identifier);
703        assert_ne!(tokens[0].keyword, NOT_A_KEYWORD);
704        assert_eq!(classes(tokens[0].keyword), 0);
705
706        let tokens = all("banana");
707        assert_eq!(tokens[0].keyword, NOT_A_KEYWORD);
708        assert_eq!(classes(tokens[0].keyword), 0);
709    }
710
711    #[test]
712    fn the_classes_come_back_off_the_index() {
713        assert_eq!(classes(lookup("select")) & RESERVED, RESERVED);
714        assert_eq!(classes(lookup("abort")) & UNRESERVED, UNRESERVED);
715        assert_eq!(lookup("supercalifragilistic"), NOT_A_KEYWORD);
716        assert_eq!(lookup(""), NOT_A_KEYWORD);
717    }
718
719    #[test]
720    fn a_quoted_identifier_is_never_a_keyword() {
721        let tokens = all(r#""select""#);
722        assert_eq!(tokens[0].kind, Kind::QuotedIdentifier);
723        assert_eq!(tokens[0].keyword, NOT_A_KEYWORD);
724        assert_eq!(scan(r#""a""b""#), [(Kind::QuotedIdentifier, r#""a""b""#)]);
725    }
726
727    #[test]
728    fn a_dollar_is_an_identifier_character_after_the_first_byte() {
729        assert_eq!(scan("a$b"), [(Kind::Identifier, "a$b")]);
730    }
731
732    #[test]
733    fn any_byte_above_ascii_is_an_identifier_character() {
734        // Nothing decodes UTF-8 here. The arithmetic in `is_word_char` says every byte at or
735        // above 0x80 continues a word, which is how an identifier in any script gets through
736        // without the tokenizer knowing what script it is.
737        assert_eq!(scan("SELECT café"), [(Kind::Keyword, "SELECT"), (Kind::Identifier, "café")]);
738    }
739
740    #[test]
741    fn a_number_swallows_more_than_a_number() {
742        // Every one of these is a single NUMBER token upstream and the parser deals with what it
743        // means. `1.2.3` in particular is not three things.
744        assert_eq!(scan("1.2.3"), [(Kind::Number, "1.2.3")]);
745        assert_eq!(scan("1_000"), [(Kind::Number, "1_000")]);
746        assert_eq!(scan("1e5"), [(Kind::Number, "1e5")]);
747        assert_eq!(scan(".1e5"), [(Kind::Number, ".1e5")]);
748        assert_eq!(scan("1e-5"), [(Kind::Number, "1e-5")]);
749        assert_eq!(scan("1.e5"), [(Kind::Number, "1.e5")]);
750    }
751
752    #[test]
753    fn a_trailing_e_stays_on_the_number() {
754        // `SELECT 1e` is one NUMBER token "1e" and not `1` aliased `e`, because the tokenizer
755        // takes the `e` and only the parser is in a position to object.
756        assert_eq!(scan("SELECT 1e"), [(Kind::Keyword, "SELECT"), (Kind::Number, "1e")]);
757        assert_eq!(scan("SELECT 1e+"), [(Kind::Keyword, "SELECT"), (Kind::Number, "1e+")]);
758    }
759
760    #[test]
761    fn what_the_number_cannot_use_it_gives_back() {
762        // The backtrack at the end of NUMERIC only keeps digits and dots, so `1e+ 1` unwinds the
763        // whole exponent it started and the `e` comes back as a name.
764        assert_eq!(
765            scan("SELECT 1e+ 1"),
766            [
767                (Kind::Keyword, "SELECT"),
768                (Kind::Number, "1"),
769                (Kind::Identifier, "e"),
770                (Kind::Operator, "+"),
771                (Kind::Number, "1"),
772            ]
773        );
774        assert_eq!(scan("1_"), [(Kind::Number, "1"), (Kind::Identifier, "_")]);
775        assert_eq!(scan("1__0"), [(Kind::Number, "1"), (Kind::Identifier, "__0")]);
776        // There is no hex literal. `0x1F` is zero and then a name, which is worth knowing before
777        // somebody reports it as a bug in the parser.
778        assert_eq!(scan("0x1F"), [(Kind::Number, "0"), (Kind::Identifier, "x1F")]);
779        assert_eq!(scan(".e100"), [(Kind::Number, "."), (Kind::Identifier, "e100")]);
780    }
781
782    #[test]
783    fn a_minus_never_joins_an_operator_run() {
784        // `-` is a single byte operator, which is what makes `SELECT 1 =- 1` a subtraction of a
785        // negative one rather than a call to an operator named `=-`.
786        assert_eq!(texts("1-1"), ["1", "-", "1"]);
787        assert_eq!(texts("SELECT 1 =- 1"), ["SELECT", "1", "=", "-", "1"]);
788        assert_eq!(texts("(a,b)"), ["(", "a", ",", "b", ")"]);
789    }
790
791    #[test]
792    fn the_postgres_plus_rule_decides_where_a_run_ends() {
793        // An operator cannot end in `+` unless it contains one of ~ ! @ # % ^ & | ` ?. This is
794        // the rule that lets a user defined operator be told apart from an operator and a sign.
795        assert_eq!(texts("SELECT 1 =+ 1"), ["SELECT", "1", "=", "+", "1"]);
796        assert_eq!(texts("SELECT 1 !=+ 1"), ["SELECT", "1", "!=+", "1"]);
797        assert_eq!(texts("SELECT 1 =++ 1"), ["SELECT", "1", "=", "+", "+", "1"]);
798        assert_eq!(texts("SELECT 1 ++ 1"), ["SELECT", "1", "+", "+", "1"]);
799    }
800
801    #[test]
802    fn the_special_operators_are_checked_before_the_run() {
803        assert_eq!(texts("a->>'b'"), ["a", "->>", "'b'"]);
804        assert_eq!(texts("a->'b'"), ["a", "->", "'b'"]);
805        assert_eq!(texts("a::b"), ["a", "::", "b"]);
806        assert_eq!(texts("a//b"), ["a", "//", "b"]);
807        assert_eq!(texts("2**3"), ["2", "**", "3"]);
808        // But only when what follows is not itself an operator character, in which case the
809        // maximal run wins and it is one token.
810        assert_eq!(texts("a::=b"), ["a", "::=", "b"]);
811    }
812
813    #[test]
814    fn a_block_comment_can_end_an_operator_run() {
815        assert_eq!(texts("1+/*c*/2"), ["1", "+", "2"]);
816    }
817
818    #[test]
819    fn a_comment_is_not_a_token_but_a_block_one_leaves_a_mark() {
820        assert_eq!(texts("SELECT --x\n1"), ["SELECT", "1"]);
821        assert_eq!(texts("SELECT /*x*/ 1"), ["SELECT", "1"]);
822        assert_eq!(texts("SELECT --x"), ["SELECT"]);
823
824        let tokens = all("SELECT /*x*/ 1");
825        assert!(tokens[1].flags.has(Flags::BLOCK_COMMENT));
826        assert!(!tokens[1].flags.has(Flags::NEWLINE));
827
828        // A line comment is not a block comment, and the newline that ends it still counts.
829        let tokens = all("SELECT --x\n1");
830        assert!(!tokens[1].flags.has(Flags::BLOCK_COMMENT));
831        assert!(tokens[1].flags.has(Flags::NEWLINE));
832    }
833
834    #[test]
835    fn the_first_token_is_preceded_by_nothing() {
836        let tokens = all("\n/*x*/ SELECT");
837        assert_eq!(tokens[0].flags, Flags::default());
838    }
839
840    #[test]
841    fn block_comments_nest() {
842        // The reason this matters is commenting out a block that already contains a comment. In
843        // PostgreSQL it works, in most SQL dialects it does not, and DuckDB followed PostgreSQL.
844        assert_eq!(texts("SELECT /* a /* b */ c */ 1"), ["SELECT", "1"]);
845        assert_eq!(
846            message("SELECT /* a /* b */ 1"),
847            "unterminated /* comment at or near \"/* a /* b */ 1\""
848        );
849    }
850
851    #[test]
852    fn a_string_keeps_its_quotes_and_its_escapes() {
853        assert_eq!(scan("'it''s'"), [(Kind::String, "'it''s'")]);
854        assert_eq!(scan("''"), [(Kind::String, "''")]);
855        assert_eq!(texts("'a' 'b'"), ["'a'", "'b'"]);
856    }
857
858    #[test]
859    fn only_the_e_prefix_changes_how_a_string_is_read() {
860        // In an E string a backslash escapes the next byte, so the doubled quote at the end is
861        // one escaped quote and then the close. Without the prefix the backslash is an ordinary
862        // character, the doubled quote is an escape, and the literal runs off the end.
863        assert_eq!(scan(r"E'\''"), [(Kind::String, r"E'\''")]);
864        assert_eq!(message(r"'\''"), "unterminated string literal");
865        for prefix in ["X", "x", "B", "b", "N", "n", "E", "e"] {
866            let query = format!("{prefix}'a'");
867            assert_eq!(tokenize(&query).expect("tokenizes")[0].kind, Kind::String);
868        }
869        // The quote has to be the very next byte, otherwise it is a name and then a string.
870        assert_eq!(texts("x 'a'"), ["x", "'a'"]);
871    }
872
873    #[test]
874    fn a_dollar_quoted_string_is_one_token_and_its_tag_has_to_match() {
875        assert_eq!(scan("$$abc$$"), [(Kind::String, "$$abc$$")]);
876        assert_eq!(scan("$tag$abc$tag$"), [(Kind::String, "$tag$abc$tag$")]);
877        assert_eq!(scan("$tag$a$other$b$tag$"), [(Kind::String, "$tag$a$other$b$tag$")]);
878        assert_eq!(scan("$$it's fine$$"), [(Kind::String, "$$it's fine$$")]);
879    }
880
881    #[test]
882    fn an_unterminated_dollar_quote_is_a_token_and_not_an_error() {
883        // The other three unterminated forms throw. This one does not, which is upstream's
884        // choice and not ours, and the flag is how the difference reaches the matcher.
885        let tokens = all("$$abc");
886        assert_eq!(tokens[0].kind, Kind::String);
887        assert!(tokens[0].flags.has(Flags::UNTERMINATED));
888        assert_eq!(tokens[0].text("$$abc"), "$$abc");
889    }
890
891    #[test]
892    fn a_parameter_is_two_tokens() {
893        // There is no parameter token kind. The grammar spells `$` followed by a number or a
894        // name, so the tokenizer never has to decide which one it is looking at.
895        assert_eq!(scan("$1"), [(Kind::Operator, "$"), (Kind::Number, "1")]);
896        assert_eq!(scan("$banana"), [(Kind::Operator, "$"), (Kind::Identifier, "banana")]);
897        assert_eq!(scan("?"), [(Kind::Operator, "?")]);
898        // A lone dollar has nothing after it to decide with and falls out as an identifier.
899        assert_eq!(scan("$"), [(Kind::Identifier, "$")]);
900    }
901
902    #[test]
903    fn a_semicolon_is_its_own_kind() {
904        assert_eq!(
905            scan("SELECT 1; SELECT 2"),
906            [
907                (Kind::Keyword, "SELECT"),
908                (Kind::Number, "1"),
909                (Kind::Terminator, ";"),
910                (Kind::Keyword, "SELECT"),
911                (Kind::Number, "2"),
912            ]
913        );
914        assert_eq!(scan(";"), [(Kind::Terminator, ";")]);
915    }
916
917    #[test]
918    fn the_four_errors_are_the_four_upstream_throws() {
919        assert_eq!(message("SELECT /* x"), "unterminated /* comment at or near \"/* x\"");
920        assert_eq!(message("SELECT 'x"), "unterminated string literal");
921        assert_eq!(message("SELECT \"x"), "unterminated quoted identifier");
922        assert_eq!(message("SELECT \"\""), "zero-length delimited identifier");
923        // And an error points at where the trouble started, not at the end of the query.
924        assert_eq!(tokenize("SELECT 'x").expect_err("fails").span().map(|s| s.start), Some(7));
925    }
926
927    #[test]
928    fn every_span_lands_where_the_text_is() {
929        let query = "SELECT a, /*c*/ 'b' || $$d$$ FROM t;";
930        for token in tokenize(query).expect("tokenizes") {
931            assert!(token.end as usize <= query.len());
932            assert!(token.start <= token.end);
933            if token.kind != Kind::EndOfInput {
934                assert!(!token.text(query).is_empty());
935            }
936        }
937    }
938
939    #[test]
940    fn a_real_query_comes_out_the_way_it_reads() {
941        assert_eq!(
942            scan("SELECT count(*) FROM t WHERE x > 5 AND y::VARCHAR = 'a';"),
943            [
944                (Kind::Keyword, "SELECT"),
945                (Kind::Identifier, "count"),
946                (Kind::Operator, "("),
947                (Kind::Operator, "*"),
948                (Kind::Operator, ")"),
949                (Kind::Keyword, "FROM"),
950                (Kind::Identifier, "t"),
951                (Kind::Keyword, "WHERE"),
952                (Kind::Identifier, "x"),
953                (Kind::Operator, ">"),
954                (Kind::Number, "5"),
955                (Kind::Keyword, "AND"),
956                (Kind::Identifier, "y"),
957                (Kind::Operator, "::"),
958                (Kind::Keyword, "VARCHAR"),
959                (Kind::Operator, "="),
960                (Kind::String, "'a'"),
961                (Kind::Terminator, ";"),
962            ]
963        );
964    }
965}