Skip to main content

mentedb_query/
lexer.rs

1//! Hand-written lexer for MQL.
2
3use mentedb_core::error::{MenteError, MenteResult};
4
5#[derive(Debug, Clone, PartialEq)]
6pub struct Token {
7    pub kind: TokenKind,
8    pub lexeme: String,
9    pub position: usize,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13pub enum TokenKind {
14    // Statements
15    Recall,
16    Relate,
17    Forget,
18    Consolidate,
19    Traverse,
20
21    // Clauses
22    Where,
23    And,
24    Or,
25    Not,
26    Near,
27    Within,
28    Limit,
29    OrderBy,
30    As,
31    Of,
32    From,
33    To,
34    With,
35
36    // Keywords
37    Agent,
38    Space,
39    Type,
40    Tag,
41    Salience,
42    Confidence,
43    Created,
44    Accessed,
45    Depth,
46    Hops,
47    Memories,
48    By,
49    EdgeType,
50
51    // Operators
52    Eq,        // =
53    Neq,       // !=
54    Gt,        // >
55    Lt,        // <
56    Gte,       // >=
57    Lte,       // <=
58    SimilarTo, // ~>
59    Arrow,     // ->
60
61    // Punctuation
62    LParen,
63    RParen,
64    LBracket,
65    RBracket,
66    Comma,
67    Dot,
68    Colon,
69    Semicolon,
70
71    // Literals
72    StringLit,
73    IntegerLit,
74    FloatLit,
75    Identifier,
76    UuidLit,
77
78    Eof,
79}
80
81pub fn tokenize(input: &str) -> MenteResult<Vec<Token>> {
82    let mut tokens = Vec::new();
83    let bytes = input.as_bytes();
84    let len = bytes.len();
85    let mut pos = 0;
86
87    while pos < len {
88        // Skip whitespace
89        if bytes[pos].is_ascii_whitespace() {
90            pos += 1;
91            continue;
92        }
93
94        let start = pos;
95
96        // String literal
97        if bytes[pos] == b'"' {
98            pos += 1;
99            while pos < len && bytes[pos] != b'"' {
100                if bytes[pos] == b'\\' {
101                    pos += 1; // skip escaped char
102                }
103                pos += 1;
104            }
105            if pos >= len {
106                return Err(MenteError::Query("unterminated string literal".into()));
107            }
108            pos += 1; // closing quote
109            let lexeme = input[start..pos].to_string();
110            tokens.push(Token {
111                kind: TokenKind::StringLit,
112                lexeme,
113                position: start,
114            });
115            continue;
116        }
117
118        // Two-char operators
119        if pos + 1 < len {
120            let two = &input[start..start + 2];
121            let kind = match two {
122                "!=" => Some(TokenKind::Neq),
123                ">=" => Some(TokenKind::Gte),
124                "<=" => Some(TokenKind::Lte),
125                "~>" => Some(TokenKind::SimilarTo),
126                "->" => Some(TokenKind::Arrow),
127                _ => None,
128            };
129            if let Some(k) = kind {
130                tokens.push(Token {
131                    kind: k,
132                    lexeme: two.to_string(),
133                    position: start,
134                });
135                pos += 2;
136                continue;
137            }
138        }
139
140        // Single-char operators/punctuation
141        let single = match bytes[pos] {
142            b'=' => Some(TokenKind::Eq),
143            b'>' => Some(TokenKind::Gt),
144            b'<' => Some(TokenKind::Lt),
145            b'(' => Some(TokenKind::LParen),
146            b')' => Some(TokenKind::RParen),
147            b'[' => Some(TokenKind::LBracket),
148            b']' => Some(TokenKind::RBracket),
149            b',' => Some(TokenKind::Comma),
150            b'.' => Some(TokenKind::Dot),
151            b':' => Some(TokenKind::Colon),
152            b';' => Some(TokenKind::Semicolon),
153            _ => None,
154        };
155        if let Some(k) = single {
156            tokens.push(Token {
157                kind: k,
158                lexeme: input[start..start + 1].to_string(),
159                position: start,
160            });
161            pos += 1;
162            continue;
163        }
164
165        // Try UUID first: if we see a hex digit, speculatively scan for UUID pattern
166        if bytes[pos].is_ascii_hexdigit() {
167            let saved = pos;
168            // Consume alphanumeric + hyphens to check for UUID
169            while pos < len
170                && (bytes[pos].is_ascii_alphanumeric() || bytes[pos] == b'_' || bytes[pos] == b'-')
171            {
172                pos += 1;
173            }
174            let candidate = &input[saved..pos];
175            if is_uuid_like(candidate) {
176                tokens.push(Token {
177                    kind: TokenKind::UuidLit,
178                    lexeme: candidate.to_string(),
179                    position: start,
180                });
181                continue;
182            }
183            // Not a UUID — reset and fall through to number/identifier parsing
184            pos = saved;
185        }
186
187        // Numbers (may start with - for negative)
188        if bytes[pos].is_ascii_digit()
189            || (bytes[pos] == b'-' && pos + 1 < len && bytes[pos + 1].is_ascii_digit())
190        {
191            if bytes[pos] == b'-' {
192                pos += 1;
193            }
194            while pos < len && bytes[pos].is_ascii_digit() {
195                pos += 1;
196            }
197            let mut is_float = false;
198            if pos < len && bytes[pos] == b'.' && pos + 1 < len && bytes[pos + 1].is_ascii_digit() {
199                is_float = true;
200                pos += 1;
201                while pos < len && bytes[pos].is_ascii_digit() {
202                    pos += 1;
203                }
204            }
205            let lexeme = input[start..pos].to_string();
206            let kind = if is_float {
207                TokenKind::FloatLit
208            } else {
209                TokenKind::IntegerLit
210            };
211            tokens.push(Token {
212                kind,
213                lexeme,
214                position: start,
215            });
216            continue;
217        }
218
219        // Identifiers, keywords
220        if bytes[pos].is_ascii_alphanumeric() || bytes[pos] == b'_' {
221            while pos < len
222                && (bytes[pos].is_ascii_alphanumeric() || bytes[pos] == b'_' || bytes[pos] == b'-')
223            {
224                pos += 1;
225            }
226            let lexeme = input[start..pos].to_string();
227
228            let kind = match lexeme.to_lowercase().as_str() {
229                "recall" => TokenKind::Recall,
230                "relate" => TokenKind::Relate,
231                "forget" => TokenKind::Forget,
232                "consolidate" => TokenKind::Consolidate,
233                "traverse" => TokenKind::Traverse,
234                "where" => TokenKind::Where,
235                "and" => TokenKind::And,
236                "or" => TokenKind::Or,
237                "not" => TokenKind::Not,
238                "near" => TokenKind::Near,
239                "within" => TokenKind::Within,
240                "limit" => TokenKind::Limit,
241                "order" => TokenKind::OrderBy,
242                "as" => TokenKind::As,
243                "of" => TokenKind::Of,
244                "from" => TokenKind::From,
245                "to" => TokenKind::To,
246                "with" => TokenKind::With,
247                "agent" => TokenKind::Agent,
248                "space" => TokenKind::Space,
249                "type" => TokenKind::Type,
250                "tag" => TokenKind::Tag,
251                "salience" => TokenKind::Salience,
252                "confidence" => TokenKind::Confidence,
253                "created" => TokenKind::Created,
254                "accessed" => TokenKind::Accessed,
255                "depth" => TokenKind::Depth,
256                "hops" => TokenKind::Hops,
257                "memories" => TokenKind::Memories,
258                "by" => TokenKind::By,
259                "edge_type" => TokenKind::EdgeType,
260                _ => TokenKind::Identifier,
261            };
262            tokens.push(Token {
263                kind,
264                lexeme,
265                position: start,
266            });
267            continue;
268        }
269
270        return Err(MenteError::Query(format!(
271            "unexpected character '{}' at position {}",
272            bytes[pos] as char, pos
273        )));
274    }
275
276    tokens.push(Token {
277        kind: TokenKind::Eof,
278        lexeme: String::new(),
279        position: pos,
280    });
281    Ok(tokens)
282}
283
284fn is_uuid_like(s: &str) -> bool {
285    // UUID format: 8-4-4-4-12 hex chars (with dashes)
286    if s.len() != 36 {
287        return false;
288    }
289    let parts: Vec<&str> = s.split('-').collect();
290    if parts.len() != 5 {
291        return false;
292    }
293    let expected_lens = [8, 4, 4, 4, 12];
294    for (part, &expected) in parts.iter().zip(&expected_lens) {
295        if part.len() != expected || !part.chars().all(|c| c.is_ascii_hexdigit()) {
296            return false;
297        }
298    }
299    true
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    #[test]
307    fn test_recall_statement_tokens() {
308        let tokens = tokenize("RECALL memories WHERE type = episodic LIMIT 10").unwrap();
309        assert_eq!(tokens[0].kind, TokenKind::Recall);
310        assert_eq!(tokens[1].kind, TokenKind::Memories);
311        assert_eq!(tokens[2].kind, TokenKind::Where);
312        assert_eq!(tokens[3].kind, TokenKind::Type);
313        assert_eq!(tokens[4].kind, TokenKind::Eq);
314        assert_eq!(tokens[5].kind, TokenKind::Identifier);
315        assert_eq!(tokens[5].lexeme, "episodic");
316        assert_eq!(tokens[6].kind, TokenKind::Limit);
317        assert_eq!(tokens[7].kind, TokenKind::IntegerLit);
318        assert_eq!(tokens[8].kind, TokenKind::Eof);
319    }
320
321    #[test]
322    fn test_string_literal() {
323        let tokens = tokenize(r#"content ~> "database migration""#).unwrap();
324        assert_eq!(tokens[0].kind, TokenKind::Identifier);
325        assert_eq!(tokens[1].kind, TokenKind::SimilarTo);
326        assert_eq!(tokens[2].kind, TokenKind::StringLit);
327        assert_eq!(tokens[2].lexeme, r#""database migration""#);
328    }
329
330    #[test]
331    fn test_operators() {
332        let tokens = tokenize("= != > < >= <= ~> ->").unwrap();
333        let kinds: Vec<TokenKind> = tokens.iter().map(|t| t.kind).collect();
334        assert_eq!(
335            kinds,
336            vec![
337                TokenKind::Eq,
338                TokenKind::Neq,
339                TokenKind::Gt,
340                TokenKind::Lt,
341                TokenKind::Gte,
342                TokenKind::Lte,
343                TokenKind::SimilarTo,
344                TokenKind::Arrow,
345                TokenKind::Eof,
346            ]
347        );
348    }
349
350    #[test]
351    fn test_uuid_token() {
352        let tokens = tokenize("550e8400-e29b-41d4-a716-446655440000").unwrap();
353        assert_eq!(tokens[0].kind, TokenKind::UuidLit);
354    }
355
356    #[test]
357    fn test_float_literal() {
358        let tokens = tokenize("0.1 42 3.14").unwrap();
359        assert_eq!(tokens[0].kind, TokenKind::FloatLit);
360        assert_eq!(tokens[1].kind, TokenKind::IntegerLit);
361        assert_eq!(tokens[2].kind, TokenKind::FloatLit);
362    }
363
364    #[test]
365    fn test_vector_literal() {
366        let tokens = tokenize("[0.1, 0.2, 0.3]").unwrap();
367        assert_eq!(tokens[0].kind, TokenKind::LBracket);
368        assert_eq!(tokens[1].kind, TokenKind::FloatLit);
369        assert_eq!(tokens[2].kind, TokenKind::Comma);
370        assert_eq!(tokens[5].kind, TokenKind::FloatLit);
371        assert_eq!(tokens[6].kind, TokenKind::RBracket);
372    }
373
374    #[test]
375    fn test_punctuation() {
376        let tokens = tokenize("( ) [ ] , . : ;").unwrap();
377        let kinds: Vec<TokenKind> = tokens.iter().map(|t| t.kind).collect();
378        assert_eq!(
379            kinds,
380            vec![
381                TokenKind::LParen,
382                TokenKind::RParen,
383                TokenKind::LBracket,
384                TokenKind::RBracket,
385                TokenKind::Comma,
386                TokenKind::Dot,
387                TokenKind::Colon,
388                TokenKind::Semicolon,
389                TokenKind::Eof,
390            ]
391        );
392    }
393}