diff --git a/src/parser.rs b/src/parser.rs index 10d0c7d..233f3a8 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -208,22 +208,23 @@ impl Parse for FunctionSignature { } #[derive(Debug)] -pub struct Block(Vec); +pub struct Block(Vec, Option); impl Parse for Block { fn parse(mut stream: TokenStream) -> Result { let mut statements = Vec::new(); + let mut return_stmt = None; stream.expect(Token::BraceOpen)?; while !matches!(stream.peek(), Some(Token::BraceClose)) { let statement = stream.parse()?; - if let BlockLevelStatement::Return(_) = &statement { - statements.push(statement); + if let BlockLevelStatement::Return(e) = &statement { + return_stmt = Some(e.clone()); break; // Return has to be the last statement } statements.push(statement); } stream.expect(Token::BraceClose)?; - Ok(Block(statements)) + Ok(Block(statements, return_stmt)) } }