From db6b6e980b9dd2fa5609ce832a86d18f8108cd21 Mon Sep 17 00:00:00 2001
From: Edgar Luque <git@edgarluque.com>
Date: Sun, 4 Feb 2024 15:57:26 +0100
Subject: [PATCH] handle unit return

---
 lib/edlang_codegen_mlir/src/codegen.rs | 19 +++++++++++++------
 lib/edlang_ir/src/lib.rs               |  2 +-
 lib/edlang_lowering/src/lib.rs         | 11 +++++++++--
 programs/simple.ed                     |  6 ++++++
 4 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/lib/edlang_codegen_mlir/src/codegen.rs b/lib/edlang_codegen_mlir/src/codegen.rs
index 5419c4297..e53fb4c26 100644
--- a/lib/edlang_codegen_mlir/src/codegen.rs
+++ b/lib/edlang_codegen_mlir/src/codegen.rs
@@ -136,19 +136,26 @@ fn compile_module(ctx: &ModuleCompileCtx, module: &ir::ModuleBody) {
 }
 
 fn compile_fn_signature(ctx: &ModuleCompileCtx, body: &ir::Body) {
-    let (args, ret_type) = { (body.get_args(), body.ret_type.clone().unwrap()) };
+    let name = ctx.ctx.symbols.get(&body.def_id).unwrap();
+    info!("compiling fn sig: {}", name);
+
+    let (args, ret_type) = { (body.get_args(), body.ret_type.clone()) };
 
     let args: Vec<BasicMetadataTypeEnum> = args
         .iter()
         .map(|x| compile_basic_type(ctx, &x.ty).into())
         .collect();
-    let ret_type = compile_basic_type(ctx, &ret_type);
-    let name = ctx.ctx.symbols.get(&body.def_id).unwrap();
-    info!("compiling fn sig: {}", name);
+    // let ret_type = compile_basic_type(ctx, &ret_type);
+
+    let fn_type = if let ir::TypeKind::Unit = ret_type.kind {
+        ctx.ctx.context.void_type().fn_type(&args, false)
+    } else {
+        compile_basic_type(ctx, &ret_type).fn_type(&args, false)
+    };
 
     ctx.module.add_function(
         name,
-        ret_type.fn_type(&args, false),
+        fn_type,
         Some(if body.is_extern {
             inkwell::module::Linkage::AvailableExternally
         } else if body.is_pub {
@@ -381,7 +388,7 @@ fn compile_type<'a>(
                     .functions
                     .get(def_id)
                     .unwrap();
-                (fn_body.get_args(), fn_body.ret_type.clone().unwrap())
+                (fn_body.get_args(), fn_body.ret_type.clone())
             };
 
             let args: Vec<BasicMetadataTypeEnum> = args
diff --git a/lib/edlang_ir/src/lib.rs b/lib/edlang_ir/src/lib.rs
index f99ca173c..97067ec81 100644
--- a/lib/edlang_ir/src/lib.rs
+++ b/lib/edlang_ir/src/lib.rs
@@ -36,7 +36,7 @@ pub struct Body {
     pub def_id: DefId,
     pub is_pub: bool,
     pub is_extern: bool,
-    pub ret_type: Option<TypeInfo>,
+    pub ret_type: TypeInfo,
     pub locals: SmallVec<[Local; 4]>,
     pub blocks: SmallVec<[BasicBlock; 8]>,
     pub fn_span: Span,
diff --git a/lib/edlang_lowering/src/lib.rs b/lib/edlang_lowering/src/lib.rs
index 196a38b73..6a787cf26 100644
--- a/lib/edlang_lowering/src/lib.rs
+++ b/lib/edlang_lowering/src/lib.rs
@@ -122,7 +122,14 @@ fn lower_function(
 
     let body = ir::Body {
         def_id,
-        ret_type: func.return_type.as_ref().map(|x| lower_type(&mut ctx, x)),
+        ret_type: func
+            .return_type
+            .as_ref()
+            .map(|x| lower_type(&mut ctx, x))
+            .unwrap_or_else(|| TypeInfo {
+                span: None,
+                kind: ir::TypeKind::Unit,
+            }),
         locals: Default::default(),
         blocks: Default::default(),
         fn_span: func.span,
@@ -543,7 +550,7 @@ fn lower_value(
 fn lower_return(builder: &mut BodyBuilder, info: &ast::ReturnStmt) {
     let ret_type = builder.body.ret_type.clone();
     if let Some(value) = &info.value {
-        let rvalue = lower_expr(builder, value, ret_type.as_ref());
+        let rvalue = lower_expr(builder, value, Some(&ret_type));
         let ret_local = builder.ret_local.unwrap();
 
         builder.statements.push(Statement {
diff --git a/programs/simple.ed b/programs/simple.ed
index 7ec1c8162..93fddc209 100644
--- a/programs/simple.ed
+++ b/programs/simple.ed
@@ -5,4 +5,10 @@ mod Main {
         x = 4;
         return x;
     }
+
+     pub fn a() {
+        let mut x: i32 = 2;
+        x = 4;
+        return;
+    }
 }