1use bumpalo::Bump;
2use chumsky::{
3 IterParser, Parser,
4 error::Rich,
5 extra::{self, SimpleState},
6 input::ValueInput,
7 pratt::{infix, left, right},
8 primitive::just,
9 recursive::recursive,
10 select,
11 span::SimpleSpan,
12};
13use lasso::Rodeo;
14
15use crate::{
16 ast::{BinOp, ExprKind, Expr, Ident, Literal, Pat, UnaryOp},
17 token::Token,
18};
19
20type MapExtra<'a, 'b, I, E> = chumsky::input::MapExtra<'a, 'b, I, E>;
21type Extra<'tok, 'src> = extra::Full<Rich<'tok, Token<'src>>, SimpleState<Rodeo>, ()>;
22
23type BoxedParser<'tok, 'src, 'bump, I, O> = chumsky::Boxed<'tok, 'tok, I, O, Extra<'tok, 'src>>;
24
25fn literal_parser<'tok, 'src, 'bump, I>(
26 bump: &'bump Bump,
27) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
28where
29 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
30 'src: 'tok,
31 'bump: 'tok,
32{
33 select! {
34 Token::Int(num) => Literal::Int(num),
35 Token::Float(num) => Literal::Float(num),
36 Token::True => Literal::Bool(true),
37 Token::False => Literal::Bool(false),
38 }
39 .map_with(|lit, e: &mut MapExtra<'_, '_, I, Extra<'tok, 'src>>| {
40 ExprKind::literal(bump, e.span(), lit)
41 })
42 .boxed()
43}
44
45fn ident_parser<'tok, 'src, 'bump, I>() -> BoxedParser<'tok, 'src, 'bump, I, Ident>
46where
47 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
48 'src: 'tok,
49 'bump: 'tok,
50{
51 select! { Token::Ident(name) => name }
52 .map_with(|name, e: &mut MapExtra<'_, '_, I, Extra<'tok, 'src>>| {
53 Ident(e.state().get_or_intern(name))
54 })
55 .boxed()
56}
57
58fn ident_expr_parser<'tok, 'src, 'bump, I>(
59 bump: &'bump Bump,
60) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
61where
62 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
63 'src: 'tok,
64 'bump: 'tok,
65{
66 ident_parser()
67 .map_with(|name, e: &mut MapExtra<'_, '_, I, Extra<'tok, 'src>>| {
68 ExprKind::ident(bump, e.span(), name)
69 })
70 .boxed()
71}
72
73fn pat_parser<'tok, 'src, 'bump, I>() -> BoxedParser<'tok, 'src, 'bump, I, Pat<'bump>>
74where
75 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
76 'src: 'tok,
77 'bump: 'tok,
78{
79 let wildcard = just(Token::Ident("_"))
80 .map_with(|_, _: &mut MapExtra<'_, '_, I, Extra<'tok, 'src>>| Pat::Wildcard);
81
82 let pat_lit = select! {
83 Token::Int(n) => Pat::Lit(Literal::Int(n)),
84 Token::Float(f) => Pat::Lit(Literal::Float(f)),
85 Token::True => Pat::Lit(Literal::Bool(true)),
86 Token::False => Pat::Lit(Literal::Bool(false)),
87 };
88
89 let pat_var = ident_parser().map(|name| Pat::Var(name));
90
91 wildcard.or(pat_lit).or(pat_var).boxed()
92}
93
94fn atom_parser<'tok, 'src, 'bump, I>(
95 bump: &'bump Bump,
96 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
97) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
98where
99 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
100 'src: 'tok,
101 'bump: 'tok,
102{
103 let group = expr.delimited_by(just(Token::LParen), just(Token::RParen));
104
105 literal_parser(bump)
106 .or(ident_expr_parser(bump))
107 .or(group)
108 .boxed()
109}
110
111fn call_parser<'tok, 'src, 'bump, I>(
112 bump: &'bump Bump,
113 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
114) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
115where
116 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
117 'src: 'tok,
118 'bump: 'tok,
119{
120 let atom = atom_parser(bump, expr);
121 atom.clone()
122 .foldl_with(atom.repeated(), |func, arg, e| {
123 ExprKind::app(bump, e.span(), func, arg)
124 })
125 .boxed()
126}
127
128fn unary_parser<'tok, 'src, 'bump, I>(
129 bump: &'bump Bump,
130 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
131) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
132where
133 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
134 'src: 'tok,
135 'bump: 'tok,
136{
137 let call = call_parser(bump, expr);
138
139 let neg = just(Token::Minus).repeated().foldr_with(call, |_, rhs, e| {
140 ExprKind::unaryop(bump, e.span(), UnaryOp::Neg, rhs)
141 });
142
143 just(Token::Bang)
144 .repeated()
145 .foldr_with(neg, |_, rhs, e| {
146 ExprKind::unaryop(bump, e.span(), UnaryOp::Not, rhs)
147 })
148 .boxed()
149}
150
151fn pratt_parser<'tok, 'src, 'bump, I>(
152 bump: &'bump Bump,
153 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
154) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
155where
156 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
157 'src: 'tok,
158 'bump: 'tok,
159{
160 let unary = unary_parser(bump, expr);
161
162 unary
163 .pratt((
164 infix(right(4), just(Token::Caret), |lhs, _, rhs, e| {
165 ExprKind::binop(bump, e.span(), BinOp::Pow, lhs, rhs)
166 }),
167 infix(left(3), just(Token::Star), |lhs, _, rhs, e| {
168 ExprKind::binop(bump, e.span(), BinOp::Mul, lhs, rhs)
169 }),
170 infix(left(3), just(Token::Slash), |lhs, _, rhs, e| {
171 ExprKind::binop(bump, e.span(), BinOp::Div, lhs, rhs)
172 }),
173 infix(left(2), just(Token::Plus), |lhs, _, rhs, e| {
174 ExprKind::binop(bump, e.span(), BinOp::Add, lhs, rhs)
175 }),
176 infix(left(2), just(Token::Minus), |lhs, _, rhs, e| {
177 ExprKind::binop(bump, e.span(), BinOp::Sub, lhs, rhs)
178 }),
179 infix(left(1), just(Token::EqEq), |lhs, _, rhs, e| {
180 ExprKind::binop(bump, e.span(), BinOp::Eq, lhs, rhs)
181 }),
182 infix(left(1), just(Token::BangEq), |lhs, _, rhs, e| {
183 ExprKind::binop(bump, e.span(), BinOp::NotEq, lhs, rhs)
184 }),
185 infix(left(1), just(Token::Lt), |lhs, _, rhs, e| {
186 ExprKind::binop(bump, e.span(), BinOp::Lt, lhs, rhs)
187 }),
188 infix(left(1), just(Token::Gt), |lhs, _, rhs, e| {
189 ExprKind::binop(bump, e.span(), BinOp::Gt, lhs, rhs)
190 }),
191 infix(left(1), just(Token::LtEq), |lhs, _, rhs, e| {
192 ExprKind::binop(bump, e.span(), BinOp::Le, lhs, rhs)
193 }),
194 infix(left(1), just(Token::GtEq), |lhs, _, rhs, e| {
195 ExprKind::binop(bump, e.span(), BinOp::Ge, lhs, rhs)
196 }),
197 infix(left(0), just(Token::AndAnd), |lhs, _, rhs, e| {
198 ExprKind::binop(bump, e.span(), BinOp::And, lhs, rhs)
199 }),
200 infix(left(0), just(Token::PipePipe), |lhs, _, rhs, e| {
201 ExprKind::binop(bump, e.span(), BinOp::Or, lhs, rhs)
202 }),
203 ))
204 .boxed()
205}
206
207fn if_parser<'tok, 'src, 'bump, I>(
208 bump: &'bump Bump,
209 pratt_expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
210 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
211) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
212where
213 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
214 'src: 'tok,
215 'bump: 'tok,
216{
217 just(Token::If)
218 .ignore_then(pratt_expr.clone())
219 .then_ignore(just(Token::Then))
220 .then(pratt_expr)
221 .then_ignore(just(Token::Else))
222 .then(expr)
223 .map_with(|((cond, then_expr), else_expr), e| {
224 ExprKind::if_expr(bump, e.span(), cond, then_expr, else_expr)
225 })
226 .boxed()
227}
228
229fn match_parser<'tok, 'src, 'bump, I>(
230 bump: &'bump Bump,
231 pratt_expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
232 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
233) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
234where
235 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
236 'src: 'tok,
237 'bump: 'tok,
238{
239 let arm = just(Token::Pipe)
240 .ignore_then(pat_parser())
241 .then_ignore(just(Token::StrongArrow))
242 .then(expr);
243
244 just(Token::Match)
245 .ignore_then(pratt_expr)
246 .then(arm.repeated().at_least(1).collect::<Vec<_>>())
247 .map_with(|(scrutinee, arms), e| {
248 let arms = bump.alloc_slice_fill_iter(arms.into_iter());
249 ExprKind::match_expr(bump, e.span(), scrutinee, arms)
250 })
251 .boxed()
252}
253
254fn let_parser<'tok, 'src, 'bump, I>(
255 bump: &'bump Bump,
256 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
257) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
258where
259 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
260 'src: 'tok,
261 'bump: 'tok,
262{
263 just(Token::Let)
264 .then(just(Token::Rec).or_not())
265 .then(ident_parser())
266 .then(ident_parser().repeated().collect::<Vec<_>>())
267 .then_ignore(just(Token::Eq))
268 .then(expr.clone())
269 .then_ignore(just(Token::In))
270 .then(expr)
271 .map_with(|(((((_, rec), name), params), value), body), e| {
272 let value = params.into_iter().rev().fold(value, |body, param| {
273 ExprKind::lambda(bump, e.span(), param, body)
274 });
275 ExprKind::let_expr(bump, e.span(), name, value, body, rec.is_some())
276 })
277 .boxed()
278}
279
280fn lambda_parser<'tok, 'src, 'bump, I>(
281 bump: &'bump Bump,
282 expr: BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>,
283) -> BoxedParser<'tok, 'src, 'bump, I, Expr<'bump>>
284where
285 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
286 'src: 'tok,
287 'bump: 'tok,
288{
289 just(Token::Pipe)
290 .ignore_then(ident_parser())
291 .then_ignore(just(Token::Pipe))
292 .then(expr)
293 .map_with(|(param, body), e| ExprKind::lambda(bump, e.span(), param, body))
294 .boxed()
295}
296
297pub fn parser<'tok, 'src, 'bump, I>(
298 bump: &'bump Bump,
299) -> impl Parser<'tok, I, Expr<'bump>, Extra<'tok, 'src>>
300where
301 I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
302 'src: 'tok,
303 'bump: 'tok,
304{
305 recursive(|expr| {
306 let expr_boxed = expr.clone().boxed();
307 let pratt = pratt_parser(bump, expr_boxed.clone());
308
309 if_parser(bump, pratt.clone(), expr_boxed.clone())
310 .or(match_parser(bump, pratt.clone(), expr_boxed.clone()))
311 .or(let_parser(bump, expr_boxed.clone()))
312 .or(lambda_parser(bump, expr_boxed))
313 .or(pratt)
314 })
315}