shrew_ir/
parser.rs

1// Parser — Recursive descent parser for the .sw language
2//
3// The parser consumes a Vec<Token> (from the Lexer) and produces an AST
4// (Program). It's a classic hand-written recursive descent parser — one
5// method per grammar production.
6//
7// ERROR RECOVERY: For now we use "panic mode" — on error, we return Err
8// immediately. A future improvement could skip to the next @directive
9// and collect multiple errors per parse.
10//
11// OPERATOR PRECEDENCE (lowest to highest):
12//   1. ?? (null coalescing)
13//   2. || (logical or)
14//   3. && (logical and)
15//   4. | (bitwise or)
16//   5. ^ (bitwise xor)
17//   6. & (bitwise and)
18//   7. == != (equality)
19//   8. < > <= >= (comparison)
20//   9. << >> (shift)
21//  10. + - (additive)
22//  11. * / % (multiplicative)
23//  12. ** (power, right-associative)
24//  13. - ! ~ (unary prefix)
25//  14. . [] () (postfix: member, index, call)
26
27use crate::ast::*;
28use crate::error::{Error, ErrorKind, Result};
29use crate::token::{Span, Token, TokenKind};
30
31pub struct Parser {
32    tokens: Vec<Token>,
33    pos: usize,
34}
35
36impl Parser {
37    pub fn new(tokens: Vec<Token>) -> Self {
38        Self { tokens, pos: 0 }
39    }
40
41    /// Parse the full program.
42    pub fn parse_program(&mut self) -> Result<Program> {
43        let mut items = Vec::new();
44        while !self.at_eof() {
45            items.push(self.parse_top_level()?);
46        }
47        Ok(Program { items })
48    }
49
50    // Top-level dispatch
51
52    fn parse_top_level(&mut self) -> Result<TopLevel> {
53        match self.peek_kind() {
54            TokenKind::AtImport => Ok(TopLevel::Import(self.parse_import()?)),
55            TokenKind::AtModel => Ok(TopLevel::Metadata(self.parse_metadata()?)),
56            TokenKind::AtConfig => Ok(TopLevel::Config(self.parse_config()?)),
57            TokenKind::AtTypes => Ok(TopLevel::Types(self.parse_types()?)),
58            TokenKind::AtGraph => Ok(TopLevel::Graph(self.parse_graph()?)),
59            TokenKind::AtCustomOp => Ok(TopLevel::CustomOp(self.parse_custom_op()?)),
60            TokenKind::AtTraining => Ok(TopLevel::Training(self.parse_training()?)),
61            TokenKind::AtInference => Ok(TopLevel::Inference(self.parse_inference()?)),
62            TokenKind::AtMetrics => Ok(TopLevel::Metrics(self.parse_metrics()?)),
63            TokenKind::AtLogging => Ok(TopLevel::Logging(self.parse_logging()?)),
64            TokenKind::AtVisualizations => Ok(TopLevel::Visualization(self.parse_visualization()?)),
65            _ => Err(self.error_unexpected("a directive (@model, @graph, etc.)")),
66        }
67    }
68
69    // @import
70
71    fn parse_import(&mut self) -> Result<ImportStmt> {
72        let span = self.expect(TokenKind::AtImport)?.span;
73        let path = self.expect_string()?;
74        let alias = if self.check(&TokenKind::As) {
75            self.advance();
76            Some(self.expect_ident()?)
77        } else {
78            None
79        };
80        self.expect(TokenKind::Semi)?;
81        Ok(ImportStmt { path, alias, span })
82    }
83  
84    // @model
85
86    fn parse_metadata(&mut self) -> Result<MetadataBlock> {
87        let span = self.expect(TokenKind::AtModel)?.span;
88        self.expect(TokenKind::LBrace)?;
89        let mut fields = Vec::new();
90        while !self.check(&TokenKind::RBrace) {
91            fields.push(self.parse_field()?);
92        }
93        self.expect(TokenKind::RBrace)?;
94        Ok(MetadataBlock { fields, span })
95    }
96
97    fn parse_field(&mut self) -> Result<Field> {
98        let span = self.current_span();
99        let key = self.expect_key()?;
100        self.expect(TokenKind::Colon)?;
101        let value = self.parse_literal()?;
102        self.expect(TokenKind::Semi)?;
103        Ok(Field { key, value, span })
104    }
105
106    // @config
107
108    fn parse_config(&mut self) -> Result<ConfigBlock> {
109        let span = self.expect(TokenKind::AtConfig)?.span;
110        self.expect(TokenKind::LBrace)?;
111        let mut fields = Vec::new();
112        while !self.check(&TokenKind::RBrace) {
113            fields.push(self.parse_expr_field()?);
114        }
115        self.expect(TokenKind::RBrace)?;
116        Ok(ConfigBlock { fields, span })
117    }
118
119    fn parse_expr_field(&mut self) -> Result<ExprField> {
120        let span = self.current_span();
121        let key = self.expect_key()?;
122        self.expect(TokenKind::Colon)?;
123        let value = self.parse_expr()?;
124        self.expect(TokenKind::Semi)?;
125        Ok(ExprField { key, value, span })
126    }
127  
128    // @types 
129
130    fn parse_types(&mut self) -> Result<TypesBlock> {
131        let span = self.expect(TokenKind::AtTypes)?.span;
132        self.expect(TokenKind::LBrace)?;
133        let mut defs = Vec::new();
134        while !self.check(&TokenKind::RBrace) {
135            defs.push(self.parse_type_def()?);
136        }
137        self.expect(TokenKind::RBrace)?;
138        Ok(TypesBlock { defs, span })
139    }
140
141    fn parse_type_def(&mut self) -> Result<TypeDef> {
142        let span = self.current_span();
143        self.expect(TokenKind::Type)?;
144        let name = self.expect_key()?;
145        self.expect(TokenKind::Eq)?;
146        let ty = self.parse_type_expr()?;
147        self.expect(TokenKind::Semi)?;
148        Ok(TypeDef { name, ty, span })
149    }
150 
151    // Type expressions 
152
153    fn parse_type_expr(&mut self) -> Result<TypeExpr> {
154        let span = self.current_span();
155
156        // Tensor<[dims], dtype>
157        if self.check(&TokenKind::Tensor) {
158            return self.parse_tensor_type();
159        }
160
161        // Tuple: (Type, Type, ...)
162        if self.check(&TokenKind::LParen) {
163            self.advance();
164            let mut types = Vec::new();
165            if !self.check(&TokenKind::RParen) {
166                types.push(self.parse_type_expr()?);
167                while self.check(&TokenKind::Comma) {
168                    self.advance();
169                    types.push(self.parse_type_expr()?);
170                }
171            }
172            self.expect(TokenKind::RParen)?;
173            return Ok(TypeExpr::Tuple(types, span));
174        }
175
176        // List: [Type]
177        if self.check(&TokenKind::LBracket) {
178            self.advance();
179            let inner = self.parse_type_expr()?;
180            self.expect(TokenKind::RBracket)?;
181            return Ok(TypeExpr::List(Box::new(inner), span));
182        }
183
184        // Dynamic: ?
185        if self.check(&TokenKind::Question) {
186            self.advance();
187            return Ok(TypeExpr::Dynamic(span));
188        }
189
190        // Scalar dtype
191        if let Some(dt) = self.try_parse_dtype() {
192            return Ok(TypeExpr::Scalar(dt, span));
193        }
194
195        // Concrete int dimension
196        if let TokenKind::IntLit(n) = self.peek_kind() {
197            self.advance();
198            return Ok(TypeExpr::IntDim(n, span));
199        }
200
201        // Named type alias
202        if self.is_ident_or_keyword_as_ident() {
203            let name = self.expect_key()?;
204            return Ok(TypeExpr::Named(name, span));
205        }
206
207        Err(self.error_unexpected("a type expression"))
208    }
209
210    fn parse_tensor_type(&mut self) -> Result<TypeExpr> {
211        let span = self.current_span();
212        self.expect(TokenKind::Tensor)?;
213        self.expect(TokenKind::Lt)?;
214        self.expect(TokenKind::LBracket)?;
215
216        let mut dims = Vec::new();
217        if !self.check(&TokenKind::RBracket) {
218            dims.push(self.parse_dimension()?);
219            while self.check(&TokenKind::Comma) {
220                self.advance();
221                if self.check(&TokenKind::RBracket) {
222                    break;
223                }
224                dims.push(self.parse_dimension()?);
225            }
226        }
227        self.expect(TokenKind::RBracket)?;
228        self.expect(TokenKind::Comma)?;
229        let dtype = self.parse_dtype()?;
230        self.expect(TokenKind::Gt)?;
231        Ok(TypeExpr::Tensor { dims, dtype, span })
232    }
233
234    fn parse_dimension(&mut self) -> Result<Dimension> {
235        let span = self.current_span();
236
237        if self.check(&TokenKind::Question) {
238            self.advance();
239            return Ok(Dimension::Dynamic(span));
240        }
241        if self.check(&TokenKind::Underscore) {
242            self.advance();
243            return Ok(Dimension::Inferred(span));
244        }
245        if let TokenKind::IntLit(n) = self.peek_kind() {
246            self.advance();
247            return Ok(Dimension::Concrete(n, span));
248        }
249        if self.is_ident_or_keyword_as_ident() {
250            let name = self.expect_key()?;
251            return Ok(Dimension::Named(name, span));
252        }
253        Err(self.error_unexpected("a dimension (name, integer, ?, or _)"))
254    }
255
256    fn parse_dtype(&mut self) -> Result<DTypeKind> {
257        self.try_parse_dtype()
258            .ok_or_else(|| self.error_unexpected("a dtype (f32, f64, i64, etc.)"))
259    }
260
261    fn try_parse_dtype(&mut self) -> Option<DTypeKind> {
262        let dt = match self.peek_kind() {
263            TokenKind::F16 => DTypeKind::F16,
264            TokenKind::F32 => DTypeKind::F32,
265            TokenKind::F64 => DTypeKind::F64,
266            TokenKind::Bf16 => DTypeKind::Bf16,
267            TokenKind::I8 => DTypeKind::I8,
268            TokenKind::I16 => DTypeKind::I16,
269            TokenKind::I32 => DTypeKind::I32,
270            TokenKind::I64 => DTypeKind::I64,
271            TokenKind::U8 => DTypeKind::U8,
272            TokenKind::U16 => DTypeKind::U16,
273            TokenKind::U32 => DTypeKind::U32,
274            TokenKind::U64 => DTypeKind::U64,
275            TokenKind::Bool => DTypeKind::Bool,
276            TokenKind::Complex64 => DTypeKind::Complex64,
277            TokenKind::Complex128 => DTypeKind::Complex128,
278            _ => return None,
279        };
280        self.advance();
281        Some(dt)
282    }
283
284    // @graph
285
286    fn parse_graph(&mut self) -> Result<GraphBlock> {
287        let span = self.expect(TokenKind::AtGraph)?.span;
288        let name = self.expect_key()?;
289
290        // Optional param list
291        let params = if self.check(&TokenKind::LParen) {
292            self.advance();
293            let ps = self.parse_param_list()?;
294            self.expect(TokenKind::RParen)?;
295            ps
296        } else {
297            Vec::new()
298        };
299
300        // Optional return type
301        let return_type = if self.check(&TokenKind::Arrow) {
302            self.advance();
303            Some(self.parse_type_expr()?)
304        } else {
305            None
306        };
307
308        self.expect(TokenKind::LBrace)?;
309        let mut body = Vec::new();
310        while !self.check(&TokenKind::RBrace) {
311            body.push(self.parse_graph_stmt()?);
312        }
313        self.expect(TokenKind::RBrace)?;
314
315        Ok(GraphBlock {
316            name,
317            params,
318            return_type,
319            body,
320            span,
321        })
322    }
323
324    fn parse_param_list(&mut self) -> Result<Vec<ParamDef>> {
325        let mut params = Vec::new();
326        if !self.check(&TokenKind::RParen) {
327            params.push(self.parse_param_def()?);
328            while self.check(&TokenKind::Comma) {
329                self.advance();
330                if self.check(&TokenKind::RParen) {
331                    break;
332                }
333                params.push(self.parse_param_def()?);
334            }
335        }
336        Ok(params)
337    }
338
339    fn parse_param_def(&mut self) -> Result<ParamDef> {
340        let span = self.current_span();
341        let name = self.expect_key()?;
342        self.expect(TokenKind::Colon)?;
343        let ty = self.parse_type_expr()?;
344        let optional = if self.check(&TokenKind::Question) {
345            self.advance();
346            true
347        } else {
348            false
349        };
350        Ok(ParamDef {
351            name,
352            ty,
353            optional,
354            span,
355        })
356    }
357
358    fn parse_graph_stmt(&mut self) -> Result<GraphStmt> {
359        match self.peek_kind() {
360            TokenKind::Input => {
361                let span = self.advance().span;
362                let name = self.expect_key()?;
363                self.expect(TokenKind::Colon)?;
364                let ty = self.parse_type_expr()?;
365                let optional = if self.check(&TokenKind::Question) {
366                    self.advance();
367                    true
368                } else {
369                    false
370                };
371                self.expect(TokenKind::Semi)?;
372                Ok(GraphStmt::Input(InputDecl {
373                    name,
374                    ty,
375                    optional,
376                    span,
377                }))
378            }
379            TokenKind::Output => {
380                let span = self.advance().span;
381                // output [name :] expr ;
382                let (name, expr) = if self.is_ident_or_keyword_as_ident()
383                    && self.peek_ahead_kind(1) == TokenKind::Colon
384                {
385                    let n = self.expect_key()?;
386                    self.expect(TokenKind::Colon)?;
387                    (Some(n), self.parse_expr()?)
388                } else {
389                    (None, self.parse_expr()?)
390                };
391                self.expect(TokenKind::Semi)?;
392                Ok(GraphStmt::Output(OutputDecl { name, expr, span }))
393            }
394            TokenKind::Param => {
395                let span = self.advance().span;
396                let name = self.expect_key()?;
397                self.expect(TokenKind::Colon)?;
398                let ty = self.parse_type_expr()?;
399                let attrs = if self.check(&TokenKind::LBrace) {
400                    self.advance();
401                    let mut attrs = Vec::new();
402                    while !self.check(&TokenKind::RBrace) {
403                        let aspan = self.current_span();
404                        let key = self.expect_key()?;
405                        self.expect(TokenKind::Colon)?;
406                        let value = self.parse_expr()?;
407                        self.expect(TokenKind::Semi)?;
408                        attrs.push(ParamAttr {
409                            key,
410                            value,
411                            span: aspan,
412                        });
413                    }
414                    self.expect(TokenKind::RBrace)?;
415                    attrs
416                } else {
417                    Vec::new()
418                };
419                self.expect(TokenKind::Semi)?;
420                Ok(GraphStmt::Param(ParamDecl {
421                    name,
422                    ty,
423                    attrs,
424                    span,
425                }))
426            }
427            TokenKind::Node => {
428                let span = self.advance().span;
429                let name = self.expect_key()?;
430                let ty = if self.check(&TokenKind::Colon) {
431                    self.advance();
432                    // Could be a type_expr OR a node_body starting with {
433                    if self.check(&TokenKind::LBrace) {
434                        None // the { is the node body, not a dict type
435                    } else {
436                        Some(self.parse_type_expr()?)
437                    }
438                } else {
439                    None
440                };
441                let stmts = if self.check(&TokenKind::LBrace) {
442                    self.advance();
443                    let mut stmts = Vec::new();
444                    while !self.check(&TokenKind::RBrace) {
445                        stmts.push(self.parse_node_stmt()?);
446                    }
447                    self.expect(TokenKind::RBrace)?;
448                    stmts
449                } else {
450                    Vec::new()
451                };
452                self.expect(TokenKind::Semi)?;
453                Ok(GraphStmt::Node(NodeDecl {
454                    name,
455                    ty,
456                    stmts,
457                    span,
458                }))
459            }
460            TokenKind::AtAssert => {
461                let span = self.advance().span;
462                let condition = self.parse_expr()?;
463                let message = if self.check(&TokenKind::Comma) {
464                    self.advance();
465                    Some(self.expect_string()?)
466                } else {
467                    None
468                };
469                self.expect(TokenKind::Semi)?;
470                Ok(GraphStmt::Assert(AssertStmt {
471                    condition,
472                    message,
473                    span,
474                }))
475            }
476            TokenKind::AtCheck => {
477                let span = self.advance().span;
478                let name = self.expect_key()?;
479                self.expect(TokenKind::LBrace)?;
480                let mut conditions = Vec::new();
481                while !self.check(&TokenKind::RBrace) {
482                    // Each line is: assert expr [, "msg"] ;
483                    let aspan = self.current_span();
484                    // "assert" is not a keyword in our grammar but appears
485                    // inside @check. We'll accept it as an ident or keyword.
486                    let word = self.expect_key()?;
487                    if word != "assert" {
488                        return Err(Error::new(
489                            ErrorKind::UnexpectedToken {
490                                expected: "assert".to_string(),
491                                got: word,
492                            },
493                            aspan,
494                        ));
495                    }
496                    let condition = self.parse_expr()?;
497                    let message = if self.check(&TokenKind::Comma) {
498                        self.advance();
499                        Some(self.expect_string()?)
500                    } else {
501                        None
502                    };
503                    self.expect(TokenKind::Semi)?;
504                    conditions.push(AssertStmt {
505                        condition,
506                        message,
507                        span: aspan,
508                    });
509                }
510                self.expect(TokenKind::RBrace)?;
511                Ok(GraphStmt::Check(CheckBlock {
512                    name,
513                    conditions,
514                    span,
515                }))
516            }
517            _ => Err(self.error_unexpected(
518                "a graph statement (input, output, param, node, @assert, @check)",
519            )),
520        }
521    }
522
523    fn parse_node_stmt(&mut self) -> Result<NodeStmt> {
524        let span = self.current_span();
525
526        // @hint
527        if self.check(&TokenKind::AtHint) {
528            self.advance();
529            let kind = match self.peek_kind() {
530                TokenKind::RecomputeInBackward => {
531                    self.advance();
532                    HintKind::RecomputeInBackward
533                }
534                TokenKind::MustPreserve => {
535                    self.advance();
536                    HintKind::MustPreserve
537                }
538                TokenKind::InPlace => {
539                    self.advance();
540                    HintKind::InPlace
541                }
542                TokenKind::NoGrad => {
543                    self.advance();
544                    HintKind::NoGrad
545                }
546                _ if self.is_ident() => {
547                    let name = self.expect_ident()?;
548                    HintKind::Custom(name)
549                }
550                _ => return Err(self.error_unexpected("a hint type")),
551            };
552            self.expect(TokenKind::Semi)?;
553            return Ok(NodeStmt::Hint(kind, span));
554        }
555
556        // keyword : expr ;
557        let key = self.expect_key()?;
558        self.expect(TokenKind::Colon)?;
559
560        let stmt = match key.as_str() {
561            "op" => {
562                let expr = self.parse_expr()?;
563                self.expect(TokenKind::Semi)?;
564                NodeStmt::Op(expr, span)
565            }
566            "input" => {
567                let expr = self.parse_expr()?;
568                self.expect(TokenKind::Semi)?;
569                NodeStmt::InputRef(expr, span)
570            }
571            "output" => {
572                let ty = self.parse_type_expr()?;
573                self.expect(TokenKind::Semi)?;
574                NodeStmt::OutputType(ty, span)
575            }
576            _ => {
577                let expr = self.parse_expr()?;
578                self.expect(TokenKind::Semi)?;
579                NodeStmt::Attr(key, expr, span)
580            }
581        };
582        Ok(stmt)
583    }
584
585    // @custom_op
586
587    fn parse_custom_op(&mut self) -> Result<CustomOpBlock> {
588        let span = self.expect(TokenKind::AtCustomOp)?.span;
589        let name = self.expect_key()?;
590        self.expect(TokenKind::LBrace)?;
591        let mut stmts = Vec::new();
592        while !self.check(&TokenKind::RBrace) {
593            stmts.push(self.parse_custom_op_stmt()?);
594        }
595        self.expect(TokenKind::RBrace)?;
596        Ok(CustomOpBlock { name, stmts, span })
597    }
598
599    fn parse_custom_op_stmt(&mut self) -> Result<CustomOpStmt> {
600        let span = self.current_span();
601        match self.peek_kind() {
602            TokenKind::Signature => {
603                self.advance();
604                self.expect(TokenKind::Colon)?;
605                self.expect(TokenKind::LParen)?;
606                let params = self.parse_param_list()?;
607                self.expect(TokenKind::RParen)?;
608                self.expect(TokenKind::Arrow)?;
609                let return_type = self.parse_type_expr()?;
610                self.expect(TokenKind::Semi)?;
611                Ok(CustomOpStmt::Signature {
612                    params,
613                    return_type,
614                    span,
615                })
616            }
617            TokenKind::Impl => {
618                self.advance();
619                let target = self.expect_key()?;
620                self.expect(TokenKind::LBrace)?;
621                let mut attrs = Vec::new();
622                while !self.check(&TokenKind::RBrace) {
623                    attrs.push(self.parse_expr_field()?);
624                }
625                self.expect(TokenKind::RBrace)?;
626                Ok(CustomOpStmt::Impl {
627                    target,
628                    attrs,
629                    span,
630                })
631            }
632            TokenKind::Gradient => {
633                self.advance();
634                let target = self.expect_key()?;
635                self.expect(TokenKind::LBrace)?;
636                let mut body = Vec::new();
637                while !self.check(&TokenKind::RBrace) {
638                    body.push(self.parse_custom_op_stmt()?);
639                }
640                self.expect(TokenKind::RBrace)?;
641                Ok(CustomOpStmt::Gradient { target, body, span })
642            }
643            _ => Err(self.error_unexpected("signature, impl, or gradient")),
644        }
645    }
646
647    // @training
648
649    fn parse_training(&mut self) -> Result<TrainingBlock> {
650        let span = self.expect(TokenKind::AtTraining)?.span;
651        self.expect(TokenKind::LBrace)?;
652        let mut fields = Vec::new();
653        while !self.check(&TokenKind::RBrace) {
654            fields.push(self.parse_training_field()?);
655        }
656        self.expect(TokenKind::RBrace)?;
657        Ok(TrainingBlock { fields, span })
658    }
659
660    fn parse_training_field(&mut self) -> Result<TrainingField> {
661        let span = self.current_span();
662        match self.peek_kind() {
663            TokenKind::Model => {
664                self.advance();
665                self.expect(TokenKind::Colon)?;
666                let name = self.expect_key()?;
667                self.expect(TokenKind::Semi)?;
668                Ok(TrainingField::Model(name, span))
669            }
670            TokenKind::Loss => {
671                self.advance();
672                self.expect(TokenKind::Colon)?;
673                let name = self.expect_key()?;
674                self.expect(TokenKind::Semi)?;
675                Ok(TrainingField::Loss(name, span))
676            }
677            TokenKind::Optimizer => {
678                self.advance();
679                self.expect(TokenKind::Colon)?;
680                let fields = self.parse_brace_fields()?;
681                Ok(TrainingField::Optimizer(fields, span))
682            }
683            TokenKind::LrSchedule => {
684                self.advance();
685                self.expect(TokenKind::Colon)?;
686                let fields = self.parse_brace_fields()?;
687                Ok(TrainingField::LrSchedule(fields, span))
688            }
689            TokenKind::GradClip => {
690                self.advance();
691                self.expect(TokenKind::Colon)?;
692                let fields = self.parse_brace_fields()?;
693                Ok(TrainingField::GradClip(fields, span))
694            }
695            _ => {
696                let f = self.parse_expr_field()?;
697                Ok(TrainingField::Generic(f))
698            }
699        }
700    }
701
702    // @inference
703
704    fn parse_inference(&mut self) -> Result<InferenceBlock> {
705        let span = self.expect(TokenKind::AtInference)?.span;
706        self.expect(TokenKind::LBrace)?;
707        let mut fields = Vec::new();
708        while !self.check(&TokenKind::RBrace) {
709            fields.push(self.parse_inference_field()?);
710        }
711        self.expect(TokenKind::RBrace)?;
712        Ok(InferenceBlock { fields, span })
713    }
714
715    fn parse_inference_field(&mut self) -> Result<InferenceField> {
716        let span = self.current_span();
717        match self.peek_kind() {
718            TokenKind::Model => {
719                self.advance();
720                self.expect(TokenKind::Colon)?;
721                let name = self.expect_key()?;
722                self.expect(TokenKind::Semi)?;
723                Ok(InferenceField::Model(name, span))
724            }
725            TokenKind::Optimizations => {
726                self.advance();
727                self.expect(TokenKind::Colon)?;
728                let list = self.parse_list_expr_items()?;
729                self.expect(TokenKind::Semi)?;
730                Ok(InferenceField::Optimizations(list, span))
731            }
732            TokenKind::Quantization => {
733                self.advance();
734                self.expect(TokenKind::Colon)?;
735                let fields = self.parse_brace_fields()?;
736                Ok(InferenceField::Quantization(fields, span))
737            }
738            TokenKind::Generation => {
739                self.advance();
740                self.expect(TokenKind::Colon)?;
741                let fields = self.parse_brace_fields()?;
742                Ok(InferenceField::Generation(fields, span))
743            }
744            _ => {
745                let f = self.parse_expr_field()?;
746                Ok(InferenceField::Generic(f))
747            }
748        }
749    }
750
751    // @metrics
752
753    fn parse_metrics(&mut self) -> Result<MetricsBlock> {
754        let span = self.expect(TokenKind::AtMetrics)?.span;
755        let name = self.expect_key()?;
756        self.expect(TokenKind::LBrace)?;
757        let mut defs = Vec::new();
758        while !self.check(&TokenKind::RBrace) {
759            defs.push(self.parse_metric_def()?);
760        }
761        self.expect(TokenKind::RBrace)?;
762        Ok(MetricsBlock { name, defs, span })
763    }
764
765    fn parse_metric_def(&mut self) -> Result<MetricDef> {
766        let span = self.current_span();
767        self.expect(TokenKind::Track)?;
768        let name = self.expect_key()?;
769        self.expect(TokenKind::LBrace)?;
770        let mut attrs = Vec::new();
771        while !self.check(&TokenKind::RBrace) {
772            attrs.push(self.parse_expr_field()?);
773        }
774        self.expect(TokenKind::RBrace)?;
775        Ok(MetricDef { name, attrs, span })
776    }
777 
778    // @logging
779
780    fn parse_logging(&mut self) -> Result<LoggingBlock> {
781        let span = self.expect(TokenKind::AtLogging)?.span;
782        self.expect(TokenKind::LBrace)?;
783        let mut fields = Vec::new();
784        while !self.check(&TokenKind::RBrace) {
785            fields.push(self.parse_expr_field()?);
786        }
787        self.expect(TokenKind::RBrace)?;
788        Ok(LoggingBlock { fields, span })
789    }
790 
791    // @visualizations
792
793    fn parse_visualization(&mut self) -> Result<VisualizationBlock> {
794        let span = self.expect(TokenKind::AtVisualizations)?.span;
795        self.expect(TokenKind::LBrace)?;
796        let mut plots = Vec::new();
797        while !self.check(&TokenKind::RBrace) {
798            plots.push(self.parse_plot_def()?);
799        }
800        self.expect(TokenKind::RBrace)?;
801        Ok(VisualizationBlock { plots, span })
802    }
803
804    fn parse_plot_def(&mut self) -> Result<PlotDef> {
805        let span = self.current_span();
806        self.expect(TokenKind::Plot)?;
807        let name = self.expect_key()?;
808        self.expect(TokenKind::LBrace)?;
809        let mut attrs = Vec::new();
810        while !self.check(&TokenKind::RBrace) {
811            attrs.push(self.parse_expr_field()?);
812        }
813        self.expect(TokenKind::RBrace)?;
814        Ok(PlotDef { name, attrs, span })
815    }
816 
817    // Expressions (Pratt parser with precedence climbing)
818
819    fn parse_expr(&mut self) -> Result<Expr> {
820        self.parse_expr_bp(0)
821    }
822
823    /// Pratt parser: parse expression with minimum binding power `min_bp`.
824    fn parse_expr_bp(&mut self, min_bp: u8) -> Result<Expr> {
825        let mut lhs = self.parse_prefix()?;
826
827        loop {
828            // Postfix: . [] ()
829            lhs = match self.peek_kind() {
830                TokenKind::Dot => {
831                    let span = lhs.span();
832                    self.advance();
833                    let field = self.expect_key()?;
834                    // Check for method call: expr.method(args)
835                    if self.check(&TokenKind::LParen) {
836                        self.advance();
837                        let args = self.parse_arg_list()?;
838                        self.expect(TokenKind::RParen)?;
839                        Expr::QualifiedCall {
840                            path: vec![format!("{}", self.expr_to_string(&lhs)), field],
841                            args,
842                            span,
843                        }
844                    } else {
845                        Expr::Member {
846                            object: Box::new(lhs),
847                            field,
848                            span,
849                        }
850                    }
851                }
852                TokenKind::LBracket => {
853                    let span = lhs.span();
854                    self.advance();
855                    let index = self.parse_expr()?;
856                    let end = if self.check(&TokenKind::Colon) {
857                        self.advance();
858                        Some(Box::new(self.parse_expr()?))
859                    } else {
860                        None
861                    };
862                    self.expect(TokenKind::RBracket)?;
863                    Expr::Index {
864                        object: Box::new(lhs),
865                        index: Box::new(index),
866                        end,
867                        span,
868                    }
869                }
870                _ => break,
871            };
872        }
873
874        loop {
875            let (op, bp) = match self.peek_kind() {
876                TokenKind::QuestionQuestion => (BinOp::NullCoalesce, (1, 2)),
877                TokenKind::PipePipe => (BinOp::Or, (3, 4)),
878                TokenKind::AmpAmp => (BinOp::And, (5, 6)),
879                TokenKind::Pipe => (BinOp::BitOr, (7, 8)),
880                TokenKind::Caret => (BinOp::BitXor, (9, 10)),
881                TokenKind::Amp => (BinOp::BitAnd, (11, 12)),
882                TokenKind::EqEq => (BinOp::Eq, (13, 14)),
883                TokenKind::BangEq => (BinOp::Ne, (13, 14)),
884                TokenKind::Lt => (BinOp::Lt, (15, 16)),
885                TokenKind::Gt => (BinOp::Gt, (15, 16)),
886                TokenKind::LtEq => (BinOp::Le, (15, 16)),
887                TokenKind::GtEq => (BinOp::Ge, (15, 16)),
888                TokenKind::LtLt => (BinOp::Shl, (17, 18)),
889                TokenKind::GtGt => (BinOp::Shr, (17, 18)),
890                TokenKind::Plus => (BinOp::Add, (19, 20)),
891                TokenKind::Minus => (BinOp::Sub, (19, 20)),
892                TokenKind::Star => (BinOp::Mul, (21, 22)),
893                TokenKind::Slash => (BinOp::Div, (21, 22)),
894                TokenKind::Percent => (BinOp::Mod, (21, 22)),
895                TokenKind::StarStar => (BinOp::Pow, (24, 23)), // right-assoc
896                _ => break,
897            };
898
899            let (l_bp, r_bp) = bp;
900            if l_bp < min_bp {
901                break;
902            }
903
904            self.advance();
905            let span = lhs.span();
906            let rhs = self.parse_expr_bp(r_bp)?;
907            lhs = Expr::Binary {
908                left: Box::new(lhs),
909                op,
910                right: Box::new(rhs),
911                span,
912            };
913        }
914
915        Ok(lhs)
916    }
917
918    /// Parse a prefix expression (unary or primary).
919    fn parse_prefix(&mut self) -> Result<Expr> {
920        let span = self.current_span();
921        match self.peek_kind() {
922            TokenKind::Minus => {
923                self.advance();
924                let operand = self.parse_expr_bp(25)?; // unary binds tighter than binary
925                Ok(Expr::Unary {
926                    op: UnaryOp::Neg,
927                    operand: Box::new(operand),
928                    span,
929                })
930            }
931            TokenKind::Bang => {
932                self.advance();
933                let operand = self.parse_expr_bp(25)?;
934                Ok(Expr::Unary {
935                    op: UnaryOp::Not,
936                    operand: Box::new(operand),
937                    span,
938                })
939            }
940            TokenKind::Tilde => {
941                self.advance();
942                let operand = self.parse_expr_bp(25)?;
943                Ok(Expr::Unary {
944                    op: UnaryOp::BitNot,
945                    operand: Box::new(operand),
946                    span,
947                })
948            }
949            _ => self.parse_primary(),
950        }
951    }
952
953    /// Parse a primary expression (literals, identifiers, calls, etc.).
954    fn parse_primary(&mut self) -> Result<Expr> {
955        let span = self.current_span();
956        match self.peek_kind() {
957            TokenKind::IntLit(n) => {
958                self.advance();
959                Ok(Expr::Int(n, span))
960            }
961            TokenKind::FloatLit(n) => {
962                self.advance();
963                Ok(Expr::Float(n, span))
964            }
965            TokenKind::StringLit(ref s) => {
966                let s = s.clone();
967                self.advance();
968                Ok(Expr::Str(s, span))
969            }
970            TokenKind::True => {
971                self.advance();
972                Ok(Expr::Bool(true, span))
973            }
974            TokenKind::False => {
975                self.advance();
976                Ok(Expr::Bool(false, span))
977            }
978            TokenKind::Null => {
979                self.advance();
980                Ok(Expr::Null(span))
981            }
982            TokenKind::LParen => {
983                self.advance();
984                let expr = self.parse_expr()?;
985                self.expect(TokenKind::RParen)?;
986                Ok(Expr::Paren(Box::new(expr), span))
987            }
988            TokenKind::LBracket => {
989                let items = self.parse_list_expr_items()?;
990                Ok(Expr::List(items, span))
991            }
992            TokenKind::LBrace => {
993                self.advance();
994                let mut entries = Vec::new();
995                while !self.check(&TokenKind::RBrace) {
996                    let key = self.expect_key()?;
997                    self.expect(TokenKind::Colon)?;
998                    let value = self.parse_expr()?;
999                    entries.push((key, value));
1000                    if !self.check(&TokenKind::RBrace) {
1001                        self.expect(TokenKind::Comma)?;
1002                    }
1003                }
1004                self.expect(TokenKind::RBrace)?;
1005                Ok(Expr::Dict(entries, span))
1006            }
1007            TokenKind::If => {
1008                self.advance();
1009                let cond = self.parse_expr()?;
1010                self.expect(TokenKind::LBrace)?;
1011                let then_branch = self.parse_expr()?;
1012                self.expect(TokenKind::RBrace)?;
1013                let else_branch = if self.check(&TokenKind::Else) {
1014                    self.advance();
1015                    self.expect(TokenKind::LBrace)?;
1016                    let e = self.parse_expr()?;
1017                    self.expect(TokenKind::RBrace)?;
1018                    Some(Box::new(e))
1019                } else {
1020                    None
1021                };
1022                Ok(Expr::IfExpr {
1023                    cond: Box::new(cond),
1024                    then_branch: Box::new(then_branch),
1025                    else_branch,
1026                    span,
1027                })
1028            }
1029            TokenKind::Repeat => {
1030                self.advance();
1031                self.expect(TokenKind::LParen)?;
1032                let count = self.parse_expr()?;
1033                self.expect(TokenKind::RParen)?;
1034                self.expect(TokenKind::LBrace)?;
1035                let body = self.parse_expr()?;
1036                self.expect(TokenKind::RBrace)?;
1037                Ok(Expr::RepeatExpr {
1038                    count: Box::new(count),
1039                    body: Box::new(body),
1040                    span,
1041                })
1042            }
1043            TokenKind::Pipe => {
1044                // Closure: |a, b| { expr }
1045                self.advance();
1046                let mut params = Vec::new();
1047                if !self.check(&TokenKind::Pipe) {
1048                    params.push(self.expect_ident()?);
1049                    while self.check(&TokenKind::Comma) {
1050                        self.advance();
1051                        params.push(self.expect_ident()?);
1052                    }
1053                }
1054                self.expect(TokenKind::Pipe)?;
1055                self.expect(TokenKind::LBrace)?;
1056                let body = self.parse_expr()?;
1057                self.expect(TokenKind::RBrace)?;
1058                Ok(Expr::Closure {
1059                    params,
1060                    body: Box::new(body),
1061                    span,
1062                })
1063            }
1064            _ if self.is_ident_or_keyword_as_ident() => {
1065                let name = self.consume_as_ident()?;
1066                // Check for function call
1067                if self.check(&TokenKind::LParen) {
1068                    self.advance();
1069                    let args = self.parse_arg_list()?;
1070                    self.expect(TokenKind::RParen)?;
1071                    Ok(Expr::Call {
1072                        func: name,
1073                        args,
1074                        span,
1075                    })
1076                } else if self.check(&TokenKind::ColonColon) {
1077                    // Qualified path: mod::func(args)
1078                    let mut path = vec![name];
1079                    while self.check(&TokenKind::ColonColon) {
1080                        self.advance();
1081                        path.push(self.expect_key()?);
1082                    }
1083                    if self.check(&TokenKind::LParen) {
1084                        self.advance();
1085                        let args = self.parse_arg_list()?;
1086                        self.expect(TokenKind::RParen)?;
1087                        Ok(Expr::QualifiedCall { path, args, span })
1088                    } else {
1089                        // Just a qualified ident; return last as ident
1090                        let full = path.join("::");
1091                        Ok(Expr::Ident(full, span))
1092                    }
1093                } else {
1094                    Ok(Expr::Ident(name, span))
1095                }
1096            }
1097            _ => Err(self.error_unexpected("an expression")),
1098        }
1099    }
1100
1101    fn parse_arg_list(&mut self) -> Result<Vec<Arg>> {
1102        let mut args = Vec::new();
1103        if self.check(&TokenKind::RParen) {
1104            return Ok(args);
1105        }
1106        args.push(self.parse_arg()?);
1107        while self.check(&TokenKind::Comma) {
1108            self.advance();
1109            if self.check(&TokenKind::RParen) {
1110                break;
1111            }
1112            args.push(self.parse_arg()?);
1113        }
1114        Ok(args)
1115    }
1116
1117    fn parse_arg(&mut self) -> Result<Arg> {
1118        let span = self.current_span();
1119        // Try named arg: key : expr (key can be ident or keyword)
1120        let is_key = matches!(self.peek_kind(), TokenKind::Ident(_))
1121            || self.peek_kind().keyword_str().is_some();
1122        if is_key && self.peek_ahead_kind(1) == TokenKind::Colon {
1123            let name = self.expect_key()?;
1124            self.expect(TokenKind::Colon)?;
1125            let value = self.parse_expr()?;
1126            return Ok(Arg {
1127                name: Some(name),
1128                value,
1129                span,
1130            });
1131        }
1132        let value = self.parse_expr()?;
1133        Ok(Arg {
1134            name: None,
1135            value,
1136            span,
1137        })
1138    }
1139
1140    fn parse_list_expr_items(&mut self) -> Result<Vec<Expr>> {
1141        self.expect(TokenKind::LBracket)?;
1142        let mut items = Vec::new();
1143        if !self.check(&TokenKind::RBracket) {
1144            items.push(self.parse_expr()?);
1145            while self.check(&TokenKind::Comma) {
1146                self.advance();
1147                if self.check(&TokenKind::RBracket) {
1148                    break;
1149                }
1150                items.push(self.parse_expr()?);
1151            }
1152        }
1153        self.expect(TokenKind::RBracket)?;
1154        Ok(items)
1155    }
1156
1157    
1158    // Literals
1159    
1160
1161    fn parse_literal(&mut self) -> Result<Literal> {
1162        let span = self.current_span();
1163        match self.peek_kind() {
1164            TokenKind::IntLit(n) => {
1165                self.advance();
1166                Ok(Literal::Int(n, span))
1167            }
1168            TokenKind::FloatLit(n) => {
1169                self.advance();
1170                Ok(Literal::Float(n, span))
1171            }
1172            TokenKind::StringLit(ref s) => {
1173                let s = s.clone();
1174                self.advance();
1175                Ok(Literal::Str(s, span))
1176            }
1177            TokenKind::True => {
1178                self.advance();
1179                Ok(Literal::Bool(true, span))
1180            }
1181            TokenKind::False => {
1182                self.advance();
1183                Ok(Literal::Bool(false, span))
1184            }
1185            TokenKind::Null => {
1186                self.advance();
1187                Ok(Literal::Null(span))
1188            }
1189            TokenKind::LBracket => {
1190                self.advance();
1191                let mut items = Vec::new();
1192                if !self.check(&TokenKind::RBracket) {
1193                    items.push(self.parse_literal()?);
1194                    while self.check(&TokenKind::Comma) {
1195                        self.advance();
1196                        if self.check(&TokenKind::RBracket) {
1197                            break;
1198                        }
1199                        items.push(self.parse_literal()?);
1200                    }
1201                }
1202                self.expect(TokenKind::RBracket)?;
1203                Ok(Literal::List(items, span))
1204            }
1205            TokenKind::LBrace => {
1206                self.advance();
1207                let mut entries = Vec::new();
1208                while !self.check(&TokenKind::RBrace) {
1209                    let key = self.expect_key()?;
1210                    self.expect(TokenKind::Colon)?;
1211                    let val = self.parse_literal()?;
1212                    entries.push((key, val));
1213                    if !self.check(&TokenKind::RBrace) {
1214                        self.expect(TokenKind::Comma)?;
1215                    }
1216                }
1217                self.expect(TokenKind::RBrace)?;
1218                Ok(Literal::Dict(entries, span))
1219            }
1220            _ => Err(self.error_unexpected("a literal value")),
1221        }
1222    }
1223
1224    // Helpers: { key: expr; ... } blocks
1225
1226    fn parse_brace_fields(&mut self) -> Result<Vec<ExprField>> {
1227        self.expect(TokenKind::LBrace)?;
1228        let mut fields = Vec::new();
1229        while !self.check(&TokenKind::RBrace) {
1230            fields.push(self.parse_expr_field()?);
1231        }
1232        self.expect(TokenKind::RBrace)?;
1233        Ok(fields)
1234    }
1235
1236    // Token stream helpers
1237
1238    fn peek(&self) -> &Token {
1239        &self.tokens[self.pos.min(self.tokens.len() - 1)]
1240    }
1241
1242    fn peek_kind(&self) -> TokenKind {
1243        self.peek().kind.clone()
1244    }
1245
1246    fn peek_ahead_kind(&self, offset: usize) -> TokenKind {
1247        let idx = (self.pos + offset).min(self.tokens.len() - 1);
1248        self.tokens[idx].kind.clone()
1249    }
1250
1251    fn current_span(&self) -> Span {
1252        self.peek().span
1253    }
1254
1255    fn at_eof(&self) -> bool {
1256        matches!(self.peek_kind(), TokenKind::Eof)
1257    }
1258
1259    fn check(&self, kind: &TokenKind) -> bool {
1260        std::mem::discriminant(&self.peek_kind()) == std::mem::discriminant(kind)
1261    }
1262
1263    fn advance(&mut self) -> Token {
1264        let tok = self.tokens[self.pos.min(self.tokens.len() - 1)].clone();
1265        if self.pos < self.tokens.len() - 1 {
1266            self.pos += 1;
1267        }
1268        tok
1269    }
1270
1271    fn expect(&mut self, kind: TokenKind) -> Result<Token> {
1272        if self.check(&kind) {
1273            Ok(self.advance())
1274        } else {
1275            Err(Error::new(
1276                ErrorKind::UnexpectedToken {
1277                    expected: format!("{}", kind),
1278                    got: format!("{}", self.peek_kind()),
1279                },
1280                self.current_span(),
1281            ))
1282        }
1283    }
1284
1285    fn expect_ident(&mut self) -> Result<String> {
1286        match self.peek_kind() {
1287            TokenKind::Ident(s) => {
1288                self.advance();
1289                Ok(s)
1290            }
1291            _ => Err(self.error_unexpected("an identifier")),
1292        }
1293    }
1294
1295    /// Consume an identifier OR keyword token as a string key.
1296    /// Used in field-key positions where keywords like `init`, `source`,
1297    /// `frozen`, `op`, `backend`, etc. can appear as dictionary keys.
1298    fn expect_key(&mut self) -> Result<String> {
1299        match self.peek_kind() {
1300            TokenKind::Ident(s) => {
1301                self.advance();
1302                Ok(s)
1303            }
1304            ref kind if kind.keyword_str().is_some() => {
1305                let s = kind.keyword_str().unwrap().to_string();
1306                self.advance();
1307                Ok(s)
1308            }
1309            _ => Err(self.error_unexpected("a key (identifier or keyword)")),
1310        }
1311    }
1312
1313    fn expect_string(&mut self) -> Result<String> {
1314        match self.peek_kind() {
1315            TokenKind::StringLit(s) => {
1316                self.advance();
1317                Ok(s)
1318            }
1319            _ => Err(self.error_unexpected("a string literal")),
1320        }
1321    }
1322
1323    fn is_ident(&self) -> bool {
1324        matches!(self.peek_kind(), TokenKind::Ident(_))
1325    }
1326
1327    /// Check if current token is an identifier or a keyword that can double
1328    /// as an identifier in expression context.
1329    fn is_ident_or_keyword_as_ident(&self) -> bool {
1330        matches!(self.peek_kind(), TokenKind::Ident(_)) || self.peek_kind().keyword_str().is_some()
1331    }
1332
1333    /// Consume the current token as an identifier string, even if it's a keyword
1334    /// that can be used as an ident in expression context.
1335    fn consume_as_ident(&mut self) -> Result<String> {
1336        self.expect_key()
1337    }
1338
1339    fn error_unexpected(&self, expected: &str) -> Error {
1340        Error::new(
1341            ErrorKind::UnexpectedToken {
1342                expected: expected.to_string(),
1343                got: format!("{}", self.peek_kind()),
1344            },
1345            self.current_span(),
1346        )
1347    }
1348
1349    fn expr_to_string(&self, expr: &Expr) -> String {
1350        match expr {
1351            Expr::Ident(name, _) => name.clone(),
1352            _ => "<expr>".to_string(),
1353        }
1354    }
1355}
1356
1357// Public convenience function
1358
1359/// Parse a .sw source string into an AST Program.
1360pub fn parse(source: &str) -> Result<Program> {
1361    let tokens = crate::lexer::Lexer::new(source).tokenize()?;
1362    let mut parser = Parser::new(tokens);
1363    parser.parse_program()
1364}
shrew_ir/parser.rs

shrew_ir/
parser.rs
