Skip to main content

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
38pub fn tokenize_with_config(input: &str, config: LexerConfig) -> Result<Vec<Token>, SyntaxError> {
39    if input.len() > config.max_input_bytes {
40        return Err(SyntaxError::new(
41            SyntaxErrorCode::InputTooLarge,
42            Span::new(0, input.len()),
43        ));
44    }
45
46    let bytes = input.as_bytes();
47    let mut tokens = Vec::new();
48    let mut position = 0;
49
50    while position < bytes.len() {
51        let byte = bytes[position];
52        if byte.is_ascii_whitespace() {
53            position += 1;
54            continue;
55        }
56
57        if byte == b'-' && bytes.get(position + 1) == Some(&b'-') {
58            position += 2;
59            while position < bytes.len() && !matches!(bytes[position], b'\n' | b'\r') {
60                position += 1;
61            }
62            continue;
63        }
64
65        if byte == b'/' && bytes.get(position + 1) == Some(&b'*') {
66            position = skip_block_comment(input, position, config.max_block_comment_depth)?;
67            continue;
68        }
69
70        let start = position;
71        let kind = match byte {
72            b'"' => {
73                position = scan_quoted(
74                    input,
75                    position,
76                    b'"',
77                    SyntaxErrorCode::UnterminatedQuotedIdentifier,
78                )?;
79                TokenKind::QuotedIdentifier
80            }
81            b'\'' => {
82                position =
83                    scan_quoted(input, position, b'\'', SyntaxErrorCode::UnterminatedString)?;
84                TokenKind::String
85            }
86            b'$' => {
87                if bytes.get(position + 1).is_some_and(u8::is_ascii_digit) {
88                    let (next, parameter) = scan_parameter(input, position)?;
89                    position = next;
90                    TokenKind::Parameter(parameter)
91                } else if let Some(next) = scan_dollar_quoted_string(input, position)? {
92                    position = next;
93                    TokenKind::String
94                } else {
95                    return Err(SyntaxError::new(
96                        SyntaxErrorCode::InvalidParameter,
97                        Span::new(position, (position + 1).min(input.len())),
98                    )
99                    .expected("a one-based parameter or dollar-quoted string"));
100                }
101            }
102            b'0'..=b'9' => {
103                position = scan_number(input, position);
104                TokenKind::Number
105            }
106            b'(' => single(&mut position, TokenKind::LeftParen),
107            b')' => single(&mut position, TokenKind::RightParen),
108            b'[' => single(&mut position, TokenKind::LeftBracket),
109            b']' => single(&mut position, TokenKind::RightBracket),
110            b'{' => single(&mut position, TokenKind::LeftBrace),
111            b'}' => single(&mut position, TokenKind::RightBrace),
112            b',' => single(&mut position, TokenKind::Comma),
113            b'.' => single(&mut position, TokenKind::Dot),
114            b';' => single(&mut position, TokenKind::Semicolon),
115            b'+' => single(&mut position, TokenKind::Plus),
116            b'*' => single(&mut position, TokenKind::Star),
117            b'%' => single(&mut position, TokenKind::Percent),
118            b'&' => single(&mut position, TokenKind::Ampersand),
119            b'^' => single(&mut position, TokenKind::Caret),
120            b'?' => single(&mut position, TokenKind::Question),
121            b':' if bytes.get(position + 1) == Some(&b':') => {
122                position += 2;
123                TokenKind::DoubleColon
124            }
125            b':' => single(&mut position, TokenKind::Colon),
126            b'|' if bytes.get(position + 1) == Some(&b'|') => {
127                position += 2;
128                TokenKind::Concat
129            }
130            b'|' => single(&mut position, TokenKind::Pipe),
131            b'=' => single(&mut position, TokenKind::Equal),
132            b'!' if bytes.get(position + 1) == Some(&b'=') => {
133                position += 2;
134                TokenKind::NotEqual
135            }
136            b'<' if bytes.get(position + 1) == Some(&b'=') => {
137                position += 2;
138                TokenKind::LessOrEqual
139            }
140            b'<' if bytes.get(position + 1) == Some(&b'>') => {
141                position += 2;
142                TokenKind::NotEqual
143            }
144            b'<' if bytes.get(position + 1) == Some(&b'-') => {
145                position += 2;
146                TokenKind::ArrowLeft
147            }
148            b'<' => single(&mut position, TokenKind::Less),
149            b'>' if bytes.get(position + 1) == Some(&b'=') => {
150                position += 2;
151                TokenKind::GreaterOrEqual
152            }
153            b'>' => single(&mut position, TokenKind::Greater),
154            b'-' if bytes.get(position + 1) == Some(&b'>') => {
155                position += 2;
156                TokenKind::ArrowRight
157            }
158            b'-' => single(&mut position, TokenKind::Minus),
159            b'/' => single(&mut position, TokenKind::Slash),
160            _ => {
161                let character = input[position..].chars().next().ok_or_else(|| {
162                    SyntaxError::new(
163                        SyntaxErrorCode::InvalidCharacter,
164                        Span::new(position, position),
165                    )
166                })?;
167                if is_identifier_start(character) {
168                    position = scan_identifier(input, position);
169                    TokenKind::Word
170                } else {
171                    let end = position + character.len_utf8();
172                    return Err(SyntaxError::new(
173                        SyntaxErrorCode::InvalidCharacter,
174                        Span::new(position, end),
175                    )
176                    .found(character.to_string()));
177                }
178            }
179        };
180
181        push_token(
182            &mut tokens,
183            Token::new(kind, Span::new(start, position)),
184            config.max_tokens,
185        )?;
186    }
187
188    tokens.push(Token::new(
189        TokenKind::End,
190        Span::new(input.len(), input.len()),
191    ));
192    Ok(tokens)
193}
194
195fn single(position: &mut usize, kind: TokenKind) -> TokenKind {
196    *position += 1;
197    kind
198}
199
200fn push_token(tokens: &mut Vec<Token>, token: Token, max_tokens: usize) -> Result<(), SyntaxError> {
201    if tokens.len() >= max_tokens {
202        return Err(SyntaxError::new(
203            SyntaxErrorCode::TokenLimitExceeded,
204            token.span,
205        ));
206    }
207    tokens.push(token);
208    Ok(())
209}
210
211fn skip_block_comment(input: &str, start: usize, max_depth: usize) -> Result<usize, SyntaxError> {
212    if max_depth == 0 {
213        return Err(SyntaxError::new(
214            SyntaxErrorCode::CommentNestingLimitExceeded,
215            Span::new(start, (start + 2).min(input.len())),
216        ));
217    }
218    let bytes = input.as_bytes();
219    let mut position = start + 2;
220    let mut depth = 1usize;
221    while position < bytes.len() {
222        match (bytes[position], bytes.get(position + 1).copied()) {
223            (b'/', Some(b'*')) => {
224                depth = depth.checked_add(1).ok_or_else(|| {
225                    SyntaxError::new(
226                        SyntaxErrorCode::CommentNestingLimitExceeded,
227                        Span::new(start, position + 2),
228                    )
229                })?;
230                if depth > max_depth {
231                    return Err(SyntaxError::new(
232                        SyntaxErrorCode::CommentNestingLimitExceeded,
233                        Span::new(start, position + 2),
234                    ));
235                }
236                position += 2;
237            }
238            (b'*', Some(b'/')) => {
239                depth -= 1;
240                position += 2;
241                if depth == 0 {
242                    return Ok(position);
243                }
244            }
245            _ => position += 1,
246        }
247    }
248    Err(SyntaxError::new(
249        SyntaxErrorCode::UnterminatedBlockComment,
250        Span::new(start, input.len()),
251    ))
252}
253
254fn scan_quoted(
255    input: &str,
256    start: usize,
257    quote: u8,
258    error_code: SyntaxErrorCode,
259) -> Result<usize, SyntaxError> {
260    let bytes = input.as_bytes();
261    let mut position = start + 1;
262    while position < bytes.len() {
263        if bytes[position] == quote {
264            if bytes.get(position + 1) == Some(&quote) {
265                position += 2;
266            } else {
267                return Ok(position + 1);
268            }
269        } else {
270            position += 1;
271        }
272    }
273    Err(SyntaxError::new(error_code, Span::new(start, input.len())))
274}
275
276fn scan_parameter(input: &str, start: usize) -> Result<(usize, u32), SyntaxError> {
277    let bytes = input.as_bytes();
278    let mut position = start + 1;
279    while bytes.get(position).is_some_and(u8::is_ascii_digit) {
280        position += 1;
281    }
282    let value = input[start + 1..position].parse::<u32>().map_err(|_| {
283        SyntaxError::new(
284            SyntaxErrorCode::InvalidParameter,
285            Span::new(start, position),
286        )
287    })?;
288    if value == 0 {
289        return Err(SyntaxError::new(
290            SyntaxErrorCode::InvalidParameter,
291            Span::new(start, position),
292        )
293        .expected("a one-based PostgreSQL parameter"));
294    }
295    Ok((position, value))
296}
297
298fn scan_dollar_quoted_string(input: &str, start: usize) -> Result<Option<usize>, SyntaxError> {
299    let bytes = input.as_bytes();
300    let mut delimiter_end = start + 1;
301    while let Some(byte) = bytes.get(delimiter_end) {
302        if *byte == b'$' {
303            let delimiter = &input[start..=delimiter_end];
304            let content_start = delimiter_end + 1;
305            return input[content_start..]
306                .find(delimiter)
307                .map(|offset| Some(content_start + offset + delimiter.len()))
308                .ok_or_else(|| {
309                    SyntaxError::new(
310                        SyntaxErrorCode::UnterminatedDollarQuotedString,
311                        Span::new(start, input.len()),
312                    )
313                });
314        }
315        if !byte.is_ascii_alphanumeric() && *byte != b'_' {
316            return Ok(None);
317        }
318        delimiter_end += 1;
319    }
320    Ok(None)
321}
322
323fn scan_number(input: &str, start: usize) -> usize {
324    let bytes = input.as_bytes();
325    let mut position = start;
326    while bytes.get(position).is_some_and(u8::is_ascii_digit) {
327        position += 1;
328    }
329    if bytes.get(position) == Some(&b'.') && bytes.get(position + 1).is_some_and(u8::is_ascii_digit)
330    {
331        position += 1;
332        while bytes.get(position).is_some_and(u8::is_ascii_digit) {
333            position += 1;
334        }
335    }
336    if matches!(bytes.get(position), Some(b'e' | b'E')) {
337        let exponent = position;
338        position += 1;
339        if matches!(bytes.get(position), Some(b'+' | b'-')) {
340            position += 1;
341        }
342        let digit_start = position;
343        while bytes.get(position).is_some_and(u8::is_ascii_digit) {
344            position += 1;
345        }
346        if digit_start == position {
347            return exponent;
348        }
349    }
350    position
351}
352
353fn scan_identifier(input: &str, start: usize) -> usize {
354    let mut end = start;
355    for (offset, character) in input[start..].char_indices() {
356        if offset == 0 {
357            if !is_identifier_start(character) {
358                break;
359            }
360        } else if !is_identifier_continue(character) {
361            break;
362        }
363        end = start + offset + character.len_utf8();
364    }
365    end
366}
367
368fn is_identifier_start(character: char) -> bool {
369    character == '_' || character.is_alphabetic()
370}
371
372fn is_identifier_continue(character: char) -> bool {
373    is_identifier_start(character) || character.is_numeric() || character == '$'
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379
380    fn kinds(input: &str) -> Vec<TokenKind> {
381        tokenize(input)
382            .expect("valid SQL")
383            .into_iter()
384            .map(|token| token.kind)
385            .collect()
386    }
387
388    #[test]
389    fn tokenizes_postgres_sql_pgq_surface() {
390        assert_eq!(
391            kinds(
392                "SELECT p.name FROM GRAPH_TABLE (g MATCH (p IS person)-[e IS knows]->(q) \
393                 COLUMNS (p.name AS name, $1)) AS gt;",
394            ),
395            vec![
396                TokenKind::Word,
397                TokenKind::Word,
398                TokenKind::Dot,
399                TokenKind::Word,
400                TokenKind::Word,
401                TokenKind::Word,
402                TokenKind::LeftParen,
403                TokenKind::Word,
404                TokenKind::Word,
405                TokenKind::LeftParen,
406                TokenKind::Word,
407                TokenKind::Word,
408                TokenKind::Word,
409                TokenKind::RightParen,
410                TokenKind::Minus,
411                TokenKind::LeftBracket,
412                TokenKind::Word,
413                TokenKind::Word,
414                TokenKind::Word,
415                TokenKind::RightBracket,
416                TokenKind::ArrowRight,
417                TokenKind::LeftParen,
418                TokenKind::Word,
419                TokenKind::RightParen,
420                TokenKind::Word,
421                TokenKind::LeftParen,
422                TokenKind::Word,
423                TokenKind::Dot,
424                TokenKind::Word,
425                TokenKind::Word,
426                TokenKind::Word,
427                TokenKind::Comma,
428                TokenKind::Parameter(1),
429                TokenKind::RightParen,
430                TokenKind::RightParen,
431                TokenKind::Word,
432                TokenKind::Word,
433                TokenKind::Semicolon,
434                TokenKind::End,
435            ]
436        );
437    }
438
439    #[test]
440    fn skips_nested_comments_and_preserves_byte_spans() {
441        let input = "/* outer /* nested */ done */ SELECT \"na\"\"me\"";
442        let tokens = tokenize(input).expect("valid nested comment");
443        assert_eq!(tokens[0].text(input), "SELECT");
444        assert_eq!(tokens[1].kind, TokenKind::QuotedIdentifier);
445        assert_eq!(tokens[1].text(input), "\"na\"\"me\"");
446    }
447
448    #[test]
449    fn recognizes_dollar_quoted_strings_and_unicode_identifiers() {
450        let input = "SELECT 数据, $tag$MATCH (n)$tag$";
451        let tokens = tokenize(input).expect("valid PostgreSQL tokens");
452        assert_eq!(tokens[1].text(input), "数据");
453        assert_eq!(tokens[3].kind, TokenKind::String);
454    }
455
456    #[test]
457    fn rejects_zero_parameter_and_unterminated_input() {
458        let error = tokenize("SELECT $0").expect_err("zero parameter must fail");
459        assert_eq!(error.code, SyntaxErrorCode::InvalidParameter);
460
461        let error = tokenize("SELECT /*").expect_err("comment must terminate");
462        assert_eq!(error.code, SyntaxErrorCode::UnterminatedBlockComment);
463    }
464
465    #[test]
466    fn enforces_token_and_comment_depth_limits() {
467        let error = tokenize_with_config(
468            "SELECT one",
469            LexerConfig {
470                max_tokens: 1,
471                ..LexerConfig::default()
472            },
473        )
474        .expect_err("token budget must fail");
475        assert_eq!(error.code, SyntaxErrorCode::TokenLimitExceeded);
476
477        let error = tokenize_with_config(
478            "/* /* nested */ */",
479            LexerConfig {
480                max_block_comment_depth: 1,
481                ..LexerConfig::default()
482            },
483        )
484        .expect_err("comment depth must fail");
485        assert_eq!(error.code, SyntaxErrorCode::CommentNestingLimitExceeded);
486
487        let error = tokenize_with_config(
488            "/**/",
489            LexerConfig {
490                max_block_comment_depth: 0,
491                ..LexerConfig::default()
492            },
493        )
494        .expect_err("zero comment budget must fail");
495        assert_eq!(error.code, SyntaxErrorCode::CommentNestingLimitExceeded);
496    }
497}