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                "begin" => Token::Begin,
213                "commit" => Token::Commit,
214                "rollback" => Token::Rollback,
215                "view" => Token::View,
216                "materialized" => Token::Materialized,
217                "materialize" => Token::Materialized,
218                "refresh" => Token::Refresh,
219                "union" => Token::Union,
220                "having" => Token::Having,
221                "distinct" => Token::Distinct,
222                "in" => Token::In,
223                "between" => Token::Between,
224                "like" => Token::Like,
225                "count" => Token::Count,
226                "avg" => Token::Avg,
227                "sum" => Token::Sum,
228                "min" => Token::Min,
229                "max" => Token::Max,
230                "is" => Token::Is,
231                "null" => Token::Null,
232                "upper" => Token::Upper,
233                "lower" => Token::Lower,
234                "length" => Token::Length,
235                "trim" => Token::Trim,
236                "substring" => Token::Substring,
237                "concat" => Token::Concat,
238                "abs" => Token::Abs,
239                "round" => Token::Round,
240                "ceil" => Token::Ceil,
241                "floor" => Token::Floor,
242                "sqrt" => Token::Sqrt,
243                "pow" => Token::Pow,
244                "now" => Token::Now,
245                "extract" => Token::Extract,
246                "date_add" => Token::DateAdd,
247                "date_diff" => Token::DateDiff,
248                "cast" => Token::Cast,
249                "case" => Token::Case,
250                "when" => Token::When,
251                "then" => Token::Then,
252                "else" => Token::Else,
253                "end" => Token::End,
254                "over" => Token::Over,
255                "partition" => Token::Partition,
256                "row_number" => Token::RowNumber,
257                "rank" => Token::Rank,
258                "dense_rank" => Token::DenseRank,
259                "alter" => Token::Alter,
260                "drop" => Token::Drop,
261                "add" => Token::Add,
262                "column" => Token::Column,
263                "explain" => Token::Explain,
264                "true" => Token::BoolLit(true),
265                "false" => Token::BoolLit(false),
266                _ => Token::Ident(word),
267            };
268            tokens.push(token);
269            continue;
270        }
271
272        // Two-char operators
273        if pos + 1 < chars.len() {
274            let two: String = chars[pos..pos + 2].iter().collect();
275            match two.as_str() {
276                ":=" => {
277                    tokens.push(Token::Assign);
278                    pos += 2;
279                    continue;
280                }
281                "->" => {
282                    tokens.push(Token::Arrow);
283                    pos += 2;
284                    continue;
285                }
286                "!=" => {
287                    tokens.push(Token::Neq);
288                    pos += 2;
289                    continue;
290                }
291                "<=" => {
292                    tokens.push(Token::Lte);
293                    pos += 2;
294                    continue;
295                }
296                ">=" => {
297                    tokens.push(Token::Gte);
298                    pos += 2;
299                    continue;
300                }
301                "??" => {
302                    tokens.push(Token::Coalesce);
303                    pos += 2;
304                    continue;
305                }
306                _ => {}
307            }
308        }
309
310        // Single-char operators
311        let token = match chars[pos] {
312            '=' => Token::Eq,
313            '<' => Token::Lt,
314            '>' => Token::Gt,
315            '|' => Token::Pipe,
316            '+' => Token::Plus,
317            '-' => Token::Minus,
318            '*' => Token::Star,
319            '/' => Token::Slash,
320            '{' => Token::LBrace,
321            '}' => Token::RBrace,
322            '(' => Token::LParen,
323            ')' => Token::RParen,
324            ',' => Token::Comma,
325            ':' => Token::Colon,
326            '.' => Token::Dot,
327            c => {
328                return Err(LexError {
329                    message: format!("unexpected character: {c}"),
330                    position: pos,
331                })
332            }
333        };
334        tokens.push(token);
335        pos += 1;
336    }
337
338    tokens.push(Token::Eof);
339    Ok(tokens)
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345    use crate::token::Token;
346
347    #[test]
348    fn test_lex_simple_query() {
349        let tokens = lex("User filter .age > 30").unwrap();
350        assert_eq!(
351            tokens,
352            vec![
353                Token::Ident("User".into()),
354                Token::Filter,
355                Token::DotIdent("age".into()),
356                Token::Gt,
357                Token::IntLit(30),
358                Token::Eof,
359            ]
360        );
361    }
362
363    #[test]
364    fn test_lex_projection() {
365        let tokens = lex("User { name, email }").unwrap();
366        assert_eq!(
367            tokens,
368            vec![
369                Token::Ident("User".into()),
370                Token::LBrace,
371                Token::Ident("name".into()),
372                Token::Comma,
373                Token::Ident("email".into()),
374                Token::RBrace,
375                Token::Eof,
376            ]
377        );
378    }
379
380    #[test]
381    fn test_lex_insert() {
382        let tokens = lex(r#"insert User { name := "Alice", age := 30 }"#).unwrap();
383        assert_eq!(
384            tokens,
385            vec![
386                Token::Insert,
387                Token::Ident("User".into()),
388                Token::LBrace,
389                Token::Ident("name".into()),
390                Token::Assign,
391                Token::StringLit("Alice".into()),
392                Token::Comma,
393                Token::Ident("age".into()),
394                Token::Assign,
395                Token::IntLit(30),
396                Token::RBrace,
397                Token::Eof,
398            ]
399        );
400    }
401
402    #[test]
403    fn test_lex_params() {
404        let tokens = lex("User filter .age > $min_age").unwrap();
405        assert_eq!(
406            tokens,
407            vec![
408                Token::Ident("User".into()),
409                Token::Filter,
410                Token::DotIdent("age".into()),
411                Token::Gt,
412                Token::Param("min_age".into()),
413                Token::Eof,
414            ]
415        );
416    }
417
418    #[test]
419    fn test_lex_string_with_escapes() {
420        let tokens = lex(r#""hello \"world\"""#).unwrap();
421        assert_eq!(
422            tokens,
423            vec![Token::StringLit("hello \"world\"".into()), Token::Eof,]
424        );
425    }
426
427    #[test]
428    fn test_lex_aggregation() {
429        let tokens = lex("count(User)").unwrap();
430        assert_eq!(
431            tokens,
432            vec![
433                Token::Count,
434                Token::LParen,
435                Token::Ident("User".into()),
436                Token::RParen,
437                Token::Eof,
438            ]
439        );
440    }
441
442    /// Regression for issue #24: an integer literal with more digits than
443    /// i64 can hold previously reached `s.parse::<i64>().unwrap()` and
444    /// panicked. It must return a `LexError` instead.
445    #[test]
446    fn test_lex_intlit_overflow_returns_err() {
447        // 22 digits — well past i64::MAX (19 digits).
448        let err = lex("4444444441111111144444").expect_err("must error, not panic");
449        assert!(
450            err.message.contains("integer literal out of range"),
451            "unexpected message: {}",
452            err.message
453        );
454        assert_eq!(err.position, 0);
455    }
456
457    /// Same bug, reached via the exact fuzzer reproducer from the
458    /// libFuzzer artifact attached to issue #24 (base64
459    /// `YXMJCQkJCQkJCQkJCQkJNDQ0NDQ0NDQ0MTExMTExMTQ0NDQJCQkJCQk=`).
460    #[test]
461    fn test_lex_fuzz_repro_issue_24() {
462        let input = "as\t\t\t\t\t\t\t\t\t\t\t\t\t44444444411111114444\t\t\t\t\t\t";
463        let err = lex(input).expect_err("fuzz reproducer must now error, not panic");
464        assert!(err.message.contains("integer literal"));
465    }
466}