From 1d9277e200e0b43ba1da70fbfa599dc16496c65e Mon Sep 17 00:00:00 2001
From: Slendi <slendi@socopon.com>
Date: Wed, 28 Feb 2024 15:07:07 +0200
Subject: [PATCH] Add return statement

Signed-off-by: Slendi <slendi@socopon.com>
---
 ast.odin              | 11 +++++++++++
 parser.odin           | 13 +++++++++++++
 test_type_checker.cat |  2 +-
 type_checker.odin     | 38 +++++++++++++++++++++++++++++++++++++-
 4 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/ast.odin b/ast.odin
index 79d7fcd..5960dd9 100644
--- a/ast.odin
+++ b/ast.odin
@@ -20,6 +20,7 @@ NodeKind :: enum {
 	VariableDeclaration,
 	Cast,
 	BitwiseCast,
+	Ret,
 
 	//
 	Function,
@@ -301,6 +302,16 @@ node_create_bitwise_cast :: proc(
 	return
 }
 
+node_create_ret :: proc(range: TextRange, value: ^Node) -> (ret: ^Node) {
+	ret = new(Node)
+	ret^ = {
+		kind = .Ret,
+		range = range,
+		children = {value},
+	}
+	return
+}
+
 node_create_struct_enum_or_union :: proc(
 	range: TextRange,
 	kind: NodeKind,
diff --git a/parser.odin b/parser.odin
index 355428e..76277b1 100644
--- a/parser.odin
+++ b/parser.odin
@@ -101,6 +101,8 @@ parser_parse_statement :: proc(parser: ^Parser) -> (ret: ^Node) {
 		ret = parser_parse_enum_definition(parser)
 	} else if parser.tok.kind == .Union {
 		ret = parser_parse_union_definition(parser)
+	} else if parser.tok.kind == .Ret {
+		ret = parser_parse_ret(parser)
 	} else {
 		ret = parser_parse_expression(parser)
 		expect(parser, .Semicolon)
@@ -111,6 +113,17 @@ parser_parse_statement :: proc(parser: ^Parser) -> (ret: ^Node) {
 	return
 }
 
+@(private = "file")
+parser_parse_ret :: proc(parser: ^Parser) -> ^Node {
+	range := parser.tok.range
+	expect(parser, .Ret)
+	if parser.tok.kind == .Semicolon {
+		return node_create_ret(range, nil)
+	}
+	ret := parser_parse_expression(parser)
+	return node_create_ret(range, ret)
+}
+
 @(private = "file")
 parser_parse_struct_definition :: proc(parser: ^Parser) -> ^Node {
 	range := parser.tok.range
diff --git a/test_type_checker.cat b/test_type_checker.cat
index d4042c2..a78e1ff 100644
--- a/test_type_checker.cat
+++ b/test_type_checker.cat
@@ -2,7 +2,7 @@ let asdf := 123.0
 let poop :: 12 as f32 + 2.0 * asdf
 
 fn name(a b: f32) i32 {
-    a + b
+    ret (a + b) as i32
 }
 
 name 123 as f32 456.0
diff --git a/type_checker.odin b/type_checker.odin
index a8e5d36..2900fe1 100644
--- a/type_checker.odin
+++ b/type_checker.odin
@@ -7,6 +7,7 @@ Scope :: struct {
     function_definitions: map[int]^FunctionType, // A map to nodes which are the function definitions
     variable_definitions: map[int]^Type, // A map to types
     variable_mutability_definitions: map[int]bool, // A map to a variable's mutability
+    function_return_type: ^Type,
 }
 
 @(private = "file")
@@ -66,6 +67,8 @@ scope_enter :: proc() {
     append(&scope_stack, Scope{})
     scope_stack[len(scope_stack) - 1].function_definitions = make(map[int]^FunctionType)
     scope_stack[len(scope_stack) - 1].variable_definitions = make(map[int]^Type)
+    scope_stack[len(scope_stack) - 1].variable_mutability_definitions = make(map[int]bool)
+    scope_stack[len(scope_stack) - 1].function_return_type = nil
 }
 
 scope_leave :: proc() {
@@ -99,6 +102,15 @@ scope_function_lookup :: proc(name: [dynamic]u8) -> ^FunctionType {
     return nil
 }
 
+scope_function_return_type_lookup :: proc() -> ^Type {
+    #reverse for &scope in scope_stack {
+        if scope.function_return_type != nil {
+            return scope.function_return_type
+        }
+    }
+    return nil
+}
+
 type_check_function_call :: proc(ast: ^Node, parent_ast: ^Node, must_be_function := true) -> ^FunctionType {
     name : [dynamic]u8
     if ast.kind == .FunctionCall {
@@ -230,6 +242,28 @@ type_check :: proc(ast: ^Node, parent_ast: ^Node) {
     case .UnaryExpression:
         // FIXME: Verify that the operation is possible
         type_check(ast.children[0], ast)
+    case .Ret:
+        function_return_type := scope_function_return_type_lookup()
+        if function_return_type == nil {
+            append(&g_message_list,
+                message_create(
+                    .Error,
+                    fmt.aprintf("Return statement outside of function"),
+                    ast.range,
+                ),
+            )
+        } else {
+            type_check(ast.children[0], ast)
+            if !compare_types(function_return_type, ast.children[0].return_type) {
+                append(&g_message_list,
+                    message_create(
+                        .Error,
+                        fmt.aprintf("Type mismatch: {} and {}", function_return_type, ast.children[0].return_type),
+                        ast.range,
+                    ),
+                )
+            }
+        }
     case .Cast:
         type_check(ast.children[0], ast)
         type_to := ast_to_type(ast.children[1])
@@ -260,15 +294,17 @@ type_check :: proc(ast: ^Node, parent_ast: ^Node) {
         scope_stack[len(scope_stack) - 1].variable_definitions[get_character_sum_of_dyn_arr(&ast.children[0].value.([dynamic]u8))] = ast.return_type
     case .Function:
         scope_enter()
+        ast.return_type = ast_to_type(ast.children[0])
+        scope_stack[len(scope_stack) - 1].function_return_type = ast.return_type
         for child, i in ast.children {
             if i < 2 {
                 continue
             }
             type_check(child, ast)
             scope_stack[len(scope_stack) - 1].variable_definitions[get_character_sum_of_dyn_arr(&child.children[0].value.([dynamic]u8))] = child.return_type
+            scope_stack[len(scope_stack) - 1].variable_mutability_definitions[get_character_sum_of_dyn_arr(&child.children[0].value.([dynamic]u8))] = true
         }
         type_check(ast.children[1], ast)
-        ast.return_type = ast.children[0].return_type
         scope_leave()
     case:
         fmt.panicf("Unhandled node kind in type_check: {}", ast.kind)