hawdb_sql_syntax/
lexer.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::{Span, SyntaxError, SyntaxErrorCode, Token, TokenKind};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct LexerConfig {
19    pub max_input_bytes: usize,
20    pub max_tokens: usize,
21    pub max_block_comment_depth: usize,
22}
23
24impl Default for LexerConfig {
25    fn default() -> Self {
26        Self {
27            max_input_bytes: 16 * 1024 * 1024,
28            max_tokens: 1_000_000,
29            max_block_comment_depth: 64,
30        }
31    }
32}
33
34pub fn tokenize(input: &str) -> Result<Vec<Token>, SyntaxError> {
35    tokenize_with_config(input, LexerConfig::default())
36}
37
38/// Reads only the first non-comment token for explicit statement-family dispatch.
39/// The rest of the input is validated by the selected parser. The full-input
40/// byte limit applies to complete tokenization, not to statement dispatch.
41pub fn first_token(input: &str) -> Result<Token, SyntaxError> {
42    Ok(tokenize_inner(input, LexerConfig::default(), true)?[0])
43}
44
45pub fn tokenize_with_config(input: &str, config: LexerConfig) -> Result<Vec<Token>, SyntaxError> {
46    tokenize_inner(input, config, false)
47}
48
49fn tokenize_inner(
50    input: &str,
51    config: LexerConfig,
52    first_only: bool,
53) -> Result<Vec<Token>, SyntaxError> {
54    // Dispatch only reads a prefix. Applying full-input admission here would
55    // also restrict legacy-routed DDL/DML before their parser is selected.
56    if !first_only && input.len() > config.max_input_bytes {
57        return Err(SyntaxError::new(
58            SyntaxErrorCode::InputTooLarge,
59            Span::new(0, input.len()),
60        ));
61    }
62
63    let bytes = input.as_bytes();
64    let mut tokens = Vec::new();
65    let mut position = 0;
66
67    while position < bytes.len() {
68        let byte = bytes[position];
69        if byte.is_ascii_whitespace() {
70            position += 1;
71            continue;
72        }
73
74        if byte == b'-' && bytes.get(position + 1) == Some(&b'-') {
75            position += 2;
76            while position < bytes.len() && !matches!(bytes[position], b'\n' | b'\r') {
77                position += 1;
78            }
79            continue;
80        }
81
82        if byte == b'/' && bytes.get(position + 1) == Some(&b'*') {
83            position = skip_block_comment(input, position, config.max_block_comment_depth)?;
84            continue;
85        }
86
87        let start = position;
88        let kind = match byte {
89            b'u' | b'U'
90                if bytes.get(position + 1) == Some(&b'&')
91                    && bytes.get(position + 2) == Some(&b'\'') =>
92            {
93                position = scan_quoted(
94                    input,
95                    position + 2,
96                    b'\'',
97                    SyntaxErrorCode::UnterminatedString,
98                )?;
99                TokenKind::String
100            }
101            b'e' | b'E' if bytes.get(position + 1) == Some(&b'\'') => {
102                position = scan_escaped_string(input, position)?;
103                TokenKind::String
104            }
105            b'"' => {
106                position = scan_quoted(
107                    input,
108                    position,
109                    b'"',
110                    SyntaxErrorCode::UnterminatedQuotedIdentifier,
111                )?;
112                TokenKind::QuotedIdentifier
113            }
114            b'\'' => {
115                position =
116                    scan_quoted(input, position, b'\'', SyntaxErrorCode::UnterminatedString)?;
117                TokenKind::String
118            }
119            b'$' => {
120                if bytes.get(position + 1).is_some_and(u8::is_ascii_digit) {
121                    let (next, parameter) = scan_parameter(input, position)?;
122                    position = next;
123                    TokenKind::Parameter(parameter)
124                } else if let Some(next) = scan_dollar_quoted_string(input, position)? {
125                    position = next;
126                    TokenKind::String
127                } else {
128                    return Err(SyntaxError::new(
129                        SyntaxErrorCode::InvalidParameter,
130                        Span::new(position, (position + 1).min(input.len())),
131                    )
132                    .expected("a one-based parameter or dollar-quoted string"));
133                }
134            }
135            b'0'..=b'9' => {
136                position = scan_number(input, position);
137                TokenKind::Number
138            }
139            b'(' => single(&mut position, TokenKind::LeftParen),
140            b')' => single(&mut position, TokenKind::RightParen),
141            b'[' => single(&mut position, TokenKind::LeftBracket),
142            b']' => single(&mut position, TokenKind::RightBracket),
143            b'{' => single(&mut position, TokenKind::LeftBrace),
144            b'}' => single(&mut position, TokenKind::RightBrace),
145            b',' => single(&mut position, TokenKind::Comma),
146            b'.' if bytes.get(position + 1).is_some_and(u8::is_ascii_digit) => {
147                position = scan_number(input, position);
148                TokenKind::Number
149            }
150            b'.' => single(&mut position, TokenKind::Dot),
151            b';' => single(&mut position, TokenKind::Semicolon),
152            b'+' => single(&mut position, TokenKind::Plus),
153            b'*' => single(&mut position, TokenKind::Star),
154            b'%' => single(&mut position, TokenKind::Percent),
155            b'&' => single(&mut position, TokenKind::Ampersand),
156            b'^' => single(&mut position, TokenKind::Caret),
157            b'?' => single(&mut position, TokenKind::Question),
158            b':' if bytes.get(position + 1) == Some(&b':') => {
159                position += 2;
160                TokenKind::DoubleColon
161            }
162            b':' => single(&mut position, TokenKind::Colon),
163            b'|' if bytes.get(position + 1) == Some(&b'|') => {
164                position += 2;
165                TokenKind::Concat
166            }
167            b'|' => single(&mut position, TokenKind::Pipe),
168            b'=' if bytes.get(position + 1) == Some(&b'=') => {
169                position += 2;
170                TokenKind::DoubleEqual
171            }
172            b'=' => single(&mut position, TokenKind::Equal),
173            b'!' if bytes.get(position + 1) == Some(&b'=') => {
174                position += 2;
175                TokenKind::NotEqual
176            }
177            b'<' if bytes.get(position + 1) == Some(&b'=') => {
178                position += 2;
179                TokenKind::LessOrEqual
180            }
181            b'<' if bytes.get(position + 1) == Some(&b'>') => {
182                position += 2;
183                TokenKind::NotEqual
184            }
185            b'<' if bytes.get(position + 1) == Some(&b'-') => {
186                position += 2;
187                TokenKind::ArrowLeft
188            }
189            b'<' => single(&mut position, TokenKind::Less),
190            b'>' if bytes.get(position + 1) == Some(&b'=') => {
191                position += 2;
192                TokenKind::GreaterOrEqual
193            }
194            b'>' => single(&mut position, TokenKind::Greater),
195            b'-' if bytes.get(position + 1) == Some(&b'>') => {
196                position += 2;
197                TokenKind::ArrowRight
198            }
199            b'-' => single(&mut position, TokenKind::Minus),
200            b'/' => single(&mut position, TokenKind::Slash),
201            _ => {
202                let character = input[position..].chars().next().ok_or_else(|| {
203                    SyntaxError::new(
204                        SyntaxErrorCode::InvalidCharacter,
205                        Span::new(position, position),
206                    )
207                })?;
208                if is_identifier_start(character) {
209                    position = scan_identifier(input, position);
210                    TokenKind::Word
211                } else {
212                    let end = position + character.len_utf8();
213                    return Err(SyntaxError::new(
214                        SyntaxErrorCode::InvalidCharacter,
215                        Span::new(position, end),
216                    )
217                    .found(character.to_string()));
218                }
219            }
220        };
221
222        if kind == TokenKind::Number {
223            for position in start..position {
224                if bytes[position] == b'_'
225                    && (!position
226                        .checked_sub(1)
227                        .is_some_and(|previous| bytes[previous].is_ascii_digit())
228                        || !bytes.get(position + 1).is_some_and(u8::is_ascii_digit))
229                {
230                    return Err(SyntaxError::new(
231                        SyntaxErrorCode::InvalidCharacter,
232                        Span::new(position, position + 1),
233                    )
234                    .expected("a digit on each side of a numeric separator")
235                    .found("_"));
236                }
237            }
238        }
239
240        push_token(
241            &mut tokens,
242            Token::new(kind, Span::new(start, position)),
243            config.max_tokens,
244        )?;
245        if first_only {
246            return Ok(tokens);
247        }
248    }
249
250    tokens.push(Token::new(
251        TokenKind::End,
252        Span::new(input.len(), input.len()),
253    ));
254    Ok(tokens)
255}
256
257fn single(position: &mut usize, kind: TokenKind) -> TokenKind {
258    *position += 1;
259    kind
260}
261
262fn push_token(tokens: &mut Vec<Token>, token: Token, max_tokens: usize) -> Result<(), SyntaxError> {
263    if tokens.len() >= max_tokens {
264        return Err(SyntaxError::new(
265            SyntaxErrorCode::TokenLimitExceeded,
266            token.span,
267        ));
268    }
269    tokens.push(token);
270    Ok(())
271}
272
273fn skip_block_comment(input: &str, start: usize, max_depth: usize) -> Result<usize, SyntaxError> {
274    if max_depth == 0 {
275        return Err(SyntaxError::new(
276            SyntaxErrorCode::CommentNestingLimitExceeded,
277            Span::new(start, (start + 2).min(input.len())),
278        ));
279    }
280    let bytes = input.as_bytes();
281    let mut position = start + 2;
282    let mut depth = 1usize;
283    while position < bytes.len() {
284        match (bytes[position], bytes.get(position + 1).copied()) {
285            (b'/', Some(b'*')) => {
286                depth = depth.checked_add(1).ok_or_else(|| {
287                    SyntaxError::new(
288                        SyntaxErrorCode::CommentNestingLimitExceeded,
289                        Span::new(start, position + 2),
290                    )
291                })?;
292                if depth > max_depth {
293                    return Err(SyntaxError::new(
294                        SyntaxErrorCode::CommentNestingLimitExceeded,
295                        Span::new(start, position + 2),
296                    ));
297                }
298                position += 2;
299            }
300            (b'*', Some(b'/')) => {
301                depth -= 1;
302                position += 2;
303                if depth == 0 {
304                    return Ok(position);
305                }
306            }
307            _ => position += 1,
308        }
309    }
310    Err(SyntaxError::new(
311        SyntaxErrorCode::UnterminatedBlockComment,
312        Span::new(start, input.len()),
313    ))
314}
315
316fn scan_quoted(
317    input: &str,
318    start: usize,
319    quote: u8,
320    error_code: SyntaxErrorCode,
321) -> Result<usize, SyntaxError> {
322    let bytes = input.as_bytes();
323    let mut position = start + 1;
324    while position < bytes.len() {
325        if bytes[position] == quote {
326            if bytes.get(position + 1) == Some(&quote) {
327                position += 2;
328            } else {
329                return Ok(position + 1);
330            }
331        } else {
332            position += 1;
333        }
334    }
335    Err(SyntaxError::new(error_code, Span::new(start, input.len())))
336}
337
338fn scan_escaped_string(input: &str, start: usize) -> Result<usize, SyntaxError> {
339    let bytes = input.as_bytes();
340    let mut position = start + 2;
341    while position < bytes.len() {
342        match bytes[position] {
343            b'\\' => position += 2,
344            b'\'' if bytes.get(position + 1) == Some(&b'\'') => position += 2,
345            b'\'' => return Ok(position + 1),
346            _ => position += 1,
347        }
348    }
349    Err(SyntaxError::new(
350        SyntaxErrorCode::UnterminatedString,
351        Span::new(start, input.len()),
352    ))
353}
354
355fn scan_parameter(input: &str, start: usize) -> Result<(usize, u32), SyntaxError> {
356    let bytes = input.as_bytes();
357    let mut position = start + 1;
358    while bytes.get(position).is_some_and(u8::is_ascii_digit) {
359        position += 1;
360    }
361    let value = input[start + 1..position].parse::<u32>().map_err(|_| {
362        SyntaxError::new(
363            SyntaxErrorCode::InvalidParameter,
364            Span::new(start, position),
365        )
366    })?;
367    if value == 0 {
368        return Err(SyntaxError::new(
369            SyntaxErrorCode::InvalidParameter,
370            Span::new(start, position),
371        )
372        .expected("a one-based PostgreSQL parameter"));
373    }
374    Ok((position, value))
375}
376
377fn scan_dollar_quoted_string(input: &str, start: usize) -> Result<Option<usize>, SyntaxError> {
378    let bytes = input.as_bytes();
379    let mut delimiter_end = start + 1;
380    while let Some(byte) = bytes.get(delimiter_end) {
381        if *byte == b'$' {
382            let delimiter = &input[start..=delimiter_end];
383            let content_start = delimiter_end + 1;
384            return input[content_start..]
385                .find(delimiter)
386                .map(|offset| Some(content_start + offset + delimiter.len()))
387                .ok_or_else(|| {
388                    SyntaxError::new(
389                        SyntaxErrorCode::UnterminatedDollarQuotedString,
390                        Span::new(start, input.len()),
391                    )
392                });
393        }
394        if !byte.is_ascii_alphanumeric() && *byte != b'_' {
395            return Ok(None);
396        }
397        delimiter_end += 1;
398    }
399    Ok(None)
400}
401
402fn scan_number(input: &str, start: usize) -> usize {
403    let bytes = input.as_bytes();
404    let mut position = start;
405    while bytes
406        .get(position)
407        .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'_')
408    {
409        position += 1;
410    }
411    if bytes.get(position) == Some(&b'.') && bytes.get(position + 1) != Some(&b'.') {
412        position += 1;
413        while bytes
414            .get(position)
415            .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'_')
416        {
417            position += 1;
418        }
419    }
420    if matches!(bytes.get(position), Some(b'e' | b'E')) {
421        let exponent = position;
422        position += 1;
423        if matches!(bytes.get(position), Some(b'+' | b'-')) {
424            position += 1;
425        }
426        let digit_start = position;
427        while bytes
428            .get(position)
429            .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'_')
430        {
431            position += 1;
432        }
433        if digit_start == position {
434            return exponent;
435        }
436    }
437    position
438}
439
440fn scan_identifier(input: &str, start: usize) -> usize {
441    let mut end = start;
442    for (offset, character) in input[start..].char_indices() {
443        if offset == 0 {
444            if !is_identifier_start(character) {
445                break;
446            }
447        } else if !is_identifier_continue(character) {
448            break;
449        }
450        end = start + offset + character.len_utf8();
451    }
452    end
453}
454
455fn is_identifier_start(character: char) -> bool {
456    character == '_' || character.is_alphabetic()
457}
458
459fn is_identifier_continue(character: char) -> bool {
460    is_identifier_start(character) || character.is_numeric() || character == '$'
461}
462
463#[cfg(test)]
464mod tests {
465    use super::*;
466
467    fn kinds(input: &str) -> Vec<TokenKind> {
468        tokenize(input)
469            .expect("valid SQL")
470            .into_iter()
471            .map(|token| token.kind)
472            .collect()
473    }
474
475    #[test]
476    fn tokenizes_postgres_sql_pgq_surface() {
477        assert_eq!(
478            kinds(
479                "SELECT p.name FROM GRAPH_TABLE (g MATCH (p IS person)-[e IS knows]->(q) \
480                 COLUMNS (p.name AS name, $1)) AS gt;",
481            ),
482            vec![
483                TokenKind::Word,
484                TokenKind::Word,
485                TokenKind::Dot,
486                TokenKind::Word,
487                TokenKind::Word,
488                TokenKind::Word,
489                TokenKind::LeftParen,
490                TokenKind::Word,
491                TokenKind::Word,
492                TokenKind::LeftParen,
493                TokenKind::Word,
494                TokenKind::Word,
495                TokenKind::Word,
496                TokenKind::RightParen,
497                TokenKind::Minus,
498                TokenKind::LeftBracket,
499                TokenKind::Word,
500                TokenKind::Word,
501                TokenKind::Word,
502                TokenKind::RightBracket,
503                TokenKind::ArrowRight,
504                TokenKind::LeftParen,
505                TokenKind::Word,
506                TokenKind::RightParen,
507                TokenKind::Word,
508                TokenKind::LeftParen,
509                TokenKind::Word,
510                TokenKind::Dot,
511                TokenKind::Word,
512                TokenKind::Word,
513                TokenKind::Word,
514                TokenKind::Comma,
515                TokenKind::Parameter(1),
516                TokenKind::RightParen,
517                TokenKind::RightParen,
518                TokenKind::Word,
519                TokenKind::Word,
520                TokenKind::Semicolon,
521                TokenKind::End,
522            ]
523        );
524    }
525
526    #[test]
527    fn skips_nested_comments_and_preserves_byte_spans() {
528        let input = "/* outer /* nested */ done */ SELECT \"na\"\"me\"";
529        let tokens = tokenize(input).expect("valid nested comment");
530        assert_eq!(tokens[0].text(input), "SELECT");
531        assert_eq!(tokens[1].kind, TokenKind::QuotedIdentifier);
532        assert_eq!(tokens[1].text(input), "\"na\"\"me\"");
533    }
534
535    #[test]
536    fn recognizes_dollar_quoted_strings_and_unicode_identifiers() {
537        let input = "SELECT 数据, $tag$MATCH (n)$tag$";
538        let tokens = tokenize(input).expect("valid PostgreSQL tokens");
539        assert_eq!(tokens[1].text(input), "数据");
540        assert_eq!(tokens[3].kind, TokenKind::String);
541    }
542
543    #[test]
544    fn rejects_zero_parameter_and_unterminated_input() {
545        let error = tokenize("SELECT $0").expect_err("zero parameter must fail");
546        assert_eq!(error.code, SyntaxErrorCode::InvalidParameter);
547
548        let error = tokenize("SELECT /*").expect_err("comment must terminate");
549        assert_eq!(error.code, SyntaxErrorCode::UnterminatedBlockComment);
550    }
551
552    #[test]
553    fn enforces_token_and_comment_depth_limits() {
554        let error = tokenize_with_config(
555            "SELECT one",
556            LexerConfig {
557                max_tokens: 1,
558                ..LexerConfig::default()
559            },
560        )
561        .expect_err("token budget must fail");
562        assert_eq!(error.code, SyntaxErrorCode::TokenLimitExceeded);
563
564        let error = tokenize_with_config(
565            "/* /* nested */ */",
566            LexerConfig {
567                max_block_comment_depth: 1,
568                ..LexerConfig::default()
569            },
570        )
571        .expect_err("comment depth must fail");
572        assert_eq!(error.code, SyntaxErrorCode::CommentNestingLimitExceeded);
573
574        let error = tokenize_with_config(
575            "/**/",
576            LexerConfig {
577                max_block_comment_depth: 0,
578                ..LexerConfig::default()
579            },
580        )
581        .expect_err("zero comment budget must fail");
582        assert_eq!(error.code, SyntaxErrorCode::CommentNestingLimitExceeded);
583    }
584}
hawdb_sql_syntax/lexer.rs

hawdb_sql_syntax/
lexer.rs
