Skip to main content

mentedb_query/
parser.rs

1//! Hand-written recursive descent parser for MQL.
2
3use mentedb_core::edge::EdgeType;
4use mentedb_core::error::{MenteError, MenteResult};
5use mentedb_core::memory::MemoryType;
6use uuid::Uuid;
7
8use crate::ast::*;
9use crate::lexer::{Token, TokenKind};
10use mentedb_core::types::MemoryId;
11
12pub struct Parser<'a> {
13    tokens: &'a [Token],
14    pos: usize,
15}
16
17impl<'a> Parser<'a> {
18    pub fn new(tokens: &'a [Token]) -> Self {
19        Self { tokens, pos: 0 }
20    }
21
22    pub fn parse(tokens: &[Token]) -> MenteResult<Statement> {
23        let mut parser = Parser::new(tokens);
24        parser.parse_statement()
25    }
26
27    fn peek(&self) -> &Token {
28        &self.tokens[self.pos.min(self.tokens.len() - 1)]
29    }
30
31    fn advance(&mut self) -> &Token {
32        let tok = &self.tokens[self.pos.min(self.tokens.len() - 1)];
33        if self.pos < self.tokens.len() {
34            self.pos += 1;
35        }
36        tok
37    }
38
39    fn expect(&mut self, kind: TokenKind) -> MenteResult<&Token> {
40        let tok = self.peek();
41        if tok.kind != kind {
42            return Err(MenteError::Query(format!(
43                "expected {:?}, found {:?} ('{}') at position {}",
44                kind, tok.kind, tok.lexeme, tok.position
45            )));
46        }
47        Ok(self.advance())
48    }
49
50    fn at(&self, kind: TokenKind) -> bool {
51        self.peek().kind == kind
52    }
53
54    fn parse_statement(&mut self) -> MenteResult<Statement> {
55        match self.peek().kind {
56            TokenKind::Recall => self.parse_recall(),
57            TokenKind::Relate => self.parse_relate(),
58            TokenKind::Forget => self.parse_forget(),
59            TokenKind::Consolidate => self.parse_consolidate(),
60            TokenKind::Traverse => self.parse_traverse(),
61            _ => Err(MenteError::Query(format!(
62                "expected statement keyword, found {:?} at position {}",
63                self.peek().kind,
64                self.peek().position
65            ))),
66        }
67    }
68
69    fn parse_recall(&mut self) -> MenteResult<Statement> {
70        self.advance(); // RECALL
71
72        // Optional "memories" keyword
73        if self.at(TokenKind::Memories) {
74            self.advance();
75        }
76
77        let mut near = None;
78        let mut filters = Vec::new();
79        let mut limit = None;
80        let mut order_by = None;
81
82        // NEAR [vector]
83        if self.at(TokenKind::Near) {
84            self.advance();
85            near = Some(self.parse_vector()?);
86        }
87
88        // WHERE clause
89        if self.at(TokenKind::Where) {
90            self.advance();
91            filters = self.parse_filters()?;
92        }
93
94        // ORDER BY field
95        if self.at(TokenKind::OrderBy) {
96            self.advance();
97            // consume optional "BY"
98            if self.at(TokenKind::By) {
99                self.advance();
100            }
101            let field = self.parse_field()?;
102            let descending = false; // default ascending
103            order_by = Some(OrderBy { field, descending });
104        }
105
106        // LIMIT n
107        if self.at(TokenKind::Limit) {
108            self.advance();
109            let tok = self.advance();
110            let n: usize = tok
111                .lexeme
112                .parse()
113                .map_err(|_| MenteError::Query(format!("invalid limit value: {}", tok.lexeme)))?;
114            limit = Some(n);
115        }
116
117        // AS OF <timestamp> — point-in-time temporal filter. Added as a ValidAt
118        // filter so it flows through the same pipeline as WHERE clauses.
119        if self.at(TokenKind::As) {
120            self.advance();
121            self.expect(TokenKind::Of)?;
122            let tok = self.advance();
123            let t: i64 = tok.lexeme.parse().map_err(|_| {
124                MenteError::Query(format!("invalid AS OF timestamp: {}", tok.lexeme))
125            })?;
126            filters.push(Filter {
127                field: Field::ValidAt,
128                op: Operator::Eq,
129                value: Value::Integer(t),
130            });
131        }
132
133        Ok(Statement::Recall(RecallStatement {
134            filters,
135            near,
136            limit,
137            order_by,
138        }))
139    }
140
141    fn parse_relate(&mut self) -> MenteResult<Statement> {
142        self.advance(); // RELATE
143
144        let source = self.parse_uuid()?;
145        self.expect(TokenKind::Arrow)?;
146        let target = self.parse_uuid()?;
147        self.expect(TokenKind::As)?;
148        let edge_type = self.parse_edge_type()?;
149
150        let mut weight = None;
151        if self.at(TokenKind::With) {
152            self.advance();
153            // expect "weight = <float>"
154            self.expect(TokenKind::Identifier)?; // "weight"
155            self.expect(TokenKind::Eq)?;
156            let tok = self.advance();
157            let w: f32 = tok
158                .lexeme
159                .parse()
160                .map_err(|_| MenteError::Query(format!("invalid weight value: {}", tok.lexeme)))?;
161            weight = Some(w);
162        }
163
164        Ok(Statement::Relate(RelateStatement {
165            source,
166            target,
167            edge_type,
168            weight,
169        }))
170    }
171
172    fn parse_forget(&mut self) -> MenteResult<Statement> {
173        self.advance(); // FORGET
174        let target = self.parse_uuid()?;
175        Ok(Statement::Forget(ForgetStatement { target }))
176    }
177
178    fn parse_consolidate(&mut self) -> MenteResult<Statement> {
179        self.advance(); // CONSOLIDATE
180        let mut filters = Vec::new();
181        if self.at(TokenKind::Where) {
182            self.advance();
183            filters = self.parse_filters()?;
184        }
185        Ok(Statement::Consolidate(ConsolidateStatement { filters }))
186    }
187
188    fn parse_traverse(&mut self) -> MenteResult<Statement> {
189        self.advance(); // TRAVERSE
190        let start = self.parse_uuid()?;
191
192        self.expect(TokenKind::Depth)?;
193        let tok = self.advance();
194        let depth: usize = tok
195            .lexeme
196            .parse()
197            .map_err(|_| MenteError::Query(format!("invalid depth value: {}", tok.lexeme)))?;
198
199        let mut edge_filter = None;
200        if self.at(TokenKind::Where) {
201            self.advance();
202            // edge_type = <type>
203            self.expect(TokenKind::EdgeType)?;
204            self.expect(TokenKind::Eq)?;
205            let et = self.parse_edge_type()?;
206            edge_filter = Some(vec![et]);
207        }
208
209        Ok(Statement::Traverse(TraverseStatement {
210            start,
211            depth,
212            edge_filter,
213        }))
214    }
215
216    fn parse_filters(&mut self) -> MenteResult<Vec<Filter>> {
217        let mut filters = vec![self.parse_filter()?];
218        while self.at(TokenKind::And) {
219            self.advance();
220            filters.push(self.parse_filter()?);
221        }
222        Ok(filters)
223    }
224
225    fn parse_filter(&mut self) -> MenteResult<Filter> {
226        let field = self.parse_field()?;
227        let op = self.parse_operator()?;
228        let value = self.parse_value(&field)?;
229        Ok(Filter { field, op, value })
230    }
231
232    fn parse_field(&mut self) -> MenteResult<Field> {
233        let tok = self.advance();
234        match tok.kind {
235            TokenKind::Identifier if tok.lexeme.eq_ignore_ascii_case("content") => {
236                Ok(Field::Content)
237            }
238            TokenKind::Type => Ok(Field::Type),
239            TokenKind::Tag => Ok(Field::Tag),
240            TokenKind::Agent => Ok(Field::Agent),
241            TokenKind::Space => Ok(Field::Space),
242            TokenKind::Salience => Ok(Field::Salience),
243            TokenKind::Confidence => Ok(Field::Confidence),
244            TokenKind::Created => Ok(Field::Created),
245            TokenKind::Accessed => Ok(Field::Accessed),
246            _ => Err(MenteError::Query(format!(
247                "expected field name, found '{}' at position {}",
248                tok.lexeme, tok.position
249            ))),
250        }
251    }
252
253    fn parse_operator(&mut self) -> MenteResult<Operator> {
254        let tok = self.advance();
255        match tok.kind {
256            TokenKind::Eq => Ok(Operator::Eq),
257            TokenKind::Neq => Ok(Operator::Neq),
258            TokenKind::Gt => Ok(Operator::Gt),
259            TokenKind::Lt => Ok(Operator::Lt),
260            TokenKind::Gte => Ok(Operator::Gte),
261            TokenKind::Lte => Ok(Operator::Lte),
262            TokenKind::SimilarTo => Ok(Operator::SimilarTo),
263            _ => Err(MenteError::Query(format!(
264                "expected operator, found '{}' at position {}",
265                tok.lexeme, tok.position
266            ))),
267        }
268    }
269
270    fn parse_value(&mut self, field: &Field) -> MenteResult<Value> {
271        // For Type field, parse as MemoryType
272        if *field == Field::Type {
273            return self.parse_memory_type_value();
274        }
275
276        let tok = self.advance();
277        match tok.kind {
278            TokenKind::StringLit => {
279                // Strip surrounding quotes
280                let inner = tok.lexeme[1..tok.lexeme.len() - 1].to_string();
281                // Check if this looks like a UUID inside quotes
282                if let Ok(uuid) = inner.parse::<MemoryId>() {
283                    return Ok(Value::Uuid(uuid.into()));
284                }
285                Ok(Value::Text(inner))
286            }
287            TokenKind::IntegerLit => {
288                let n: i64 = tok
289                    .lexeme
290                    .parse()
291                    .map_err(|_| MenteError::Query(format!("invalid integer: {}", tok.lexeme)))?;
292                Ok(Value::Integer(n))
293            }
294            TokenKind::FloatLit => {
295                let n: f64 = tok
296                    .lexeme
297                    .parse()
298                    .map_err(|_| MenteError::Query(format!("invalid float: {}", tok.lexeme)))?;
299                Ok(Value::Number(n))
300            }
301            TokenKind::UuidLit => {
302                let uuid: Uuid = tok
303                    .lexeme
304                    .parse()
305                    .map_err(|_| MenteError::Query(format!("invalid UUID: {}", tok.lexeme)))?;
306                Ok(Value::Uuid(uuid))
307            }
308            TokenKind::Identifier => {
309                let lower = tok.lexeme.to_lowercase();
310                match lower.as_str() {
311                    "true" => Ok(Value::Bool(true)),
312                    "false" => Ok(Value::Bool(false)),
313                    _ => Ok(Value::Text(tok.lexeme.clone())),
314                }
315            }
316            TokenKind::LBracket => {
317                // put back and parse as vector
318                self.pos -= 1;
319                let v = self.parse_vector()?;
320                Ok(Value::Vector(v))
321            }
322            _ => Err(MenteError::Query(format!(
323                "expected value, found '{}' at position {}",
324                tok.lexeme, tok.position
325            ))),
326        }
327    }
328
329    fn parse_memory_type_value(&mut self) -> MenteResult<Value> {
330        let tok = self.advance();
331        let name = match tok.kind {
332            TokenKind::Identifier | TokenKind::StringLit => {
333                if tok.kind == TokenKind::StringLit {
334                    tok.lexeme[1..tok.lexeme.len() - 1].to_string()
335                } else {
336                    tok.lexeme.clone()
337                }
338            }
339            _ => {
340                return Err(MenteError::Query(format!(
341                    "expected memory type, found '{}' at position {}",
342                    tok.lexeme, tok.position
343                )));
344            }
345        };
346
347        let mt = match name.to_lowercase().as_str() {
348            "episodic" => MemoryType::Episodic,
349            "semantic" => MemoryType::Semantic,
350            "procedural" => MemoryType::Procedural,
351            "antipattern" | "anti_pattern" => MemoryType::AntiPattern,
352            "reasoning" => MemoryType::Reasoning,
353            "correction" => MemoryType::Correction,
354            _ => {
355                return Err(MenteError::Query(format!("unknown memory type: {}", name)));
356            }
357        };
358        Ok(Value::MemoryType(mt))
359    }
360
361    fn parse_edge_type(&mut self) -> MenteResult<EdgeType> {
362        let tok = self.advance();
363        let name = match tok.kind {
364            TokenKind::Identifier | TokenKind::StringLit => {
365                if tok.kind == TokenKind::StringLit {
366                    tok.lexeme[1..tok.lexeme.len() - 1].to_string()
367                } else {
368                    tok.lexeme.clone()
369                }
370            }
371            _ => {
372                return Err(MenteError::Query(format!(
373                    "expected edge type, found '{}' at position {}",
374                    tok.lexeme, tok.position
375                )));
376            }
377        };
378
379        match name.to_lowercase().as_str() {
380            "caused" => Ok(EdgeType::Caused),
381            "before" => Ok(EdgeType::Before),
382            "related" => Ok(EdgeType::Related),
383            "contradicts" => Ok(EdgeType::Contradicts),
384            "supports" => Ok(EdgeType::Supports),
385            "supersedes" => Ok(EdgeType::Supersedes),
386            "derived" => Ok(EdgeType::Derived),
387            "partof" | "part_of" => Ok(EdgeType::PartOf),
388            _ => Err(MenteError::Query(format!("unknown edge type: {}", name))),
389        }
390    }
391
392    fn parse_uuid(&mut self) -> MenteResult<MemoryId> {
393        let tok = self.advance();
394        match tok.kind {
395            TokenKind::UuidLit => tok
396                .lexeme
397                .parse()
398                .map_err(|_| MenteError::Query(format!("invalid UUID: {}", tok.lexeme))),
399            TokenKind::StringLit => {
400                let inner = &tok.lexeme[1..tok.lexeme.len() - 1];
401                inner.parse().map_err(|_| {
402                    MenteError::Query(format!("invalid UUID in string: {}", tok.lexeme))
403                })
404            }
405            _ => Err(MenteError::Query(format!(
406                "expected UUID, found '{}' at position {}",
407                tok.lexeme, tok.position
408            ))),
409        }
410    }
411
412    fn parse_vector(&mut self) -> MenteResult<Vec<f32>> {
413        self.expect(TokenKind::LBracket)?;
414        let mut values = Vec::new();
415        if !self.at(TokenKind::RBracket) {
416            let tok = self.advance();
417            let v: f32 = tok.lexeme.parse().map_err(|_| {
418                MenteError::Query(format!("invalid float in vector: {}", tok.lexeme))
419            })?;
420            values.push(v);
421            while self.at(TokenKind::Comma) {
422                self.advance();
423                let tok = self.advance();
424                let v: f32 = tok.lexeme.parse().map_err(|_| {
425                    MenteError::Query(format!("invalid float in vector: {}", tok.lexeme))
426                })?;
427                values.push(v);
428            }
429        }
430        self.expect(TokenKind::RBracket)?;
431        Ok(values)
432    }
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438    use crate::lexer::tokenize;
439
440    #[test]
441    fn test_parse_recall_with_type_filter() {
442        let tokens = tokenize("RECALL memories WHERE type = episodic LIMIT 5").unwrap();
443        let stmt = Parser::parse(&tokens).unwrap();
444        match stmt {
445            Statement::Recall(r) => {
446                assert_eq!(r.filters.len(), 1);
447                assert_eq!(r.filters[0].field, Field::Type);
448                assert_eq!(r.filters[0].value, Value::MemoryType(MemoryType::Episodic));
449                assert_eq!(r.limit, Some(5));
450            }
451            _ => panic!("expected Recall"),
452        }
453    }
454
455    #[test]
456    fn test_parse_recall_as_of() {
457        // AS OF <t> lowers to a ValidAt filter carrying the timestamp, on top of
458        // any WHERE filters, and coexists with LIMIT.
459        let tokens =
460            tokenize("RECALL memories WHERE type = semantic LIMIT 5 AS OF 1700000000").unwrap();
461        let stmt = Parser::parse(&tokens).unwrap();
462        match stmt {
463            Statement::Recall(r) => {
464                assert_eq!(r.limit, Some(5));
465                assert_eq!(r.filters.len(), 2);
466                let valid_at = r
467                    .filters
468                    .iter()
469                    .find(|f| f.field == Field::ValidAt)
470                    .expect("expected a ValidAt filter from AS OF");
471                assert_eq!(valid_at.value, Value::Integer(1_700_000_000));
472            }
473            _ => panic!("expected Recall"),
474        }
475    }
476
477    #[test]
478    fn test_parse_recall_as_of_without_where() {
479        let tokens = tokenize("RECALL memories AS OF 42").unwrap();
480        let stmt = Parser::parse(&tokens).unwrap();
481        match stmt {
482            Statement::Recall(r) => {
483                assert_eq!(r.filters.len(), 1);
484                assert_eq!(r.filters[0].field, Field::ValidAt);
485                assert_eq!(r.filters[0].value, Value::Integer(42));
486            }
487            _ => panic!("expected Recall"),
488        }
489    }
490
491    #[test]
492    fn test_parse_recall_similar_to() {
493        let tokens =
494            tokenize(r#"RECALL memories WHERE content ~> "database migration" LIMIT 10"#).unwrap();
495        let stmt = Parser::parse(&tokens).unwrap();
496        match stmt {
497            Statement::Recall(r) => {
498                assert_eq!(r.filters.len(), 1);
499                assert_eq!(r.filters[0].op, Operator::SimilarTo);
500                assert_eq!(r.limit, Some(10));
501            }
502            _ => panic!("expected Recall"),
503        }
504    }
505
506    #[test]
507    fn test_parse_recall_near() {
508        let tokens = tokenize("RECALL memories NEAR [0.1, 0.2, 0.3] LIMIT 10").unwrap();
509        let stmt = Parser::parse(&tokens).unwrap();
510        match stmt {
511            Statement::Recall(r) => {
512                assert_eq!(r.near, Some(vec![0.1, 0.2, 0.3]));
513                assert_eq!(r.limit, Some(10));
514            }
515            _ => panic!("expected Recall"),
516        }
517    }
518
519    #[test]
520    fn test_parse_relate() {
521        let tokens = tokenize(
522            "RELATE 550e8400-e29b-41d4-a716-446655440000 -> 660e8400-e29b-41d4-a716-446655440000 AS caused WITH weight = 0.9"
523        ).unwrap();
524        let stmt = Parser::parse(&tokens).unwrap();
525        match stmt {
526            Statement::Relate(r) => {
527                assert_eq!(r.edge_type, EdgeType::Caused);
528                assert_eq!(r.weight, Some(0.9));
529            }
530            _ => panic!("expected Relate"),
531        }
532    }
533
534    #[test]
535    fn test_parse_forget() {
536        let tokens = tokenize("FORGET 550e8400-e29b-41d4-a716-446655440000").unwrap();
537        let stmt = Parser::parse(&tokens).unwrap();
538        match stmt {
539            Statement::Forget(f) => {
540                assert_eq!(
541                    f.target,
542                    "550e8400-e29b-41d4-a716-446655440000"
543                        .parse::<MemoryId>()
544                        .unwrap()
545                );
546            }
547            _ => panic!("expected Forget"),
548        }
549    }
550
551    #[test]
552    fn test_parse_consolidate() {
553        let tokens =
554            tokenize(r#"CONSOLIDATE WHERE type = episodic AND accessed < "2024-01-01""#).unwrap();
555        let stmt = Parser::parse(&tokens).unwrap();
556        match stmt {
557            Statement::Consolidate(c) => {
558                assert_eq!(c.filters.len(), 2);
559            }
560            _ => panic!("expected Consolidate"),
561        }
562    }
563
564    #[test]
565    fn test_parse_traverse() {
566        let tokens = tokenize(
567            "TRAVERSE 550e8400-e29b-41d4-a716-446655440000 DEPTH 3 WHERE edge_type = caused",
568        )
569        .unwrap();
570        let stmt = Parser::parse(&tokens).unwrap();
571        match stmt {
572            Statement::Traverse(t) => {
573                assert_eq!(t.depth, 3);
574                assert_eq!(t.edge_filter, Some(vec![EdgeType::Caused]));
575            }
576            _ => panic!("expected Traverse"),
577        }
578    }
579}