Allow returns without value HEAD main

2 files changed, 15 insertions, 6 deletions

diff --git a/src/parse/ast/statement.rs b/src/parse/ast/statement.rs
index 7ed91dc..899190b 100644
--- a/src/parse/ast/statement.rs
+++ b/src/parse/ast/statement.rs
@@ -16,7 +16,7 @@ pub enum StatementKind {
     Print(Expression),
     Break(Option<Expression>),
     Continue,
-    Return(Expression),
+    Return(Option<Expression>),
 }
 
 impl Display for Statement {
@@ -39,8 +39,12 @@ impl Statement {
                 expression.nested_fmt(f, depth + 1)?;
             }
             StatementKind::Return(expression) => {
-                writeln!(f, "{}Return:", pad)?;
-                expression.nested_fmt(f, depth + 1)?;
+                if let Some(returned) = expression {
+                    writeln!(f, "{}Return:", pad)?;
+                    returned.nested_fmt(f, depth + 1)?;
+                } else {
+                    writeln!(f, "{}Return", pad)?;
+                }
             }
             StatementKind::Break(expression) => {
                 if let Some(returned_on_break) = expression {
diff --git a/src/parse/parser.rs b/src/parse/parser.rs
index f19dcf1..a9baf7a 100644
--- a/src/parse/parser.rs
+++ b/src/parse/parser.rs
@@ -51,11 +51,16 @@ impl<T: Iterator<Item = Token>> Parser<T> {
 
     fn return_statement(&mut self) -> Result<Statement, ParserError> {
         let return_token = consume!(self, KeywordReturn)?;
-        let expression = self.expression()?;
-        consume!(self, SemiColon)?;
+        let returned = if consume_if!(self, SemiColon).is_none() {
+            let expression = self.expression()?;
+            consume!(self, SemiColon)?;
+            Some(expression)
+        } else {
+            None
+        };
         Ok(Statement {
             at: return_token.location,
-            kind: StatementKind::Return(expression),
+            kind: StatementKind::Return(returned),
         })
     }