mq_lang/ast/
node.rs

1use super::{Program, TokenId};
2#[cfg(feature = "ast-json")]
3use crate::arena::ArenaId;
4use crate::{Ident, Shared, Token, arena::Arena, number::Number, range::Range, selector::Selector};
5#[cfg(feature = "ast-json")]
6use serde::{Deserialize, Serialize};
7use smallvec::SmallVec;
8use smol_str::SmolStr;
9use std::{
10    fmt::{self, Display, Formatter},
11    hash::{Hash, Hasher},
12};
13
14/// Represents a function parameter with an optional default value
15#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
16#[derive(PartialEq, PartialOrd, Debug, Clone)]
17pub struct Param {
18    pub ident: IdentWithToken,
19    pub default: Option<Shared<Node>>,
20}
21
22impl Display for Param {
23    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
24        write!(f, "{}", self.ident)
25    }
26}
27
28impl Param {
29    pub fn new(name: IdentWithToken) -> Self {
30        Self::with_default(name, None)
31    }
32
33    pub fn with_default(name: IdentWithToken, default_value: Option<Shared<Node>>) -> Self {
34        Self {
35            ident: name,
36            default: default_value,
37        }
38    }
39}
40
41pub type Params = SmallVec<[Param; 4]>;
42pub type Args = SmallVec<[Shared<Node>; 4]>;
43pub type Cond = (Option<Shared<Node>>, Shared<Node>);
44pub type Branches = SmallVec<[Cond; 4]>;
45pub type MatchArms = SmallVec<[MatchArm; 4]>;
46
47#[derive(PartialEq, PartialOrd, Debug, Clone)]
48#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
49pub struct Node {
50    #[cfg_attr(
51        feature = "ast-json",
52        serde(skip_serializing, skip_deserializing, default = "default_token_id")
53    )]
54    pub token_id: TokenId,
55    pub expr: Shared<Expr>,
56}
57
58#[cfg(feature = "ast-json")]
59fn default_token_id() -> TokenId {
60    ArenaId::new(0)
61}
62
63impl Node {
64    #[cfg(feature = "ast-json")]
65    pub fn to_json(&self) -> Result<String, serde_json::Error> {
66        serde_json::to_string_pretty(self)
67    }
68
69    #[cfg(feature = "ast-json")]
70    pub fn from_json(json_str: &str) -> Result<Self, serde_json::Error> {
71        serde_json::from_str(json_str)
72    }
73
74    pub fn range(&self, arena: Shared<Arena<Shared<Token>>>) -> Range {
75        match &*self.expr {
76            Expr::Block(program)
77            | Expr::Def(_, _, program)
78            | Expr::Fn(_, program)
79            | Expr::While(_, program)
80            | Expr::Loop(program)
81            | Expr::Module(_, program)
82            | Expr::Foreach(_, _, program) => {
83                let start = program
84                    .first()
85                    .map(|node| node.range(Shared::clone(&arena)).start)
86                    .unwrap_or_default();
87                let end = program
88                    .last()
89                    .map(|node| node.range(Shared::clone(&arena)).end)
90                    .unwrap_or_default();
91                Range { start, end }
92            }
93            Expr::Call(_, args) => {
94                let start = args
95                    .first()
96                    .map(|node| node.range(Shared::clone(&arena)).start)
97                    .unwrap_or_default();
98                let end = args
99                    .last()
100                    .map(|node| node.range(Shared::clone(&arena)).end)
101                    .unwrap_or_default();
102                Range { start, end }
103            }
104            Expr::CallDynamic(callable, args) => {
105                let start = callable.range(Shared::clone(&arena)).start;
106                let end = args
107                    .last()
108                    .map(|node| node.range(Shared::clone(&arena)).end)
109                    .unwrap_or_else(|| callable.range(Shared::clone(&arena)).end);
110                Range { start, end }
111            }
112            Expr::Macro(_, params, block) => {
113                let start = params
114                    .first()
115                    .and_then(|param| param.ident.token.as_ref().map(|t| t.range))
116                    .unwrap_or(block.range(Shared::clone(&arena)))
117                    .start;
118                let end = block.range(arena).end;
119                Range { start, end }
120            }
121            Expr::Let(_, node)
122            | Expr::Var(_, node)
123            | Expr::Assign(_, node)
124            | Expr::Quote(node)
125            | Expr::Unquote(node) => node.range(Shared::clone(&arena)),
126            Expr::If(nodes) => {
127                if let (Some(first), Some(last)) = (nodes.first(), nodes.last()) {
128                    let start = first.1.range(Shared::clone(&arena));
129                    let end = last.1.range(Shared::clone(&arena));
130                    Range {
131                        start: start.start,
132                        end: end.end,
133                    }
134                } else {
135                    // Fallback to token range if no branches exist
136                    arena[self.token_id].range
137                }
138            }
139            Expr::Match(value, arms) => {
140                let start = value.range(Shared::clone(&arena)).start;
141                let end = arms
142                    .last()
143                    .map(|arm| arm.body.range(Shared::clone(&arena)).end)
144                    .unwrap_or_else(|| arena[self.token_id].range.end);
145                Range { start, end }
146            }
147            Expr::Paren(node) => node.range(Shared::clone(&arena)),
148            Expr::Try(try_expr, catch_expr) => {
149                let start = try_expr.range(Shared::clone(&arena)).start;
150                let end = catch_expr.range(Shared::clone(&arena)).end;
151                Range { start, end }
152            }
153            Expr::And(expr1, expr2) | Expr::Or(expr1, expr2) => {
154                let start = expr1.range(Shared::clone(&arena)).start;
155                let end = expr2.range(Shared::clone(&arena)).end;
156                Range { start, end }
157            }
158            Expr::Break(Some(value_node)) => {
159                let start = arena[self.token_id].range.start;
160                let end = value_node.range(Shared::clone(&arena)).end;
161                Range { start, end }
162            }
163            Expr::Literal(_)
164            | Expr::Ident(_)
165            | Expr::Selector(_)
166            | Expr::Include(_)
167            | Expr::Import(_)
168            | Expr::InterpolatedString(_)
169            | Expr::QualifiedAccess(_, _)
170            | Expr::Nodes
171            | Expr::Self_
172            | Expr::Break(None)
173            | Expr::Continue => arena[self.token_id].range,
174        }
175    }
176
177    pub fn is_nodes(&self) -> bool {
178        matches!(*self.expr, Expr::Nodes)
179    }
180}
181
182#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
183#[derive(PartialEq, Debug, Eq, Clone)]
184pub struct IdentWithToken {
185    pub name: Ident,
186    #[cfg_attr(feature = "ast-json", serde(skip_serializing_if = "Option::is_none", default))]
187    pub token: Option<Shared<Token>>,
188}
189
190impl Hash for IdentWithToken {
191    fn hash<H: Hasher>(&self, state: &mut H) {
192        self.name.hash(state);
193    }
194}
195
196impl Ord for IdentWithToken {
197    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
198        self.name.cmp(&other.name)
199    }
200}
201
202impl PartialOrd for IdentWithToken {
203    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
204        Some(self.cmp(other))
205    }
206}
207
208impl IdentWithToken {
209    pub fn new(name: &str) -> Self {
210        Self::new_with_token(name, None)
211    }
212
213    pub fn new_with_token(name: &str, token: Option<Shared<Token>>) -> Self {
214        Self {
215            name: name.into(),
216            token,
217        }
218    }
219}
220
221impl Display for IdentWithToken {
222    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
223        write!(f, "{}", self.name)
224    }
225}
226
227#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
228#[derive(Debug, Clone, PartialOrd, PartialEq)]
229pub enum StringSegment {
230    Text(String),
231    Expr(Shared<Node>),
232    Env(SmolStr),
233    Self_,
234}
235
236#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
237#[derive(PartialEq, PartialOrd, Debug, Clone)]
238pub enum Pattern {
239    Literal(Literal),
240    Ident(IdentWithToken),
241    Wildcard,
242    Array(Vec<Pattern>),
243    ArrayRest(Vec<Pattern>, IdentWithToken), // patterns before .., rest binding
244    Dict(Vec<(IdentWithToken, Pattern)>),
245    Type(Ident), // :string, :number, etc.
246}
247
248#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
249#[derive(PartialEq, PartialOrd, Debug, Clone)]
250pub struct MatchArm {
251    pub pattern: Pattern,
252    pub guard: Option<Shared<Node>>,
253    pub body: Shared<Node>,
254}
255
256#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
257#[derive(PartialEq, PartialOrd, Debug, Clone)]
258pub enum Literal {
259    String(String),
260    Number(Number),
261    Symbol(Ident),
262    Bool(bool),
263    None,
264}
265
266impl From<&str> for Literal {
267    fn from(s: &str) -> Self {
268        Literal::String(s.to_owned())
269    }
270}
271
272#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
273#[derive(PartialEq, PartialOrd, Debug, Clone)]
274pub enum AccessTarget {
275    Call(IdentWithToken, Args),
276    Ident(IdentWithToken),
277}
278
279impl Display for Literal {
280    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
281        match self {
282            Literal::String(s) => write!(f, "{}", s),
283            Literal::Number(n) => write!(f, "{}", n),
284            Literal::Symbol(i) => write!(f, "{}", i),
285            Literal::Bool(b) => write!(f, "{}", b),
286            Literal::None => write!(f, "none"),
287        }
288    }
289}
290
291#[cfg_attr(feature = "ast-json", derive(Serialize, Deserialize))]
292#[derive(PartialEq, PartialOrd, Debug, Clone)]
293pub enum Expr {
294    Block(Program),
295    Call(IdentWithToken, Args),
296    CallDynamic(Shared<Node>, Args),
297    Def(IdentWithToken, Params, Program),
298    Macro(IdentWithToken, Params, Shared<Node>),
299    Fn(Params, Program),
300    Let(IdentWithToken, Shared<Node>),
301    Loop(Program),
302    Var(IdentWithToken, Shared<Node>),
303    Assign(IdentWithToken, Shared<Node>),
304    And(Shared<Node>, Shared<Node>),
305    Or(Shared<Node>, Shared<Node>),
306    Literal(Literal),
307    Ident(IdentWithToken),
308    InterpolatedString(Vec<StringSegment>),
309    Selector(Selector),
310    While(Shared<Node>, Program),
311    Foreach(IdentWithToken, Shared<Node>, Program),
312    If(Branches),
313    Match(Shared<Node>, MatchArms),
314    Include(Literal),
315    Import(Literal),
316    Module(IdentWithToken, Program),
317    QualifiedAccess(Vec<IdentWithToken>, AccessTarget),
318    Self_,
319    Nodes,
320    Paren(Shared<Node>),
321    Quote(Shared<Node>),
322    Unquote(Shared<Node>),
323    Try(Shared<Node>, Shared<Node>),
324    Break(Option<Shared<Node>>),
325    Continue,
326}
327
328#[cfg(feature = "debugger")]
329impl Display for Expr {
330    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
331        match self {
332            Expr::Call(ident, args) => {
333                write!(f, "{}(", ident)?;
334                for (i, arg) in args.iter().enumerate() {
335                    if i > 0 {
336                        write!(f, ", ")?;
337                    }
338                    write!(f, "{}", arg.expr)?;
339                }
340                write!(f, ")")
341            }
342            Expr::CallDynamic(callable, args) => {
343                write!(f, "{}(", callable.expr)?;
344                for (i, arg) in args.iter().enumerate() {
345                    if i > 0 {
346                        write!(f, ", ")?;
347                    }
348                    write!(f, "{}", arg.expr)?;
349                }
350                write!(f, ")")
351            }
352            _ => write!(f, ""),
353        }
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360    use crate::{Position, TokenKind, arena::ArenaId};
361    use rstest::rstest;
362    use smallvec::smallvec;
363
364    fn create_token(range: Range) -> Shared<Token> {
365        Shared::new(Token {
366            range,
367            kind: TokenKind::Eof,
368            module_id: ArenaId::new(0),
369        })
370    }
371
372    #[rstest]
373    #[case(
374        Expr::CallDynamic(
375            Shared::new(Node {
376                token_id: ArenaId::new(1),
377                expr: Shared::new(Expr::Literal(Literal::String("callee".to_string()))),
378            }),
379            smallvec![
380                Shared::new(Node {
381                    token_id: ArenaId::new(0),
382                    expr: Shared::new(Expr::Literal(Literal::String("arg1".to_string()))),
383                }),
384                Shared::new(Node {
385                    token_id: ArenaId::new(1),
386                    expr: Shared::new(Expr::Literal(Literal::String("arg2".to_string()))),
387                }),
388            ]
389        ),
390        vec![
391            (0, Range { start: Position::new(1, 1), end: Position::new(1, 5) }),
392            (1, Range { start: Position::new(2, 1), end: Position::new(2, 5) }),
393        ],
394        Range { start: Position::new(2, 1), end: Position::new(2, 5) }
395    )]
396    #[case(
397        Expr::Match(
398            Shared::new(Node {
399                token_id: ArenaId::new(0),
400                expr: Shared::new(Expr::Literal(Literal::String("val".to_string()))),
401            }),
402            smallvec![
403                MatchArm {
404                    pattern: Pattern::Literal(Literal::String("a".to_string())),
405                    guard: None,
406                    body: Shared::new(Node {
407                        token_id: ArenaId::new(1),
408                        expr: Shared::new(Expr::Literal(Literal::String("body1".to_string()))),
409                    }),
410                },
411                MatchArm {
412                    pattern: Pattern::Literal(Literal::String("b".to_string())),
413                    guard: None,
414                    body: Shared::new(Node {
415                        token_id: ArenaId::new(2),
416                        expr: Shared::new(Expr::Literal(Literal::String("body2".to_string()))),
417                    }),
418                },
419            ]
420        ),
421        vec![
422            (0, Range { start: Position::new(10, 1), end: Position::new(10, 5) }),
423            (1, Range { start: Position::new(11, 1), end: Position::new(11, 5) }),
424            (2, Range { start: Position::new(12, 1), end: Position::new(12, 5) }),
425        ],
426        Range { start: Position::new(10, 1), end: Position::new(12, 5) }
427    )]
428    #[case(
429        Expr::Try(
430            Shared::new(Node {
431                token_id: ArenaId::new(0),
432                expr: Shared::new(Expr::Literal(Literal::String("try".to_string()))),
433            }),
434            Shared::new(Node {
435                token_id: ArenaId::new(1),
436                expr: Shared::new(Expr::Literal(Literal::String("catch".to_string()))),
437            })
438        ),
439        vec![
440            (0, Range { start: Position::new(20, 1), end: Position::new(20, 5) }),
441            (1, Range { start: Position::new(21, 1), end: Position::new(21, 5) }),
442        ],
443        Range { start: Position::new(20, 1), end: Position::new(21, 5) }
444    )]
445    #[case(
446        Expr::Let(
447            IdentWithToken::new("x"),
448            Shared::new(Node {
449                token_id: ArenaId::new(0),
450                expr: Shared::new(Expr::Literal(Literal::String("letval".to_string()))),
451            })
452        ),
453        vec![
454            (0, Range { start: Position::new(30, 1), end: Position::new(30, 5) }),
455        ],
456        Range { start: Position::new(30, 1), end: Position::new(30, 5) }
457    )]
458    #[case(
459        Expr::Paren(
460            Shared::new(Node {
461                token_id: ArenaId::new(0),
462                expr: Shared::new(Expr::Literal(Literal::String("paren".to_string()))),
463            })
464        ),
465        vec![
466            (0, Range { start: Position::new(40, 1), end: Position::new(40, 5) }),
467        ],
468        Range { start: Position::new(40, 1), end: Position::new(40, 5) }
469    )]
470    #[case(
471        Expr::Block(vec![
472            Shared::new(Node {
473                token_id: ArenaId::new(0),
474                expr: Shared::new(Expr::Literal(Literal::String("block1".to_string()))),
475            }),
476            Shared::new(Node {
477                token_id: ArenaId::new(1),
478                expr: Shared::new(Expr::Literal(Literal::String("block2".to_string()))),
479            }),
480        ]),
481        vec![
482            (0, Range { start: Position::new(50, 1), end: Position::new(50, 5) }),
483            (1, Range { start: Position::new(51, 1), end: Position::new(51, 5) }),
484        ],
485        Range { start: Position::new(50, 1), end: Position::new(51, 5) }
486    )]
487    #[case(
488        Expr::Def(
489            IdentWithToken::new("f"),
490            smallvec![],
491            vec![
492                Shared::new(Node {
493                    token_id: ArenaId::new(0),
494                    expr: Shared::new(Expr::Literal(Literal::String("def1".to_string()))),
495                }),
496                Shared::new(Node {
497                    token_id: ArenaId::new(1),
498                    expr: Shared::new(Expr::Literal(Literal::String("def2".to_string()))),
499                }),
500            ]
501        ),
502        vec![
503            (0, Range { start: Position::new(60, 1), end: Position::new(60, 5) }),
504            (1, Range { start: Position::new(61, 1), end: Position::new(61, 5) }),
505        ],
506        Range { start: Position::new(60, 1), end: Position::new(61, 5) }
507    )]
508    #[case(
509        Expr::Fn(
510            smallvec![],
511            vec![
512                Shared::new(Node {
513                    token_id: ArenaId::new(0),
514                    expr: Shared::new(Expr::Literal(Literal::String("fn1".to_string()))),
515                }),
516                Shared::new(Node {
517                    token_id: ArenaId::new(1),
518                    expr: Shared::new(Expr::Literal(Literal::String("fn2".to_string()))),
519                }),
520            ]
521        ),
522        vec![
523            (0, Range { start: Position::new(70, 1), end: Position::new(70, 5) }),
524            (1, Range { start: Position::new(71, 1), end: Position::new(71, 5) }),
525        ],
526        Range { start: Position::new(70, 1), end: Position::new(71, 5) }
527    )]
528    #[case(
529        Expr::While(
530            Shared::new(Node {
531                token_id: ArenaId::new(0),
532                expr: Shared::new(Expr::Literal(Literal::String("cond".to_string()))),
533            }),
534            vec![
535                Shared::new(Node {
536                    token_id: ArenaId::new(1),
537                    expr: Shared::new(Expr::Literal(Literal::String("while1".to_string()))),
538                }),
539                Shared::new(Node {
540                    token_id: ArenaId::new(2),
541                    expr: Shared::new(Expr::Literal(Literal::String("while2".to_string()))),
542                }),
543            ]
544        ),
545        vec![
546            (0, Range { start: Position::new(81, 1), end: Position::new(81, 5) }),
547            (1, Range { start: Position::new(82, 1), end: Position::new(82, 5) }),
548            (2, Range { start: Position::new(82, 1), end: Position::new(82, 5) }),
549        ],
550        Range { start: Position::new(82, 1), end: Position::new(82, 5) }
551    )]
552    #[case(
553        Expr::Foreach(
554            IdentWithToken::new("item"),
555            Shared::new(Node {
556                token_id: ArenaId::new(0),
557                expr: Shared::new(Expr::Literal(Literal::String("iter".to_string()))),
558            }),
559            vec![
560                Shared::new(Node {
561                    token_id: ArenaId::new(1),
562                    expr: Shared::new(Expr::Literal(Literal::String("foreach1".to_string()))),
563                }),
564                Shared::new(Node {
565                    token_id: ArenaId::new(2),
566                    expr: Shared::new(Expr::Literal(Literal::String("foreach2".to_string()))),
567                }),
568            ]
569        ),
570        vec![
571            (0, Range { start: Position::new(101, 1), end: Position::new(101, 5) }),
572            (1, Range { start: Position::new(102, 1), end: Position::new(102, 5) }),
573            (2, Range { start: Position::new(102, 1), end: Position::new(102, 5) }),
574        ],
575        Range { start: Position::new(102, 1), end: Position::new(102, 5) }
576    )]
577    #[case(
578        Expr::If(smallvec![
579            (
580                Some(Shared::new(Node {
581                    token_id: ArenaId::new(0),
582                    expr: Shared::new(Expr::Literal(Literal::String("cond1".to_string()))),
583                })),
584                Shared::new(Node {
585                    token_id: ArenaId::new(1),
586                    expr: Shared::new(Expr::Literal(Literal::String("if1".to_string()))),
587                })
588            ),
589            (
590                Some(Shared::new(Node {
591                    token_id: ArenaId::new(2),
592                    expr: Shared::new(Expr::Literal(Literal::String("cond2".to_string()))),
593                })),
594                Shared::new(Node {
595                    token_id: ArenaId::new(3),
596                    expr: Shared::new(Expr::Literal(Literal::String("if2".to_string()))),
597                })
598            ),
599        ]),
600        vec![
601            (0, Range { start: Position::new(111, 1), end: Position::new(111, 5) }),
602            (1, Range { start: Position::new(113, 1), end: Position::new(113, 5) }),
603            (2, Range { start: Position::new(114, 1), end: Position::new(115, 5) }),
604            (3, Range { start: Position::new(116, 1), end: Position::new(117, 5) }),
605        ],
606        Range { start: Position::new(113, 1), end: Position::new(117, 5) }
607    )]
608    #[case(
609        Expr::Call(
610            IdentWithToken::new("func"),
611            smallvec![
612                Shared::new(Node {
613                    token_id: ArenaId::new(0),
614                    expr: Shared::new(Expr::Literal(Literal::String("arg1".to_string()))),
615                }),
616                Shared::new(Node {
617                    token_id: ArenaId::new(1),
618                    expr: Shared::new(Expr::Literal(Literal::String("arg2".to_string()))),
619                }),
620            ]
621        ),
622        vec![
623            (0, Range { start: Position::new(120, 1), end: Position::new(120, 5) }),
624            (1, Range { start: Position::new(121, 1), end: Position::new(121, 5) }),
625        ],
626        Range { start: Position::new(120, 1), end: Position::new(121, 5) }
627    )]
628    fn test_node_range_various_exprs(
629        #[case] expr: Expr,
630        #[case] token_ranges: Vec<(usize, Range)>,
631        #[case] expected: Range,
632    ) {
633        let mut arena = Arena::new(150);
634        for (_, range) in &token_ranges {
635            let token = create_token(*range);
636            let _ = arena.alloc(token);
637        }
638        let node = Node {
639            token_id: ArenaId::new(0),
640            expr: Shared::new(expr),
641        };
642        assert_eq!(node.range(Shared::new(arena)), expected);
643    }
644}
mq_lang/ast/node.rs

mq_lang/ast/
node.rs
