Implement `@(fini)` (opposite of `@(init)`)

8 files changed, 103 insertions, 16 deletions

diff --git a/core/runtime/core.odin b/core/runtime/core.odin
index c64ab7d3b..2d20310ae 100644
--- a/core/runtime/core.odin
+++ b/core/runtime/core.odin
@@ -507,11 +507,8 @@ Odin_Endian_Type :: type_of(ODIN_ENDIAN)
 foreign {
 	@(link_name="__$startup_runtime")
 	_startup_runtime :: proc "odin" () ---
-}
-
-@(link_name="__$cleanup_runtime")
-_cleanup_runtime :: proc() {
-	default_temp_allocator_destroy(&global_default_temp_allocator_data)
+	@(link_name="__$cleanup_runtime")
+	_cleanup_runtime :: proc "odin" () ---
 }
 
 _cleanup_runtime_contextless :: proc "contextless" () {
diff --git a/core/runtime/default_temporary_allocator.odin b/core/runtime/default_temporary_allocator.odin
index b71cd103a..85e76df6c 100644
--- a/core/runtime/default_temporary_allocator.odin
+++ b/core/runtime/default_temporary_allocator.odin
@@ -198,3 +198,9 @@ default_temp_allocator :: proc(allocator: ^Default_Temp_Allocator) -> Allocator
 		data = allocator,
 	}
 }
+
+@(fini, private)
+_destroy_temp_allocator_fini :: proc() {
+	default_temp_allocator_destroy(&global_default_temp_allocator_data)
+	print_string("fini")
+}
diff --git a/src/check_decl.cpp b/src/check_decl.cpp
index f0059424e..63e6514e0 100644
--- a/src/check_decl.cpp
+++ b/src/check_decl.cpp
@@ -816,9 +816,14 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 	if (ac.test) {
 		e->flags |= EntityFlag_Test;
 	}
-	if (ac.init) {
+	if (ac.init && ac.fini) {
+		error(e->token, "A procedure cannot be both declared as @(init) and @(fini)");
+	} else if (ac.init) {
 		e->flags |= EntityFlag_Init;
+	} else if (ac.fini) {
+		e->flags |= EntityFlag_Fini;
 	}
+
 	if (ac.set_cold) {
 		e->flags |= EntityFlag_Cold;
 	}
diff --git a/src/checker.cpp b/src/checker.cpp
index 303a2449e..5163fe675 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -1148,6 +1148,7 @@ gb_internal void init_checker_info(CheckerInfo *i) {
 	array_init(&i->variable_init_order, a);
 	array_init(&i->testing_procedures, a, 0, 0);
 	array_init(&i->init_procedures, a, 0, 0);
+	array_init(&i->fini_procedures, a, 0, 0);
 	array_init(&i->required_foreign_imports_through_force, a, 0, 0);
 
 	map_init(&i->objc_msgSend_types);
@@ -1539,7 +1540,7 @@ gb_internal void add_entity_flags_from_file(CheckerContext *c, Entity *e, Scope
 		AstPackage *pkg = c->file->pkg;
 		if (pkg->kind == Package_Init && e->kind == Entity_Procedure && e->token.string == "main") {
 			// Do nothing
-		} else if (e->flags & (EntityFlag_Test|EntityFlag_Init)) {
+		} else if (e->flags & (EntityFlag_Test|EntityFlag_Init|EntityFlag_Fini)) {
 			// Do nothing
 		} else {
 			e->flags |= EntityFlag_Lazy;
@@ -1607,7 +1608,7 @@ gb_internal bool could_entity_be_lazy(Entity *e, DeclInfo *d) {
 		return false;
 	}
 
-	if (e->flags & (EntityFlag_Test|EntityFlag_Init)) {
+	if (e->flags & (EntityFlag_Test|EntityFlag_Init|EntityFlag_Fini)) {
 		return false;
 	} else if (e->kind == Entity_Variable && e->Variable.is_export) {
 		return false;
@@ -2416,6 +2417,28 @@ gb_internal void generate_minimum_dependency_set(Checker *c, Entity *start) {
 					add_dependency_to_set(c, e);
 					array_add(&c->info.init_procedures, e);
 				}
+			} else if (e->flags & EntityFlag_Fini) {
+				Type *t = base_type(e->type);
+				GB_ASSERT(t->kind == Type_Proc);
+
+				bool is_fini = true;
+
+				if (t->Proc.param_count != 0 || t->Proc.result_count != 0) {
+					gbString str = type_to_string(t);
+					error(e->token, "@(fini) procedures must have a signature type with no parameters nor results, got %s", str);
+					gb_string_free(str);
+					is_fini = false;
+				}
+
+				if ((e->scope->flags & (ScopeFlag_File|ScopeFlag_Pkg)) == 0) {
+					error(e->token, "@(fini) procedures must be declared at the file scope");
+					is_fini = false;
+				}
+
+				if (is_fini) {
+					add_dependency_to_set(c, e);
+					array_add(&c->info.fini_procedures, e);
+				}
 			}
 			break;
 		}
@@ -2967,6 +2990,12 @@ gb_internal DECL_ATTRIBUTE_PROC(proc_decl_attribute) {
 		}
 		ac->init = true;
 		return true;
+	} else if (name == "fini") {
+		if (value != nullptr) {
+			error(value, "'%.*s' expects no parameter, or a string literal containing \"file\" or \"package\"", LIT(name));
+		}
+		ac->fini = true;
+		return true;
 	} else if (name == "deferred") {
 		if (value != nullptr) {
 			Operand o = {};
@@ -3608,6 +3637,7 @@ gb_internal void check_collect_value_decl(CheckerContext *c, Ast *decl) {
 	EntityVisiblityKind entity_visibility_kind = c->foreign_context.visibility_kind;
 	bool is_test = false;
 	bool is_init = false;
+	bool is_fini = false;
 
 	for_array(i, vd->attributes) {
 		Ast *attr = vd->attributes[i];
@@ -3667,6 +3697,8 @@ gb_internal void check_collect_value_decl(CheckerContext *c, Ast *decl) {
 				is_test = true;
 			} else if (name == "init") {
 				is_init = true;
+			} else if (name == "fini") {
+				is_fini = true;
 			}
 		}
 	}
@@ -3800,8 +3832,12 @@ gb_internal void check_collect_value_decl(CheckerContext *c, Ast *decl) {
 				if (is_test) {
 					e->flags |= EntityFlag_Test;
 				}
-				if (is_init) {
+				if (is_init && is_fini) {
+					error(name, "A procedure cannot be both declared as @(init) and @(fini)");
+				} else if (is_init) {
 					e->flags |= EntityFlag_Init;
+				} else if (is_fini) {
+					e->flags |= EntityFlag_Fini;
 				}
 			} else if (init->kind == Ast_ProcGroup) {
 				ast_node(pg, ProcGroup, init);
@@ -5627,9 +5663,14 @@ gb_internal GB_COMPARE_PROC(init_procedures_cmp) {
 	return i32_cmp(x->token.pos.offset, y->token.pos.offset);
 }
 
+gb_internal GB_COMPARE_PROC(fini_procedures_cmp) {
+	return init_procedures_cmp(b, a);
+}
+
 
-gb_internal void check_sort_init_procedures(Checker *c) {
+gb_internal void check_sort_init_and_fini_procedures(Checker *c) {
 	gb_sort_array(c->info.init_procedures.data, c->info.init_procedures.count, init_procedures_cmp);
+	gb_sort_array(c->info.fini_procedures.data, c->info.fini_procedures.count, fini_procedures_cmp);
 }
 
 gb_internal void add_type_info_for_type_definitions(Checker *c) {
@@ -5839,8 +5880,8 @@ gb_internal void check_parsed_files(Checker *c) {
 		add_type_and_value(&c->builtin_ctx, u.expr, u.info->mode, u.info->type, u.info->value);
 	}
 
-	TIME_SECTION("sort init procedures");
-	check_sort_init_procedures(c);
+	TIME_SECTION("sort init and fini procedures");
+	check_sort_init_and_fini_procedures(c);
 
 	if (c->info.intrinsics_entry_point_usage.count > 0) {
 		TIME_SECTION("check intrinsics.__entry_point usage");
diff --git a/src/checker.hpp b/src/checker.hpp
index d461b1f6e..4c3b8725a 100644
--- a/src/checker.hpp
+++ b/src/checker.hpp
@@ -117,6 +117,7 @@ struct AttributeContext {
 	bool    disabled_proc       : 1;
 	bool    test                : 1;
 	bool    init                : 1;
+	bool    fini                : 1;
 	bool    set_cold            : 1;
 	u32 optimization_mode; // ProcedureOptimizationMode
 	i64 foreign_import_priority_index;
@@ -351,6 +352,7 @@ struct CheckerInfo {
 
 	Array<Entity *> testing_procedures;
 	Array<Entity *> init_procedures;
+	Array<Entity *> fini_procedures;
 
 	Array<Entity *> definitions;
 	Array<Entity *> entities;
diff --git a/src/entity.cpp b/src/entity.cpp
index 4b0a6a3c8..0c3629b2b 100644
--- a/src/entity.cpp
+++ b/src/entity.cpp
@@ -75,6 +75,7 @@ enum EntityFlag : u64 {
 	EntityFlag_Test          = 1ull<<30,
 	EntityFlag_Init          = 1ull<<31,
 	EntityFlag_Subtype       = 1ull<<32,
+	EntityFlag_Fini          = 1ull<<33,
 	
 	EntityFlag_CustomLinkName = 1ull<<40,
 	EntityFlag_CustomLinkage_Internal = 1ull<<41,
diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index 34c1ec9b4..6d35615a3 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -1161,6 +1161,34 @@ gb_internal lbProcedure *lb_create_startup_runtime(lbModule *main_module, lbProc
 	return p;
 }
 
+gb_internal lbProcedure *lb_create_cleanup_runtime(lbModule *main_module) { // Cleanup Runtime
+	Type *proc_type = alloc_type_proc(nullptr, nullptr, 0, nullptr, 0, false, ProcCC_Odin);
+
+	lbProcedure *p = lb_create_dummy_procedure(main_module, str_lit(LB_CLEANUP_RUNTIME_PROC_NAME), proc_type);
+	p->is_startup = true;
+
+	lb_begin_procedure_body(p);
+
+	CheckerInfo *info = main_module->gen->info;
+
+	for (Entity *e : info->fini_procedures) {
+		lbValue value = lb_find_procedure_value_from_entity(main_module, e);
+		lb_emit_call(p, value, {}, ProcInlining_none);
+	}
+
+	lb_end_procedure_body(p);
+
+	if (!main_module->debug_builder && LLVMVerifyFunction(p->value, LLVMReturnStatusAction)) {
+		gb_printf_err("LLVM CODE GEN FAILED FOR PROCEDURE: %s\n", "main");
+		LLVMDumpValue(p->value);
+		gb_printf_err("\n\n\n\n");
+		LLVMVerifyFunction(p->value, LLVMAbortProcessAction);
+	}
+
+	return p;
+}
+
+
 gb_internal WORKER_TASK_PROC(lb_generate_procedures_and_types_per_module) {
 	lbModule *m = cast(lbModule *)data;
 	for (Entity *e : m->global_procedures_and_types_to_create) {
@@ -1328,6 +1356,7 @@ gb_internal WORKER_TASK_PROC(lb_llvm_function_pass_per_module) {
 	if (m == &m->gen->default_module) {
 		lb_llvm_function_pass_per_function_internal(m, m->gen->startup_type_info);
 		lb_llvm_function_pass_per_function_internal(m, m->gen->startup_runtime);
+		lb_llvm_function_pass_per_function_internal(m, m->gen->cleanup_runtime);
 		lb_llvm_function_pass_per_function_internal(m, m->gen->objc_names);
 	}
 
@@ -1674,7 +1703,7 @@ gb_internal bool lb_llvm_object_generation(lbGenerator *gen, bool do_threading)
 
 
 
-gb_internal lbProcedure *lb_create_main_procedure(lbModule *m, lbProcedure *startup_runtime) {
+gb_internal lbProcedure *lb_create_main_procedure(lbModule *m, lbProcedure *startup_runtime, lbProcedure *cleanup_runtime) {
 	LLVMPassManagerRef default_function_pass_manager = LLVMCreateFunctionPassManagerForModule(m->mod);
 	lb_populate_function_pass_manager(m, default_function_pass_manager, false, build_context.optimization_level);
 	LLVMFinalizeFunctionPassManager(default_function_pass_manager);
@@ -1793,7 +1822,7 @@ gb_internal lbProcedure *lb_create_main_procedure(lbModule *m, lbProcedure *star
 
 
 	if (call_cleanup) {
-		lbValue cleanup_runtime_value = lb_find_runtime_value(m, str_lit("_cleanup_runtime"));
+		lbValue cleanup_runtime_value = {cleanup_runtime->value, cleanup_runtime->type};
 		lb_emit_call(p, cleanup_runtime_value, {}, ProcInlining_none);
 	}
 
@@ -2330,9 +2359,13 @@ gb_internal bool lb_generate_code(lbGenerator *gen) {
 	gen->startup_type_info = lb_create_startup_type_info(default_module);
 	gen->objc_names = lb_create_objc_names(default_module);
 
-	TIME_SECTION("LLVM Runtime Startup Creation (Global Variables)");
+	TIME_SECTION("LLVM Runtime Startup Creation (Global Variables & @(init))");
 	gen->startup_runtime = lb_create_startup_runtime(default_module, gen->startup_type_info, gen->objc_names, global_variables);
 
+	TIME_SECTION("LLVM Runtime Cleanup Creation & @(fini)");
+	gen->cleanup_runtime = lb_create_cleanup_runtime(default_module);
+
+
 	if (build_context.ODIN_DEBUG) {
 		for (auto const &entry : builtin_pkg->scope->elements) {
 			Entity *e = entry.value;
@@ -2352,7 +2385,7 @@ gb_internal bool lb_generate_code(lbGenerator *gen) {
 
 	if (build_context.command_kind == Command_test && !already_has_entry_point) {
 		TIME_SECTION("LLVM main");
-		lb_create_main_procedure(default_module, gen->startup_runtime);
+		lb_create_main_procedure(default_module, gen->startup_runtime, gen->cleanup_runtime);
 	}
 
 	TIME_SECTION("LLVM Procedure Generation (missing)");
diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index 7bf287b49..0b3bc2a5a 100644
--- a/src/llvm_backend.hpp
+++ b/src/llvm_backend.hpp
@@ -208,6 +208,7 @@ struct lbGenerator {
 
 	lbProcedure *startup_type_info;
 	lbProcedure *startup_runtime;
+	lbProcedure *cleanup_runtime;
 	lbProcedure *objc_names;
 };
 
@@ -540,6 +541,7 @@ gb_internal LLVMTypeRef OdinLLVMGetArrayElementType(LLVMTypeRef type);
 gb_internal LLVMTypeRef OdinLLVMGetVectorElementType(LLVMTypeRef type);
 
 #define LB_STARTUP_RUNTIME_PROC_NAME   "__$startup_runtime"
+#define LB_CLEANUP_RUNTIME_PROC_NAME   "__$cleanup_runtime"
 #define LB_STARTUP_TYPE_INFO_PROC_NAME "__$startup_type_info"
 #define LB_TYPE_INFO_DATA_NAME       "__$type_info_data"
 #define LB_TYPE_INFO_TYPES_NAME      "__$type_info_types_data"