Skip to main content

aiken_lang/parser/definition/
function.rs

1use crate::{
2    ast,
3    expr::UntypedExpr,
4    parser::{annotation, error::ParseError, expr, pattern, token::Token, utils},
5};
6use chumsky::prelude::*;
7
8pub fn parser() -> impl Parser<Token, ast::UntypedDefinition, Error = ParseError> {
9    utils::optional_flag(Token::Pub)
10        .then_ignore(just(Token::Fn))
11        .then(select! {Token::Name {name} => name})
12        .then(
13            param(false)
14                .separated_by(just(Token::Comma))
15                .allow_trailing()
16                .delimited_by(just(Token::LeftParen), just(Token::RightParen))
17                .map_with_span(|arguments, span| (arguments, span)),
18        )
19        .then(just(Token::RArrow).ignore_then(annotation()).or_not())
20        .then(
21            expr::sequence()
22                .or_not()
23                .delimited_by(just(Token::LeftBrace), just(Token::RightBrace)),
24        )
25        .map_with_span(
26            |((((public, name), (arguments, args_span)), return_annotation), body), span| {
27                ast::UntypedDefinition::Fn(ast::Function {
28                    arguments,
29                    body: body.unwrap_or_else(|| UntypedExpr::todo(None, span)),
30                    doc: None,
31                    location: ast::Span {
32                        start: span.start,
33                        end: return_annotation
34                            .as_ref()
35                            .map(|l| l.location().end)
36                            .unwrap_or_else(|| args_span.end),
37                    },
38                    end_position: span.end - 1,
39                    name,
40                    public,
41                    return_annotation,
42                    return_type: (),
43                    on_test_failure: ast::OnTestFailure::FailImmediately,
44                })
45            },
46        )
47}
48
49pub fn param(is_validator_param: bool) -> impl Parser<Token, ast::UntypedArg, Error = ParseError> {
50    choice((
51        select! {Token::Name {name} => name}
52            .then(select! {Token::DiscardName {name} => name})
53            .map_with_span(|(label, name), span| {
54                ast::ArgBy::ByName(ast::ArgName::Discarded {
55                    label,
56                    name,
57                    location: span,
58                })
59            }),
60        select! {Token::DiscardName {name} => name}.map_with_span(|name, span| {
61            ast::ArgBy::ByName(ast::ArgName::Discarded {
62                label: name.clone(),
63                name,
64                location: span,
65            })
66        }),
67        select! {Token::Name {name} => name}
68            .then(select! {Token::Name {name} => name})
69            .map_with_span(|(label, name), span| {
70                ast::ArgBy::ByName(ast::ArgName::Named {
71                    label,
72                    name,
73                    location: span,
74                })
75            }),
76        select! {Token::Name {name} => name}.map_with_span(|name, span| {
77            ast::ArgBy::ByName(ast::ArgName::Named {
78                label: name.clone(),
79                name,
80                location: span,
81            })
82        }),
83        pattern().map(ast::ArgBy::ByPattern),
84    ))
85    .then(just(Token::Colon).ignore_then(annotation()).or_not())
86    .map_with_span(move |(by, annotation), span| ast::UntypedArg {
87        location: span,
88        annotation,
89        doc: None,
90        is_validator_param,
91        by,
92    })
93}
94
95#[cfg(test)]
96mod tests {
97    use crate::assert_definition;
98
99    #[test]
100    fn function_empty() {
101        assert_definition!(
102            r#"
103            pub fn run() {}
104            "#
105        );
106    }
107
108    #[test]
109    fn function_non_public() {
110        assert_definition!(
111            r#"
112            fn run() {}
113            "#
114        );
115    }
116
117    #[test]
118    fn function_assignment_only() {
119        assert_definition!(
120            r#"
121            fn run() {
122              let x = 1 + 1
123            }
124            "#
125        );
126    }
127
128    #[test]
129    fn function_by_pattern_no_annotation() {
130        assert_definition!(
131            r#"
132            fn foo(Foo { my_field }) {
133                my_field * 2
134            }
135            "#
136        );
137    }
138
139    #[test]
140    fn function_by_pattern_with_annotation() {
141        assert_definition!(
142            r#"
143            fn foo(Foo { my_field }: Foo) {
144                my_field * 2
145            }
146            "#
147        );
148    }
149    #[test]
150    fn function_by_pattern_with_alias() {
151        assert_definition!(
152            r#"
153            fn foo(Foo { my_field, .. } as x) {
154                my_field * x.my_other_field
155            }
156            "#
157        );
158    }
159
160    #[test]
161    fn grouping() {
162        assert_definition!(
163            r#"
164            pub fn a0p(pair: Pair<ByteArray, Dict<ByteArray, Int>>) {
165              (pair.2nd |> unsingleton).2nd
166            }
167            "#
168        );
169    }
170}