1use 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 Recall,
16 Relate,
17 Forget,
18 Consolidate,
19 Traverse,
20
21 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 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 Eq, Neq, Gt, Lt, Gte, Lte, SimilarTo, Arrow, LParen,
63 RParen,
64 LBracket,
65 RBracket,
66 Comma,
67 Dot,
68 Colon,
69 Semicolon,
70
71 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 if bytes[pos].is_ascii_whitespace() {
90 pos += 1;
91 continue;
92 }
93
94 let start = pos;
95
96 if bytes[pos] == b'"' {
98 pos += 1;
99 while pos < len && bytes[pos] != b'"' {
100 if bytes[pos] == b'\\' {
101 pos += 1; }
103 pos += 1;
104 }
105 if pos >= len {
106 return Err(MenteError::Query("unterminated string literal".into()));
107 }
108 pos += 1; 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 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 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 if bytes[pos].is_ascii_hexdigit() {
167 let saved = pos;
168 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 pos = saved;
185 }
186
187 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 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 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}