1use 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("e) {
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}