Skip to main content

ddb_core/lexer/
tokenizer.rs

1use super::types::{Token, TokenType};
2use crate::error::{DdbError, Result};
3use nom::{
4    branch::alt,
5    bytes::complete::{is_not, tag, take_while},
6    character::complete::{alpha1, alphanumeric1, char, digit1, multispace1},
7    combinator::{map, opt, recognize, value},
8    multi::many0,
9    sequence::{delimited, pair, preceded},
10    IResult,
11};
12
13/// Tokenizer for SQL queries
14pub struct Tokenizer {
15    position: usize,
16    line: usize,
17    column: usize,
18}
19
20impl Tokenizer {
21    pub fn new() -> Self {
22        Self {
23            position: 0,
24            line: 1,
25            column: 1,
26        }
27    }
28
29    /// Tokenize a SQL query string
30    pub fn tokenize(&mut self, input: &str) -> Result<Vec<Token>> {
31        let mut tokens = Vec::new();
32        let mut remaining = input;
33
34        while !remaining.is_empty() {
35            // Try to parse a token
36            match self.parse_token(remaining) {
37                Ok((rest, token_type)) => {
38                    let consumed = remaining.len() - rest.len();
39
40                    // Skip whitespace tokens unless we need them
41                    if !matches!(token_type, TokenType::Whitespace) {
42                        let token = Token::new(
43                            token_type,
44                            self.position,
45                            self.line,
46                            self.column,
47                        );
48                        tokens.push(token);
49                    }
50
51                    // Update position tracking
52                    for ch in remaining[..consumed].chars() {
53                        self.position += 1;
54                        if ch == '\n' {
55                            self.line += 1;
56                            self.column = 1;
57                        } else {
58                            self.column += 1;
59                        }
60                    }
61
62                    remaining = rest;
63                }
64                Err(_) => {
65                    return Err(DdbError::ParseError(format!(
66                        "Unexpected character at line {}, column {}: '{}'",
67                        self.line,
68                        self.column,
69                        remaining.chars().next().unwrap_or(' ')
70                    )));
71                }
72            }
73        }
74
75        // Add EOF token
76        tokens.push(Token::new(TokenType::Eof, self.position, self.line, self.column));
77
78        Ok(tokens)
79    }
80
81    fn parse_token<'a>(&self, input: &'a str) -> IResult<&'a str, TokenType> {
82        alt((
83            parse_whitespace,
84            parse_comment,
85            parse_multi_word_keyword,
86            parse_keyword,
87            parse_operator,
88            parse_delimiter,
89            parse_string,
90            parse_number,
91            parse_identifier,
92        ))(input)
93    }
94}
95
96// Free functions for nom parsers
97fn parse_whitespace(input: &str) -> IResult<&str, TokenType> {
98    map(multispace1, |_| TokenType::Whitespace)(input)
99}
100
101fn parse_comment(input: &str) -> IResult<&str, TokenType> {
102    alt((
103        // Line comment: -- comment
104        map(
105            preceded(tag("--"), is_not("\n")),
106            |s: &str| TokenType::Comment(s.to_string())
107        ),
108        // Block comment: /* comment */
109        map(
110            delimited(tag("/*"), is_not("*/"), tag("*/")),
111            |s: &str| TokenType::Comment(s.to_string())
112        ),
113    ))(input)
114}
115
116fn parse_keyword(input: &str) -> IResult<&str, TokenType> {
117    // Try to match an identifier first
118    let (remaining, word) = recognize(pair(
119        alt((alpha1, tag("_"))),
120        many0(alt((alphanumeric1, tag("_"))))
121    ))(input)?;
122
123    // Check for word boundary
124    if let Some(ch) = remaining.chars().next() {
125        if ch.is_alphanumeric() || ch == '_' {
126            return Err(nom::Err::Error(nom::error::Error::new(
127                input,
128                nom::error::ErrorKind::Tag,
129            )));
130        }
131    }
132
133    // Convert to lowercase and match against keywords
134    let token_type = match word.to_lowercase().as_str() {
135        "select" => TokenType::Select,
136        "insert" => TokenType::Insert,
137        "update" => TokenType::Update,
138        "delete" => TokenType::Delete,
139        "upsert" => TokenType::Upsert,
140        "create" => TokenType::Create,
141        "drop" => TokenType::Drop,
142        "use" => TokenType::Use,
143        "show" => TokenType::Show,
144        "set" => TokenType::Set,
145        "from" => TokenType::From,
146        "where" => TokenType::Where,
147        "limit" => TokenType::Limit,
148        "distinct" => TokenType::Distinct,
149        "as" => TokenType::As,
150        "into" => TokenType::Into,
151        "values" => TokenType::Values,
152        "table" => TokenType::Table,
153        "tables" => TokenType::Tables,
154        "columns" => TokenType::Columns,
155        "database" => TokenType::Database,
156        "begin" => TokenType::Begin,
157        "commit" => TokenType::Commit,
158        "rollback" => TokenType::Rollback,
159        "describe" => TokenType::Describe,
160        "and" => TokenType::And,
161        "or" => TokenType::Or,
162        "not" => TokenType::Not,
163        "like" => TokenType::Like,
164        "is" => TokenType::Is,
165        "null" => TokenType::Null,
166        "true" => TokenType::True,
167        "false" => TokenType::False,
168        "in" => TokenType::In,
169        "between" => TokenType::Between,
170        "asc" => TokenType::Asc,
171        "desc" => TokenType::Desc,
172        "order" => TokenType::Order,
173        "by" => TokenType::By,
174        "file" => TokenType::File,
175        "delimiter" => TokenType::Delimiter,
176        "variables" => TokenType::Variables,
177        "offset" => TokenType::Offset,
178        "join" => TokenType::Join,
179        "inner" => TokenType::Inner,
180        "left" => TokenType::Left,
181        "right" => TokenType::Right,
182        "full" => TokenType::Full,
183        "outer" => TokenType::Outer,
184        "on" => TokenType::On,
185        "having" => TokenType::Having,
186        "group" => TokenType::Group,
187        "if" => TokenType::If,
188        "exists" => TokenType::Exists,
189        _ => {
190            // Not a keyword, fail to match
191            return Err(nom::Err::Error(nom::error::Error::new(
192                input,
193                nom::error::ErrorKind::Tag,
194            )));
195        }
196    };
197
198    Ok((remaining, token_type))
199}
200
201fn parse_multi_word_keyword(input: &str) -> IResult<&str, TokenType> {
202    alt((
203        map(
204            alt((tag("ORDER BY"), tag("order by"), tag("Order By"))),
205            |_| TokenType::OrderBy
206        ),
207        map(
208            alt((tag("GROUP BY"), tag("group by"), tag("Group By"))),
209            |_| TokenType::GroupBy
210        ),
211    ))(input)
212}
213
214fn parse_operator(input: &str) -> IResult<&str, TokenType> {
215    alt((
216        value(TokenType::LessEqual, tag("<=")),
217        value(TokenType::GreaterEqual, tag(">=")),
218        value(TokenType::NotEqual, alt((tag("!="), tag("<>")))),
219        value(TokenType::Equal, tag("=")),
220        value(TokenType::LessThan, tag("<")),
221        value(TokenType::GreaterThan, tag(">")),
222    ))(input)
223}
224
225fn parse_delimiter(input: &str) -> IResult<&str, TokenType> {
226    alt((
227        value(TokenType::Comma, char(',')),
228        value(TokenType::Semicolon, char(';')),
229        value(TokenType::LeftParen, char('(')),
230        value(TokenType::RightParen, char(')')),
231        value(TokenType::LeftBracket, char('[')),
232        value(TokenType::RightBracket, char(']')),
233        value(TokenType::Dot, char('.')),
234        value(TokenType::Star, char('*')),
235    ))(input)
236}
237
238fn parse_string(input: &str) -> IResult<&str, TokenType> {
239    alt((
240        // Single-quoted string
241        map(
242            delimited(char('\''), take_while(|c| c != '\''), char('\'')),
243            |s: &str| TokenType::String(s.to_string())
244        ),
245        // Double-quoted string
246        map(
247            delimited(char('"'), take_while(|c| c != '"'), char('"')),
248            |s: &str| TokenType::String(s.to_string())
249        ),
250    ))(input)
251}
252
253fn parse_number(input: &str) -> IResult<&str, TokenType> {
254    map(
255        recognize(pair(
256            digit1,
257            opt(pair(char('.'), digit1))
258        )),
259        |s: &str| TokenType::Number(s.to_string())
260    )(input)
261}
262
263fn parse_identifier(input: &str) -> IResult<&str, TokenType> {
264    map(
265        recognize(pair(
266            alt((
267                tag("@@"),  // System variables (@@VERSION, @@ROWCOUNT, etc.)
268                alpha1,
269                tag("_")
270            )),
271            many0(alt((alphanumeric1, tag("_"))))
272        )),
273        |s: &str| TokenType::Identifier(s.to_string())
274    )(input)
275}
276
277impl Default for Tokenizer {
278    fn default() -> Self {
279        Self::new()
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    #[test]
288    fn test_tokenize_select() {
289        let mut tokenizer = Tokenizer::new();
290        let tokens = tokenizer.tokenize("SELECT * FROM users").unwrap();
291
292        assert_eq!(tokens[0].token_type, TokenType::Select);
293        assert_eq!(tokens[1].token_type, TokenType::Star);
294        assert_eq!(tokens[2].token_type, TokenType::From);
295        assert!(matches!(tokens[3].token_type, TokenType::Identifier(_)));
296    }
297
298    #[test]
299    fn test_tokenize_where_clause() {
300        let mut tokenizer = Tokenizer::new();
301        let tokens = tokenizer.tokenize("WHERE id = 123 AND name LIKE 'test%'").unwrap();
302
303        assert_eq!(tokens[0].token_type, TokenType::Where);
304        assert!(matches!(tokens[1].token_type, TokenType::Identifier(_)));
305        assert_eq!(tokens[2].token_type, TokenType::Equal);
306        assert!(matches!(tokens[3].token_type, TokenType::Number(_)));
307    }
308
309    #[test]
310    fn test_tokenize_strings() {
311        let mut tokenizer = Tokenizer::new();
312        let tokens = tokenizer.tokenize("'single' \"double\"").unwrap();
313
314        assert!(matches!(tokens[0].token_type, TokenType::String(ref s) if s == "single"));
315        assert!(matches!(tokens[1].token_type, TokenType::String(ref s) if s == "double"));
316    }
317}