From 9d48259a13e01e59a62673ba96fb9208a8b1c509 Mon Sep 17 00:00:00 2001
From: Edgar Luque <git@edgarluque.com>
Date: Sun, 26 Mar 2023 12:06:08 +0200
Subject: [PATCH] fix double return

---
 example.ed           |  1 +
 src/ast/statement.rs |  2 +-
 src/codegen.rs       | 19 ++++++++++++++-----
 src/grammar.lalrpop  |  2 +-
 src/main.rs          |  9 +++++++--
 5 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/example.ed b/example.ed
index 17ea20b34..e80bb0e81 100644
--- a/example.ed
+++ b/example.ed
@@ -5,4 +5,5 @@ fn add(a: i64, b: i64) -> i64 {
 fn main() {
     let x = 2 + 3;
     let y = add(x, 4);
+    return;
 }
diff --git a/src/ast/statement.rs b/src/ast/statement.rs
index e4b0ffb58..3a848be42 100644
--- a/src/ast/statement.rs
+++ b/src/ast/statement.rs
@@ -10,7 +10,7 @@ pub struct StatementBody {
 pub enum Statement {
     Assignment(StatementBody),
     Definition(StatementBody),
-    Return(SpanValue<Box<Expr>>),
+    Return(Option<SpanValue<Box<Expr>>>),
     Function(Function),
 }
 
diff --git a/src/codegen.rs b/src/codegen.rs
index 76aa946f7..07dea579f 100644
--- a/src/codegen.rs
+++ b/src/codegen.rs
@@ -142,11 +142,16 @@ impl<'ctx> CodeGen<'ctx> {
         }
 
         // todo: check function has return?
+        let mut has_return = false;
+
         for statement in &function.body {
+            if let Statement::Return(_) = statement.value {
+                has_return = true
+            }
             self.compile_statement(&entry_block, statement, &mut variables)?;
         }
 
-        if function.return_type.is_none() {
+        if !has_return {
             self.builder.build_return(None);
         }
 
@@ -176,10 +181,14 @@ impl<'ctx> CodeGen<'ctx> {
                 variables.insert(body.ident.0.value.clone(), result);
             }
             Statement::Return(ret) => {
-                let result = self
-                    .compile_expression(block, ret, variables)?
-                    .expect("should have result");
-                self.builder.build_return(Some(&result));
+                if let Some(ret) = ret {
+                    let result = self
+                        .compile_expression(block, ret, variables)?
+                        .expect("should have result");
+                    self.builder.build_return(Some(&result));
+                } else {
+                    self.builder.build_return(None);
+                }
             }
             Statement::Function(_function) => unreachable!(),
         };
diff --git a/src/grammar.lalrpop b/src/grammar.lalrpop
index 9ff31afa4..c4ecd6fc3 100644
--- a/src/grammar.lalrpop
+++ b/src/grammar.lalrpop
@@ -65,7 +65,7 @@ Function: Function = {
 BasicStatement: SpanValue<Statement> = {
     <l:@L> "let" <i:Identifier> "=" <e:Expr> ";" <r:@R> => SpanValue::new(l, Statement::new_assignment(i, e), r),
     <l:@L> <i:Identifier> "=" <e:Expr> ";" <r:@R> => SpanValue::new(l, Statement::new_definition(i, e), r),
-    <l:@L> "return" <e:Expr> ";" <r:@R> => SpanValue::new(l, Statement::Return(e), r),
+    <l:@L> "return" <e:Expr?> ";" <r:@R> => SpanValue::new(l, Statement::Return(e), r),
 };
 
 Statement: SpanValue<Statement> = {
diff --git a/src/main.rs b/src/main.rs
index 669dd2203..a889008d1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,7 +40,7 @@ enum Commands {
         input: PathBuf,
 
         /// Output optimized llvm ir.
-        #[arg(short, long)]
+        #[arg(long)]
         optimize: bool,
 
         /// The output file. If not specified its output will be stdout.
@@ -98,7 +98,12 @@ fn main() -> Result<()> {
             let program = ProgramData::new(&str_path, &code);
             check_program(&program, &ast);
         }
-        Commands::Compile { input, output, debug: _, optimize: _ } => {
+        Commands::Compile {
+            input,
+            output,
+            debug: _,
+            optimize: _,
+        } => {
             let code = fs::read_to_string(&input)?;
 
             let parser = grammar::ProgramParser::new();