diff --git a/src/parser/node_type.rs b/src/parser/node_type.rs index 507b1d6..14e07ec 100644 --- a/src/parser/node_type.rs +++ b/src/parser/node_type.rs @@ -27,6 +27,7 @@ pub struct Function { pub name: String, pub arguments: Vec, pub body: Statement, + pub ret_type: Option, } #[derive(Debug, Eq, PartialEq)] diff --git a/src/parser/rules.rs b/src/parser/rules.rs index b9e04e2..112d4f5 100644 --- a/src/parser/rules.rs +++ b/src/parser/rules.rs @@ -63,12 +63,18 @@ impl Parser { self.match_token(TokenKind::BraceClose)?; + let ty = match self.peek()?.kind { + TokenKind::Colon => Some(self.parse_type()?), + _ => None, + }; + let body = self.parse_block()?; Ok(Function { name: name, arguments: arguments, body: body, + ret_type: ty, }) } diff --git a/src/parser/tests.rs b/src/parser/tests.rs index baf3258..46b31ed 100644 --- a/src/parser/tests.rs +++ b/src/parser/tests.rs @@ -1,4 +1,5 @@ use crate::lexer::*; +use crate::parser::node_type::*; /** * Copyright 2020 Garrit Franke * @@ -591,3 +592,16 @@ fn test_no_function_args_without_type() { let tree = parse(tokens, Some(raw.to_string())); assert!(tree.is_err()) } + +#[test] +fn test_function_with_return_type() { + let raw = " + fn main(x: int): int { + return n + } + "; + let tokens = tokenize(raw); + let tree = parse(tokens, Some(raw.to_string())); + assert!(tree.is_ok()); + assert_eq!(tree.unwrap().func[0].ret_type, Some(Type::Int)); +}