lift_ast/
parser.rs

1use crate::ast::*;
2use crate::token::{Span, Token, TokenKind};
3use thiserror::Error;
4
5#[derive(Debug, Error)]
6pub enum ParseError {
7    #[error("Unexpected token {found} at line {line}:{col}, expected {expected}")]
8    UnexpectedToken {
9        found: String,
10        expected: String,
11        line: u32,
12        col: u32,
13    },
14
15    #[error("Unexpected end of file")]
16    UnexpectedEof,
17
18    #[error("Parse error: {0}")]
19    General(String),
20}
21
22pub struct Parser {
23    tokens: Vec<Token>,
24    pos: usize,
25    errors: Vec<ParseError>,
26}
27
28impl Parser {
29    pub fn new(tokens: Vec<Token>) -> Self {
30        Self {
31            tokens,
32            pos: 0,
33            errors: Vec::new(),
34        }
35    }
36
37    pub fn parse(&mut self) -> Result<Program, Vec<ParseError>> {
38        let mut program = Program {
39            dialect_directives: Vec::new(),
40            modules: Vec::new(),
41        };
42
43        while !self.at_end() {
44            match self.peek_kind() {
45                TokenKind::HashDialect(_) => match self.parse_dialect_directive() {
46                    Ok(d) => program.dialect_directives.push(d),
47                    Err(e) => {
48                        self.errors.push(e);
49                        self.recover();
50                    }
51                },
52                TokenKind::Module => match self.parse_module() {
53                    Ok(m) => program.modules.push(m),
54                    Err(e) => {
55                        self.errors.push(e);
56                        self.recover();
57                    }
58                },
59                TokenKind::Eof => break,
60                _ => {
61                    let tok = self.advance();
62                    self.errors.push(ParseError::UnexpectedToken {
63                        found: format!("{}", tok.kind),
64                        expected: "module or #dialect".to_string(),
65                        line: tok.span.line,
66                        col: tok.span.column,
67                    });
68                }
69            }
70        }
71
72        if self.errors.is_empty() {
73            Ok(program)
74        } else {
75            Err(std::mem::take(&mut self.errors))
76        }
77    }
78
79    fn parse_dialect_directive(&mut self) -> Result<DialectDirective, ParseError> {
80        let tok = self.advance();
81        if let TokenKind::HashDialect(name) = &tok.kind {
82            Ok(DialectDirective {
83                name: name.clone(),
84                span: tok.span,
85            })
86        } else {
87            Err(ParseError::General("Expected #dialect directive".into()))
88        }
89    }
90
91    fn parse_module(&mut self) -> Result<ModuleDecl, ParseError> {
92        let start = self.expect_kind(&TokenKind::Module)?.span;
93
94        let name = match &self.advance().kind {
95            TokenKind::AtIdent(n) => n.clone(),
96            other => {
97                return Err(ParseError::UnexpectedToken {
98                    found: format!("{}", other),
99                    expected: "@module_name".into(),
100                    line: start.line,
101                    col: start.column,
102                })
103            }
104        };
105
106        self.expect_kind(&TokenKind::LBrace)?;
107
108        let mut functions = Vec::new();
109        while !self.check_kind(&TokenKind::RBrace) && !self.at_end() {
110            match self.parse_function() {
111                Ok(f) => functions.push(f),
112                Err(e) => {
113                    self.errors.push(e);
114                    self.recover_to_func_or_rbrace();
115                }
116            }
117        }
118
119        self.expect_kind(&TokenKind::RBrace)?;
120
121        Ok(ModuleDecl {
122            name,
123            functions,
124            span: start,
125        })
126    }
127
128    fn parse_function(&mut self) -> Result<FuncDecl, ParseError> {
129        let start = self.expect_kind(&TokenKind::Func)?.span;
130
131        let name = match &self.advance().kind {
132            TokenKind::AtIdent(n) => n.clone(),
133            other => {
134                return Err(ParseError::UnexpectedToken {
135                    found: format!("{}", other),
136                    expected: "@func_name".into(),
137                    line: start.line,
138                    col: start.column,
139                })
140            }
141        };
142
143        self.expect_kind(&TokenKind::LParen)?;
144
145        let mut params = Vec::new();
146        while !self.check_kind(&TokenKind::RParen) && !self.at_end() {
147            if !params.is_empty() {
148                self.expect_kind(&TokenKind::Comma)?;
149            }
150            params.push(self.parse_param()?);
151        }
152
153        self.expect_kind(&TokenKind::RParen)?;
154
155        let mut returns = Vec::new();
156        if self.check_kind(&TokenKind::Arrow) {
157            self.advance(); // ->
158            returns = self.parse_return_types()?;
159        }
160
161        self.expect_kind(&TokenKind::LBrace)?;
162
163        let mut body = Vec::new();
164        while !self.check_kind(&TokenKind::RBrace) && !self.at_end() {
165            match self.parse_statement() {
166                Ok(s) => body.push(s),
167                Err(e) => {
168                    self.errors.push(e);
169                    self.recover_to_statement();
170                }
171            }
172        }
173
174        self.expect_kind(&TokenKind::RBrace)?;
175
176        Ok(FuncDecl {
177            name,
178            params,
179            returns,
180            body,
181            span: start,
182        })
183    }
184
185    fn parse_param(&mut self) -> Result<ParamDecl, ParseError> {
186        let tok = self.advance();
187        let name = match &tok.kind {
188            TokenKind::PercentIdent(n) => n.clone(),
189            other => {
190                return Err(ParseError::UnexpectedToken {
191                    found: format!("{}", other),
192                    expected: "%param_name".into(),
193                    line: tok.span.line,
194                    col: tok.span.column,
195                })
196            }
197        };
198
199        self.expect_kind(&TokenKind::Colon)?;
200        let ty = self.parse_type()?;
201
202        Ok(ParamDecl {
203            name,
204            ty,
205            span: tok.span,
206        })
207    }
208
209    fn parse_return_types(&mut self) -> Result<Vec<TypeExpr>, ParseError> {
210        if self.check_kind(&TokenKind::LParen) {
211            self.advance();
212            let mut types = Vec::new();
213            while !self.check_kind(&TokenKind::RParen) && !self.at_end() {
214                if !types.is_empty() {
215                    self.expect_kind(&TokenKind::Comma)?;
216                }
217                types.push(self.parse_type()?);
218            }
219            self.expect_kind(&TokenKind::RParen)?;
220            Ok(types)
221        } else {
222            let ty = self.parse_type()?;
223            Ok(vec![ty])
224        }
225    }
226
227    fn parse_type(&mut self) -> Result<TypeExpr, ParseError> {
228        match self.peek_kind() {
229            TokenKind::Tensor => {
230                self.advance();
231                self.expect_kind(&TokenKind::LAngle)?;
232                let mut shape = Vec::new();
233                loop {
234                    match self.peek_kind() {
235                        TokenKind::Integer(n) => {
236                            self.advance();
237                            shape.push(DimExpr::Constant(n as usize));
238                        }
239                        TokenKind::Ident(s) => {
240                            // Check if it's a dtype like f32, i8, etc.
241                            if is_dtype(&s) {
242                                let dtype = s;
243                                self.advance();
244                                self.expect_kind(&TokenKind::RAngle)?;
245                                return Ok(TypeExpr::Tensor(TensorTypeExpr { shape, dtype }));
246                            } else {
247                                self.advance();
248                                shape.push(DimExpr::Symbolic(s));
249                            }
250                        }
251                        TokenKind::Star => {
252                            self.advance();
253                            shape.push(DimExpr::Dynamic);
254                        }
255                        _ => {
256                            return Err(ParseError::General(format!(
257                                "Expected dimension or dtype in tensor type at line {}",
258                                self.current_span().line
259                            )));
260                        }
261                    }
262                    // After a dimension, expect 'x' separator or end
263                    if self.check_ident("x") {
264                        self.advance();
265                    } else if self.check_kind(&TokenKind::RAngle) {
266                        self.advance();
267                        return Ok(TypeExpr::Tensor(TensorTypeExpr {
268                            shape,
269                            dtype: "f32".into(),
270                        }));
271                    }
272                }
273            }
274            TokenKind::Qubit => {
275                self.advance();
276                Ok(TypeExpr::Qubit)
277            }
278            TokenKind::Bit => {
279                self.advance();
280                Ok(TypeExpr::Bit)
281            }
282            TokenKind::Void => {
283                self.advance();
284                Ok(TypeExpr::Void)
285            }
286            TokenKind::Index => {
287                self.advance();
288                Ok(TypeExpr::Index)
289            }
290            TokenKind::Hamiltonian => {
291                self.advance();
292                self.expect_kind(&TokenKind::LAngle)?;
293                let n = self.expect_integer()?;
294                self.expect_kind(&TokenKind::RAngle)?;
295                Ok(TypeExpr::Hamiltonian(n as usize))
296            }
297            TokenKind::Ident(s) if is_scalar_type(&s) => {
298                self.advance();
299                Ok(TypeExpr::Scalar(ScalarTypeExpr { name: s }))
300            }
301            _ => {
302                let tok = self.advance();
303                Err(ParseError::UnexpectedToken {
304                    found: format!("{}", tok.kind),
305                    expected: "type".into(),
306                    line: tok.span.line,
307                    col: tok.span.column,
308                })
309            }
310        }
311    }
312
313    fn parse_statement(&mut self) -> Result<Statement, ParseError> {
314        if self.check_kind(&TokenKind::Return) {
315            return self.parse_return().map(Statement::Return);
316        }
317
318        // Try to parse op assignment: %v0 = "op"(...) or %v0, %v1 = "op"(...)
319        if self.check_percent_ident() {
320            return self.parse_op_assign().map(Statement::OpAssign);
321        }
322
323        // Bare operation call: "op"(...)
324        if self.check_string_literal() {
325            return self.parse_bare_op().map(Statement::OpAssign);
326        }
327
328        let tok = self.advance();
329        Err(ParseError::UnexpectedToken {
330            found: format!("{}", tok.kind),
331            expected: "statement".into(),
332            line: tok.span.line,
333            col: tok.span.column,
334        })
335    }
336
337    fn parse_return(&mut self) -> Result<ReturnStmt, ParseError> {
338        let span = self.advance().span; // 'return'
339
340        let mut values = Vec::new();
341        while !self.check_kind(&TokenKind::RBrace)
342            && !self.at_end()
343            && !self.check_kind(&TokenKind::Eof)
344        {
345            if !values.is_empty() {
346                self.expect_kind(&TokenKind::Comma)?;
347            }
348            values.push(self.parse_operand()?);
349        }
350
351        Ok(ReturnStmt { values, span })
352    }
353
354    fn parse_op_assign(&mut self) -> Result<OpAssign, ParseError> {
355        let span = self.current_span();
356
357        // Parse result values: %v0 or %v0, %v1
358        let mut results = Vec::new();
359        loop {
360            let tok = self.advance();
361            match &tok.kind {
362                TokenKind::PercentIdent(n) => results.push(n.clone()),
363                other => {
364                    return Err(ParseError::UnexpectedToken {
365                        found: format!("{}", other),
366                        expected: "%result_name".into(),
367                        line: tok.span.line,
368                        col: tok.span.column,
369                    })
370                }
371            }
372            if self.check_kind(&TokenKind::Comma) {
373                self.advance();
374                // If next is '=' then the comma was part of a different grammar
375                if self.check_kind(&TokenKind::Equal) {
376                    break;
377                }
378            } else {
379                break;
380            }
381        }
382
383        self.expect_kind(&TokenKind::Equal)?;
384
385        // Parse operation name (string literal)
386        let op_name = match &self.advance().kind {
387            TokenKind::StringLiteral(s) => s.clone(),
388            other => {
389                return Err(ParseError::General(format!(
390                    "Expected operation name string, got {}",
391                    other
392                )))
393            }
394        };
395
396        // Parse operands
397        self.expect_kind(&TokenKind::LParen)?;
398        let mut operands = Vec::new();
399        while !self.check_kind(&TokenKind::RParen) && !self.at_end() {
400            if !operands.is_empty() {
401                self.expect_kind(&TokenKind::Comma)?;
402            }
403            operands.push(self.parse_operand()?);
404        }
405        self.expect_kind(&TokenKind::RParen)?;
406
407        // Parse optional attributes
408        let mut attrs = Vec::new();
409        if self.check_kind(&TokenKind::LBrace) {
410            attrs = self.parse_attr_dict()?;
411        }
412
413        // Parse optional type signature
414        let type_sig = if self.check_kind(&TokenKind::Colon) {
415            self.advance();
416            Some(self.parse_type_signature()?)
417        } else {
418            None
419        };
420
421        Ok(OpAssign {
422            results,
423            op_name,
424            operands,
425            attrs,
426            type_sig,
427            span,
428        })
429    }
430
431    fn parse_bare_op(&mut self) -> Result<OpAssign, ParseError> {
432        let span = self.current_span();
433
434        let op_name = match &self.advance().kind {
435            TokenKind::StringLiteral(s) => s.clone(),
436            other => {
437                return Err(ParseError::General(format!(
438                    "Expected operation name string, got {}",
439                    other
440                )))
441            }
442        };
443
444        self.expect_kind(&TokenKind::LParen)?;
445        let mut operands = Vec::new();
446        while !self.check_kind(&TokenKind::RParen) && !self.at_end() {
447            if !operands.is_empty() {
448                self.expect_kind(&TokenKind::Comma)?;
449            }
450            operands.push(self.parse_operand()?);
451        }
452        self.expect_kind(&TokenKind::RParen)?;
453
454        let mut attrs = Vec::new();
455        if self.check_kind(&TokenKind::LBrace) {
456            attrs = self.parse_attr_dict()?;
457        }
458
459        let type_sig = if self.check_kind(&TokenKind::Colon) {
460            self.advance();
461            Some(self.parse_type_signature()?)
462        } else {
463            None
464        };
465
466        Ok(OpAssign {
467            results: Vec::new(),
468            op_name,
469            operands,
470            attrs,
471            type_sig,
472            span,
473        })
474    }
475
476    fn parse_operand(&mut self) -> Result<Operand, ParseError> {
477        match self.peek_kind() {
478            TokenKind::PercentIdent(n) => {
479                self.advance();
480                Ok(Operand::Value(n))
481            }
482            TokenKind::AtIdent(n) => {
483                self.advance();
484                Ok(Operand::FuncRef(n))
485            }
486            TokenKind::Integer(v) => {
487                self.advance();
488                Ok(Operand::Literal(LiteralValue::Integer(v)))
489            }
490            TokenKind::Float(v) => {
491                self.advance();
492                Ok(Operand::Literal(LiteralValue::Float(v)))
493            }
494            TokenKind::True => {
495                self.advance();
496                Ok(Operand::Literal(LiteralValue::Bool(true)))
497            }
498            TokenKind::False => {
499                self.advance();
500                Ok(Operand::Literal(LiteralValue::Bool(false)))
501            }
502            _ => {
503                let tok = self.advance();
504                Err(ParseError::UnexpectedToken {
505                    found: format!("{}", tok.kind),
506                    expected: "operand".into(),
507                    line: tok.span.line,
508                    col: tok.span.column,
509                })
510            }
511        }
512    }
513
514    fn parse_attr_dict(&mut self) -> Result<Vec<(String, AttrValue)>, ParseError> {
515        self.expect_kind(&TokenKind::LBrace)?;
516        let mut attrs = Vec::new();
517
518        while !self.check_kind(&TokenKind::RBrace) && !self.at_end() {
519            if !attrs.is_empty() {
520                self.expect_kind(&TokenKind::Comma)?;
521            }
522            let key = match &self.advance().kind {
523                TokenKind::Ident(s) => s.clone(),
524                other => {
525                    return Err(ParseError::General(format!(
526                        "Expected attribute key, got {}",
527                        other
528                    )))
529                }
530            };
531            self.expect_kind(&TokenKind::Equal)?;
532            let value = self.parse_attr_value()?;
533            attrs.push((key, value));
534        }
535
536        self.expect_kind(&TokenKind::RBrace)?;
537        Ok(attrs)
538    }
539
540    fn parse_attr_value(&mut self) -> Result<AttrValue, ParseError> {
541        match self.peek_kind() {
542            TokenKind::Integer(v) => {
543                self.advance();
544                Ok(AttrValue::Integer(v))
545            }
546            TokenKind::Float(v) => {
547                self.advance();
548                Ok(AttrValue::Float(v))
549            }
550            TokenKind::True => {
551                self.advance();
552                Ok(AttrValue::Bool(true))
553            }
554            TokenKind::False => {
555                self.advance();
556                Ok(AttrValue::Bool(false))
557            }
558            TokenKind::StringLiteral(s) => {
559                self.advance();
560                Ok(AttrValue::String(s))
561            }
562            TokenKind::LBracket => {
563                self.advance();
564                let mut elems = Vec::new();
565                while !self.check_kind(&TokenKind::RBracket) && !self.at_end() {
566                    if !elems.is_empty() {
567                        self.expect_kind(&TokenKind::Comma)?;
568                    }
569                    elems.push(self.parse_attr_value()?);
570                }
571                self.expect_kind(&TokenKind::RBracket)?;
572                Ok(AttrValue::Array(elems))
573            }
574            _ => {
575                let tok = self.advance();
576                Err(ParseError::UnexpectedToken {
577                    found: format!("{}", tok.kind),
578                    expected: "attribute value".into(),
579                    line: tok.span.line,
580                    col: tok.span.column,
581                })
582            }
583        }
584    }
585
586    fn parse_type_signature(&mut self) -> Result<TypeSignature, ParseError> {
587        self.expect_kind(&TokenKind::LParen)?;
588        let mut inputs = Vec::new();
589        while !self.check_kind(&TokenKind::RParen) && !self.at_end() {
590            if !inputs.is_empty() {
591                self.expect_kind(&TokenKind::Comma)?;
592            }
593            inputs.push(self.parse_type()?);
594        }
595        self.expect_kind(&TokenKind::RParen)?;
596
597        self.expect_kind(&TokenKind::Arrow)?;
598
599        let outputs = self.parse_return_types()?;
600
601        Ok(TypeSignature { inputs, outputs })
602    }
603
604    // ── Helpers ──
605
606    fn peek_kind(&self) -> TokenKind {
607        if self.pos < self.tokens.len() {
608            self.tokens[self.pos].kind.clone()
609        } else {
610            TokenKind::Eof
611        }
612    }
613
614    fn check_kind(&self, kind: &TokenKind) -> bool {
615        std::mem::discriminant(&self.peek_kind()) == std::mem::discriminant(kind)
616    }
617
618    fn check_percent_ident(&self) -> bool {
619        matches!(self.peek_kind(), TokenKind::PercentIdent(_))
620    }
621
622    fn check_string_literal(&self) -> bool {
623        matches!(self.peek_kind(), TokenKind::StringLiteral(_))
624    }
625
626    fn check_ident(&self, name: &str) -> bool {
627        matches!(&self.peek_kind(), TokenKind::Ident(s) if s == name)
628    }
629
630    fn advance(&mut self) -> Token {
631        if self.pos < self.tokens.len() {
632            let tok = self.tokens[self.pos].clone();
633            self.pos += 1;
634            tok
635        } else {
636            Token {
637                kind: TokenKind::Eof,
638                span: Span {
639                    start: 0,
640                    end: 0,
641                    line: 0,
642                    column: 0,
643                },
644                text: String::new(),
645            }
646        }
647    }
648
649    fn expect_kind(&mut self, expected: &TokenKind) -> Result<Token, ParseError> {
650        let tok = self.advance();
651        if std::mem::discriminant(&tok.kind) == std::mem::discriminant(expected) {
652            Ok(tok)
653        } else {
654            Err(ParseError::UnexpectedToken {
655                found: format!("{}", tok.kind),
656                expected: format!("{}", expected),
657                line: tok.span.line,
658                col: tok.span.column,
659            })
660        }
661    }
662
663    fn expect_integer(&mut self) -> Result<i64, ParseError> {
664        let tok = self.advance();
665        match tok.kind {
666            TokenKind::Integer(v) => Ok(v),
667            _ => Err(ParseError::UnexpectedToken {
668                found: format!("{}", tok.kind),
669                expected: "integer".into(),
670                line: tok.span.line,
671                col: tok.span.column,
672            }),
673        }
674    }
675
676    fn current_span(&self) -> Span {
677        if self.pos < self.tokens.len() {
678            self.tokens[self.pos].span
679        } else {
680            Span {
681                start: 0,
682                end: 0,
683                line: 0,
684                column: 0,
685            }
686        }
687    }
688
689    fn at_end(&self) -> bool {
690        self.pos >= self.tokens.len() || matches!(self.peek_kind(), TokenKind::Eof)
691    }
692
693    fn recover(&mut self) {
694        while !self.at_end() {
695            match self.peek_kind() {
696                TokenKind::Module | TokenKind::Func | TokenKind::HashDialect(_) => return,
697                TokenKind::RBrace => {
698                    self.advance();
699                    return;
700                }
701                _ => {
702                    self.advance();
703                }
704            }
705        }
706    }
707
708    fn recover_to_func_or_rbrace(&mut self) {
709        while !self.at_end() {
710            match self.peek_kind() {
711                TokenKind::Func | TokenKind::RBrace => return,
712                _ => {
713                    self.advance();
714                }
715            }
716        }
717    }
718
719    fn recover_to_statement(&mut self) {
720        while !self.at_end() {
721            match self.peek_kind() {
722                TokenKind::PercentIdent(_)
723                | TokenKind::Return
724                | TokenKind::StringLiteral(_)
725                | TokenKind::RBrace => return,
726                _ => {
727                    self.advance();
728                }
729            }
730        }
731    }
732
733    pub fn errors(&self) -> &[ParseError] {
734        &self.errors
735    }
736}
737
738fn is_dtype(s: &str) -> bool {
739    matches!(
740        s,
741        "f64"
742            | "f32"
743            | "f16"
744            | "bf16"
745            | "fp8e4m3"
746            | "fp8e5m2"
747            | "i64"
748            | "i32"
749            | "i16"
750            | "i8"
751            | "i4"
752            | "i2"
753            | "u8"
754            | "i1"
755            | "index"
756    )
757}
758
759fn is_scalar_type(s: &str) -> bool {
760    matches!(
761        s,
762        "f64" | "f32" | "f16" | "bf16" | "i64" | "i32" | "i16" | "i8" | "u8" | "bool"
763    )
764}
765
766#[cfg(test)]
767mod tests {
768    use super::*;
769    use crate::lexer::Lexer;
770
771    fn parse_source(src: &str) -> Result<Program, Vec<ParseError>> {
772        let mut lexer = Lexer::new(src);
773        let tokens = lexer.tokenize().to_vec();
774        let mut parser = Parser::new(tokens);
775        parser.parse()
776    }
777
778    #[test]
779    fn test_parse_simple_module() {
780        let src = r#"
781#dialect tensor
782
783module @test {
784    func @relu(%x: tensor<4xf32>) -> tensor<4xf32> {
785        %out = "tensor.relu"(%x) : (tensor<4xf32>) -> tensor<4xf32>
786        return %out
787    }
788}
789"#;
790        let result = parse_source(src);
791        assert!(result.is_ok(), "Parse errors: {:?}", result.err());
792        let program = result.unwrap();
793        assert_eq!(program.dialect_directives.len(), 1);
794        assert_eq!(program.modules.len(), 1);
795        assert_eq!(program.modules[0].functions.len(), 1);
796        assert_eq!(program.modules[0].functions[0].name, "relu");
797    }
798
799    #[test]
800    fn test_parse_quantum_module() {
801        let src = r#"
802#dialect quantum
803
804module @qc {
805    func @bell(%q0: qubit, %q1: qubit) -> (qubit, qubit) {
806        %q2 = "quantum.h"(%q0) : (qubit) -> qubit
807        %q3, %q4 = "quantum.cx"(%q2, %q1) : (qubit, qubit) -> (qubit, qubit)
808        return %q3, %q4
809    }
810}
811"#;
812        let result = parse_source(src);
813        assert!(result.is_ok(), "Parse errors: {:?}", result.err());
814    }
815}
lift_ast/parser.rs

lift_ast/
parser.rs
