Skip to main content

powdb_query/
lexer.rs

1use crate::token::Token;
2
3/// Maximum allowed length for a string literal (16 MB).
4/// Prevents unbounded memory consumption from queries with multi-gigabyte strings.
5const MAX_STRING_LITERAL: usize = 16 * 1024 * 1024;
6
7#[derive(Debug)]
8pub struct LexError {
9    pub message: String,
10    pub position: usize,
11}
12
13impl std::fmt::Display for LexError {
14    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
15        write!(f, "at position {}: {}", self.position, self.message)
16    }
17}
18
19impl std::error::Error for LexError {}
20
21/// Tokenize a PowQL input string into a stream of tokens.
22///
23/// # Examples
24///
25/// ```
26/// use powdb_query::lexer::lex;
27/// use powdb_query::token::Token;
28///
29/// let tokens = lex("User filter .age > 30").unwrap();
30/// assert_eq!(tokens[0], Token::Ident("User".to_string()));
31/// assert_eq!(tokens[1], Token::Filter);
32/// assert_eq!(tokens[2], Token::DotIdent("age".to_string()));
33/// ```
34pub fn lex(input: &str) -> Result<Vec<Token>, LexError> {
35    let mut tokens = Vec::new();
36    let chars: Vec<char> = input.chars().collect();
37    let mut pos = 0;
38
39    while pos < chars.len() {
40        // Skip whitespace
41        if chars[pos].is_whitespace() {
42            pos += 1;
43            continue;
44        }
45
46        // Skip comments
47        if chars[pos] == '#' {
48            while pos < chars.len() && chars[pos] != '\n' {
49                pos += 1;
50            }
51            continue;
52        }
53
54        // Dot-ident: .fieldname
55        if chars[pos] == '.'
56            && pos + 1 < chars.len()
57            && (chars[pos + 1].is_alphabetic() || chars[pos + 1] == '_')
58        {
59            pos += 1; // skip dot
60            let start = pos;
61            while pos < chars.len() && (chars[pos].is_alphanumeric() || chars[pos] == '_') {
62                pos += 1;
63            }
64            let name: String = chars[start..pos].iter().collect();
65            tokens.push(Token::DotIdent(name));
66            continue;
67        }
68
69        // Param: $name
70        if chars[pos] == '$' {
71            pos += 1;
72            let start = pos;
73            while pos < chars.len() && (chars[pos].is_alphanumeric() || chars[pos] == '_') {
74                pos += 1;
75            }
76            let name: String = chars[start..pos].iter().collect();
77            tokens.push(Token::Param(name));
78            continue;
79        }
80
81        // String literal
82        if chars[pos] == '"' {
83            pos += 1;
84            let mut s = String::new();
85            while pos < chars.len() && chars[pos] != '"' {
86                if chars[pos] == '\\' && pos + 1 < chars.len() {
87                    match chars[pos + 1] {
88                        '"' => {
89                            s.push('"');
90                            pos += 2;
91                        }
92                        '\\' => {
93                            s.push('\\');
94                            pos += 2;
95                        }
96                        'n' => {
97                            s.push('\n');
98                            pos += 2;
99                        }
100                        't' => {
101                            s.push('\t');
102                            pos += 2;
103                        }
104                        _ => {
105                            s.push(chars[pos + 1]);
106                            pos += 2;
107                        }
108                    }
109                } else {
110                    s.push(chars[pos]);
111                    pos += 1;
112                }
113            }
114            if pos >= chars.len() {
115                return Err(LexError {
116                    message: "unterminated string".into(),
117                    position: pos,
118                });
119            }
120            pos += 1; // closing quote
121            if s.len() > MAX_STRING_LITERAL {
122                return Err(LexError {
123                    message: format!(
124                        "string literal exceeds maximum size of {}MB",
125                        MAX_STRING_LITERAL / (1024 * 1024)
126                    ),
127                    position: pos,
128                });
129            }
130            tokens.push(Token::StringLit(s));
131            continue;
132        }
133
134        // Number (int or float)
135        if chars[pos].is_ascii_digit()
136            || (chars[pos] == '-' && pos + 1 < chars.len() && chars[pos + 1].is_ascii_digit())
137        {
138            let start = pos;
139            if chars[pos] == '-' {
140                pos += 1;
141            }
142            while pos < chars.len() && chars[pos].is_ascii_digit() {
143                pos += 1;
144            }
145            if pos < chars.len()
146                && chars[pos] == '.'
147                && pos + 1 < chars.len()
148                && chars[pos + 1].is_ascii_digit()
149            {
150                pos += 1;
151                while pos < chars.len() && chars[pos].is_ascii_digit() {
152                    pos += 1;
153                }
154                let s: String = chars[start..pos].iter().collect();
155                let value = s.parse::<f64>().map_err(|_| LexError {
156                    message: format!("float literal out of range: {s}"),
157                    position: start,
158                })?;
159                tokens.push(Token::FloatLit(value));
160            } else {
161                let s: String = chars[start..pos].iter().collect();
162                let value = s.parse::<i64>().map_err(|_| LexError {
163                    message: format!("integer literal out of range for i64: {s}"),
164                    position: start,
165                })?;
166                tokens.push(Token::IntLit(value));
167            }
168            continue;
169        }
170
171        // Identifiers and keywords
172        if chars[pos].is_alphabetic() || chars[pos] == '_' {
173            let start = pos;
174            while pos < chars.len() && (chars[pos].is_alphanumeric() || chars[pos] == '_') {
175                pos += 1;
176            }
177            let word: String = chars[start..pos].iter().collect();
178            let token = match word.as_str() {
179                "type" => Token::Type,
180                "filter" => Token::Filter,
181                "order" => Token::Order,
182                "limit" => Token::Limit,
183                "offset" => Token::Offset,
184                "insert" => Token::Insert,
185                "update" => Token::Update,
186                "delete" => Token::Delete,
187                "upsert" => Token::Upsert,
188                "conflict" => Token::Conflict,
189                "select" => Token::Select,
190                "required" => Token::Required,
191                "multi" => Token::Multi,
192                "link" => Token::Link,
193                "index" => Token::Index,
194                "on" => Token::On,
195                "asc" => Token::Asc,
196                "desc" => Token::Desc,
197                "and" => Token::And,
198                "or" => Token::Or,
199                "not" => Token::Not,
200                "exists" => Token::Exists,
201                "let" => Token::Let,
202                "as" => Token::As,
203                "match" => Token::Match,
204                "group" => Token::Group,
205                "join" => Token::Join,
206                "inner" => Token::Inner,
207                "left" => Token::LeftKw,
208                "right" => Token::RightKw,
209                "outer" => Token::Outer,
210                "cross" => Token::Cross,
211                "transaction" => Token::Transaction,
212                "view" => Token::View,
213                "materialized" => Token::Materialized,
214                "materialize" => Token::Materialized,
215                "refresh" => Token::Refresh,
216                "union" => Token::Union,
217                "having" => Token::Having,
218                "distinct" => Token::Distinct,
219                "in" => Token::In,
220                "between" => Token::Between,
221                "like" => Token::Like,
222                "count" => Token::Count,
223                "avg" => Token::Avg,
224                "sum" => Token::Sum,
225                "min" => Token::Min,
226                "max" => Token::Max,
227                "is" => Token::Is,
228                "null" => Token::Null,
229                "upper" => Token::Upper,
230                "lower" => Token::Lower,
231                "length" => Token::Length,
232                "trim" => Token::Trim,
233                "substring" => Token::Substring,
234                "concat" => Token::Concat,
235                "abs" => Token::Abs,
236                "round" => Token::Round,
237                "ceil" => Token::Ceil,
238                "floor" => Token::Floor,
239                "sqrt" => Token::Sqrt,
240                "pow" => Token::Pow,
241                "now" => Token::Now,
242                "extract" => Token::Extract,
243                "date_add" => Token::DateAdd,
244                "date_diff" => Token::DateDiff,
245                "cast" => Token::Cast,
246                "case" => Token::Case,
247                "when" => Token::When,
248                "then" => Token::Then,
249                "else" => Token::Else,
250                "end" => Token::End,
251                "over" => Token::Over,
252                "partition" => Token::Partition,
253                "row_number" => Token::RowNumber,
254                "rank" => Token::Rank,
255                "dense_rank" => Token::DenseRank,
256                "alter" => Token::Alter,
257                "drop" => Token::Drop,
258                "add" => Token::Add,
259                "column" => Token::Column,
260                "explain" => Token::Explain,
261                "true" => Token::BoolLit(true),
262                "false" => Token::BoolLit(false),
263                _ => Token::Ident(word),
264            };
265            tokens.push(token);
266            continue;
267        }
268
269        // Two-char operators
270        if pos + 1 < chars.len() {
271            let two: String = chars[pos..pos + 2].iter().collect();
272            match two.as_str() {
273                ":=" => {
274                    tokens.push(Token::Assign);
275                    pos += 2;
276                    continue;
277                }
278                "->" => {
279                    tokens.push(Token::Arrow);
280                    pos += 2;
281                    continue;
282                }
283                "!=" => {
284                    tokens.push(Token::Neq);
285                    pos += 2;
286                    continue;
287                }
288                "<=" => {
289                    tokens.push(Token::Lte);
290                    pos += 2;
291                    continue;
292                }
293                ">=" => {
294                    tokens.push(Token::Gte);
295                    pos += 2;
296                    continue;
297                }
298                "??" => {
299                    tokens.push(Token::Coalesce);
300                    pos += 2;
301                    continue;
302                }
303                _ => {}
304            }
305        }
306
307        // Single-char operators
308        let token = match chars[pos] {
309            '=' => Token::Eq,
310            '<' => Token::Lt,
311            '>' => Token::Gt,
312            '|' => Token::Pipe,
313            '+' => Token::Plus,
314            '-' => Token::Minus,
315            '*' => Token::Star,
316            '/' => Token::Slash,
317            '{' => Token::LBrace,
318            '}' => Token::RBrace,
319            '(' => Token::LParen,
320            ')' => Token::RParen,
321            ',' => Token::Comma,
322            ':' => Token::Colon,
323            '.' => Token::Dot,
324            c => {
325                return Err(LexError {
326                    message: format!("unexpected character: {c}"),
327                    position: pos,
328                })
329            }
330        };
331        tokens.push(token);
332        pos += 1;
333    }
334
335    tokens.push(Token::Eof);
336    Ok(tokens)
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342    use crate::token::Token;
343
344    #[test]
345    fn test_lex_simple_query() {
346        let tokens = lex("User filter .age > 30").unwrap();
347        assert_eq!(
348            tokens,
349            vec![
350                Token::Ident("User".into()),
351                Token::Filter,
352                Token::DotIdent("age".into()),
353                Token::Gt,
354                Token::IntLit(30),
355                Token::Eof,
356            ]
357        );
358    }
359
360    #[test]
361    fn test_lex_projection() {
362        let tokens = lex("User { name, email }").unwrap();
363        assert_eq!(
364            tokens,
365            vec![
366                Token::Ident("User".into()),
367                Token::LBrace,
368                Token::Ident("name".into()),
369                Token::Comma,
370                Token::Ident("email".into()),
371                Token::RBrace,
372                Token::Eof,
373            ]
374        );
375    }
376
377    #[test]
378    fn test_lex_insert() {
379        let tokens = lex(r#"insert User { name := "Alice", age := 30 }"#).unwrap();
380        assert_eq!(
381            tokens,
382            vec![
383                Token::Insert,
384                Token::Ident("User".into()),
385                Token::LBrace,
386                Token::Ident("name".into()),
387                Token::Assign,
388                Token::StringLit("Alice".into()),
389                Token::Comma,
390                Token::Ident("age".into()),
391                Token::Assign,
392                Token::IntLit(30),
393                Token::RBrace,
394                Token::Eof,
395            ]
396        );
397    }
398
399    #[test]
400    fn test_lex_params() {
401        let tokens = lex("User filter .age > $min_age").unwrap();
402        assert_eq!(
403            tokens,
404            vec![
405                Token::Ident("User".into()),
406                Token::Filter,
407                Token::DotIdent("age".into()),
408                Token::Gt,
409                Token::Param("min_age".into()),
410                Token::Eof,
411            ]
412        );
413    }
414
415    #[test]
416    fn test_lex_string_with_escapes() {
417        let tokens = lex(r#""hello \"world\"""#).unwrap();
418        assert_eq!(
419            tokens,
420            vec![Token::StringLit("hello \"world\"".into()), Token::Eof,]
421        );
422    }
423
424    #[test]
425    fn test_lex_aggregation() {
426        let tokens = lex("count(User)").unwrap();
427        assert_eq!(
428            tokens,
429            vec![
430                Token::Count,
431                Token::LParen,
432                Token::Ident("User".into()),
433                Token::RParen,
434                Token::Eof,
435            ]
436        );
437    }
438
439    /// Regression for issue #24: an integer literal with more digits than
440    /// i64 can hold previously reached `s.parse::<i64>().unwrap()` and
441    /// panicked. It must return a `LexError` instead.
442    #[test]
443    fn test_lex_intlit_overflow_returns_err() {
444        // 22 digits — well past i64::MAX (19 digits).
445        let err = lex("4444444441111111144444").expect_err("must error, not panic");
446        assert!(
447            err.message.contains("integer literal out of range"),
448            "unexpected message: {}",
449            err.message
450        );
451        assert_eq!(err.position, 0);
452    }
453
454    /// Same bug, reached via the exact fuzzer reproducer from the
455    /// libFuzzer artifact attached to issue #24 (base64
456    /// `YXMJCQkJCQkJCQkJCQkJNDQ0NDQ0NDQ0MTExMTExMTQ0NDQJCQkJCQk=`).
457    #[test]
458    fn test_lex_fuzz_repro_issue_24() {
459        let input = "as\t\t\t\t\t\t\t\t\t\t\t\t\t44444444411111114444\t\t\t\t\t\t";
460        let err = lex(input).expect_err("fuzz reproducer must now error, not panic");
461        assert!(err.message.contains("integer literal"));
462    }
463}