pulsar_frontend/
ast.rs

1// Copyright (C) 2024 Ethan Uppal. All rights reserved.
2use super::{
3    token::{Token, TokenType},
4    ty::{StmtTypeCell, Type, TypeCell}
5};
6use crate::{
7    attribute::{Attribute, Attributes},
8    ty::StmtType
9};
10use pulsar_utils::{
11    format,
12    loc::{Loc, RegionProvider}
13};
14use std::fmt::{self, Display};
15
16pub type Param = (Token, Type);
17
18pub(crate) trait TokenRegionProvider {
19    fn start_token(&self) -> &Token;
20    fn end_token(&self) -> &Token;
21}
22
23macro_rules! implement_region_provider_for_token_provider {
24    ($T:ident) => {
25        impl RegionProvider for $T {
26            fn start(&self) -> Loc {
27                self.start_token().loc.clone()
28            }
29
30            fn end(&self) -> Loc {
31                let end_token = self.end_token();
32                let mut loc = end_token.loc.clone();
33                // tokens are always on one line
34                if end_token.ty != TokenType::Newline {
35                    let length = end_token.value.len() as isize;
36                    loc.pos += length;
37                    loc.col += length;
38                }
39                loc
40            }
41        }
42    };
43}
44
45#[derive(Clone)]
46pub enum ExprValue {
47    ConstantInt(i64),
48    /// TODO: Support `::`s
49    BoundName(Token),
50
51    MemberAccess(Box<Expr>, Token),
52
53    /// TODO: Call an `expr` or some sort of chaining of `::`
54    Call(Token, Vec<Expr>),
55
56    /// `ArrayLiteral(elements, should_continue)` is an array literal beginning
57    /// with `elements` and filling the remainder of the array with zeros if
58    /// `should_continue`.
59    ArrayLiteral(Vec<Expr>, bool),
60
61    PrefixOp(Token, Box<Expr>),
62    InfixBop(Box<Expr>, Token, Box<Expr>),
63    PostfixBop(Box<Expr>, Token, Box<Expr>, Token),
64
65    /// `HardwareMap(map_token, parallel_factor, f, arr)` is an array produced
66    /// by applying `f` elementwise to `arr` using a hardware parallelism
67    /// factor of `parallel_factor`.
68    HardwareMap(Token, usize, Token, Box<Expr>)
69}
70
71#[derive(Clone)]
72pub struct Expr {
73    pub value: ExprValue,
74    pub ty: TypeCell,
75    start_token: Token,
76    end_token: Token
77}
78
79impl Expr {
80    /// Constructs a new expression with the given `value` that ranges from
81    /// `start_token` to `end_token`.
82    pub fn new(value: ExprValue, start_token: Token, end_token: Token) -> Self {
83        Self {
84            value,
85            start_token,
86            end_token,
87            ty: TypeCell::new(Type::Unknown)
88        }
89    }
90}
91
92impl Display for Expr {
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        match &self.value {
95            ExprValue::ConstantInt(i) => {
96                write!(f, "{}", i)?;
97            }
98            ExprValue::BoundName(name) => {
99                write!(f, "{}", name.value)?;
100            }
101            ExprValue::MemberAccess(value, member) => {
102                write!(f, "{}.{}", value, member.value)?;
103            }
104            ExprValue::Call(name, args) => {
105                write!(
106                    f,
107                    "{}({})",
108                    name.value,
109                    args.iter()
110                        .map(|arg| arg.to_string())
111                        .collect::<Vec<_>>()
112                        .join(", ")
113                )?;
114            }
115            ExprValue::ArrayLiteral(elements, should_continue) => {
116                write!(
117                    f,
118                    "[{}{}]",
119                    elements
120                        .iter()
121                        .map(|ty| ty.to_string())
122                        .collect::<Vec<_>>()
123                        .join(", "),
124                    if *should_continue {
125                        format!(
126                            "{}...",
127                            if elements.is_empty() { "" } else { ", " }
128                        )
129                    } else {
130                        "".into()
131                    }
132                )?;
133            }
134            ExprValue::PrefixOp(op, rhs) => {
135                write!(f, "({} {})", op.value, rhs)?;
136            }
137            ExprValue::InfixBop(lhs, op, rhs) => {
138                write!(f, "({} {} {})", lhs, op.value, rhs)?;
139            }
140            ExprValue::PostfixBop(lhs, op1, rhs, op2) => {
141                write!(f, "({}{}{}{})", lhs, op1.value, rhs, op2.value)?;
142            }
143            ExprValue::HardwareMap(_, parallel_factor, fun, arr) => {
144                write!(f, "map<{}>({}, {})", parallel_factor, fun.value, arr)?;
145            }
146        }
147
148        let expr_ty = self.ty.as_ref();
149        if expr_ty.clone().is_known() {
150            write!(f, ": {}", expr_ty)?;
151        }
152
153        Ok(())
154    }
155}
156
157impl TokenRegionProvider for Expr {
158    fn start_token(&self) -> &Token {
159        &self.start_token
160    }
161
162    fn end_token(&self) -> &Token {
163        &self.end_token
164    }
165}
166
167implement_region_provider_for_token_provider!(Expr);
168
169#[derive(Clone)]
170pub enum NodeValue {
171    Function {
172        name: Token,
173        params: Vec<Param>,
174        ret: Type,
175        pure_token: Option<Token>,
176        body: Vec<Node>
177    },
178    LetBinding {
179        name: Token,
180        hint: Option<TypeCell>,
181        value: Box<Expr>
182    },
183    Return {
184        keyword_token: Token,
185        value: Option<Box<Expr>>
186    }
187}
188
189#[derive(Clone)]
190pub struct Node {
191    pub value: NodeValue,
192    pub ty: StmtTypeCell,
193    pub attributes: Attributes,
194    start_token: Token,
195    end_token: Token
196}
197
198impl Node {
199    /// Constructs a node node with the given `value` that ranges from
200    /// `start_token` to `end_token`.
201    pub fn new(value: NodeValue, start_token: Token, end_token: Token) -> Self {
202        Self {
203            value,
204            ty: StmtType::make_unknown(),
205            attributes: Attributes::default(),
206            start_token,
207            end_token
208        }
209    }
210
211    /// Marks this node as generated (that is, not present syntactically in the
212    /// user's source).
213    pub fn mark_generated(mut self) -> Self {
214        self.attributes.add(Attribute::Generated);
215        self
216    }
217
218    /// Pretty-prints this node at the given indentation `level`.
219    fn pretty(&self, level: usize) -> String {
220        let mut result = format::make_indent(level);
221        let content = match &self.value {
222            NodeValue::Function {
223                name,
224                params,
225                ret,
226                pure_token,
227                body
228            } => {
229                let insert_newline = if body.is_empty() { "" } else { "\n" };
230                format!(
231                    "{}func {}({}) -> {} {{{}{}{}{}}}",
232                    if pure_token.is_some() { "pure " } else { "" },
233                    name.value,
234                    params
235                        .iter()
236                        .map(|(name, ty)| format!("{}: {}", name.value, ty))
237                        .collect::<Vec<_>>()
238                        .join(", "),
239                    ret,
240                    insert_newline,
241                    body.iter()
242                        .map(|node| { node.pretty(level + 1) })
243                        .collect::<Vec<_>>()
244                        .join("\n"),
245                    insert_newline,
246                    format::make_indent(level)
247                )
248            }
249            NodeValue::LetBinding {
250                name,
251                hint: hint_opt,
252                value
253            } => {
254                let hint_str = if let Some(hint) = hint_opt {
255                    format!(": {}", hint)
256                } else {
257                    "".into()
258                };
259                format!("let {}{} = {}", name.value, hint_str, value)
260            }
261            NodeValue::Return {
262                keyword_token: _,
263                value: value_opt
264            } => {
265                format!(
266                    "return{}",
267                    if let Some(value) = value_opt {
268                        format!(" {}", value)
269                    } else {
270                        "".into()
271                    }
272                )
273            }
274        };
275        result.push_str(&content);
276        result
277    }
278}
279
280impl Display for Node {
281    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
282        self.pretty(0).fmt(f)
283    }
284}
285
286impl TokenRegionProvider for Node {
287    fn start_token(&self) -> &Token {
288        &self.start_token
289    }
290
291    fn end_token(&self) -> &Token {
292        &self.end_token
293    }
294}
295
296implement_region_provider_for_token_provider!(Node);