fix double return

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;
 }

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),
 }
 

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
+                if let Some(ret) = ret {
-                    .compile_expression(block, ret, variables)?
+                    let result = self
-                    .expect("should have result");
+                        .compile_expression(block, ret, variables)?
-                self.builder.build_return(Some(&result));
+                        .expect("should have result");
+                    self.builder.build_return(Some(&result));
+                } else {
+                    self.builder.build_return(None);
+                }
             }
             Statement::Function(_function) => unreachable!(),
         };

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> = {

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();