Codegen: Fix bug with recursive functions as variables

2 files changed, 24 insertions, 28 deletions

diff --git a/examples/8.src b/examples/8.src
index 73ea7aa..985ca77 100644
--- a/examples/8.src
+++ b/examples/8.src
@@ -1,14 +1,14 @@
-let fib = (n: i64) => i64 {
-	if n == 0 {
-		return 0;
-	};
-	if n == 1 {
-		return 1;
+let main = () => i64 {
+	let fib = (n: i64) => i64 {
+		if n == 0 {
+			return 0;
+		};
+		if n == 1 {
+			return 1;
+		};
+		return fib(n-2) + fib(n-1);
 	};
-	return fib(n-2) + fib(n-1);
-};
 
-let main = () => i64 {
 	let result = fib(30);
 	print(result);
 	return result;
diff --git a/src/codegen.zig b/src/codegen.zig
index c780713..9e06c8a 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -262,29 +262,22 @@ pub const CodeGen = struct {
                 const function_type = core.LLVMFunctionType(return_type, paramtypes.items.ptr, @intCast(paramtypes.items.len), 0) orelse return CodeGenError.CompilationError;
                 const function = core.LLVMAddFunction(self.llvm_module, try std.fmt.allocPrintZ(self.arena, "{s}", .{name orelse "unnamed_func"}), function_type) orelse return CodeGenError.CompilationError;
                 const function_entry = core.LLVMAppendBasicBlock(function, "entrypoint") orelse return CodeGenError.CompilationError;
-
-                // Needed for recursive functions
-                if (name != null) {
-                    const ptr = self.environment.get_variable(name.?);
-                    // Global fn
-                    if (ptr == null) {
-                        try self.environment.add_variable(name.?, try self.create_variable(.{
-                            .value = function,
-                            .type = function_type,
-                            .stack_level = null,
-                        }));
-                    } else {
-                        _ = core.LLVMBuildStore(self.builder, function, ptr.?.value) orelse return CodeGenError.CompilationError;
-                        ptr.?.type = function_type;
-                        try self.environment.add_variable(name.?, ptr.?);
-                    }
-                }
-
                 core.LLVMPositionBuilderAtEnd(self.builder, function_entry);
 
                 try self.environment.create_scope();
                 defer self.environment.drop_scope();
 
+                const ptr = self.environment.get_variable(name.?);
+
+                // Needed for recursive functions
+                if (name != null) {
+                    try self.environment.add_variable(name.?, try self.create_variable(.{
+                        .value = function,
+                        .type = function_type,
+                        .stack_level = null,
+                    }));
+                }
+
                 const params = try self.arena.alloc(types.LLVMValueRef, function_definition.parameters.len);
                 core.LLVMGetParams(function, params.ptr);
 
@@ -324,7 +317,10 @@ pub const CodeGen = struct {
                         .stack_level = null,
                     });
                 }
-                return self.environment.get_variable(name.?) orelse unreachable;
+
+                _ = core.LLVMBuildStore(self.builder, function, ptr.?.value) orelse return CodeGenError.CompilationError;
+                ptr.?.type = function_type;
+                return ptr.?;
             },
             .FUNCTION_CALL_STATEMENT => |*fn_call| {
                 if (name != null) {