Rewrite objc SEL/Class register handling code

4 files changed, 96 insertions, 90 deletions

diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index 6f3abc607..396b94f98 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -1126,30 +1126,51 @@ gb_internal lbProcedure *lb_create_objc_names(lbModule *main_module) {
 	return p;
 }
 
-gb_internal void lb_finalize_objc_names(lbProcedure *p) {
+gb_internal void lb_finalize_objc_names(lbGenerator *gen, lbProcedure *p) {
 	if (p == nullptr) {
 		return;
 	}
 	lbModule *m = p->module;
+	GB_ASSERT(m == &p->module->gen->default_module);
 
 	TEMPORARY_ALLOCATOR_GUARD();
 
+	StringSet handled = {};
+	string_set_init(&handled);
+	defer (string_set_destroy(&handled));
+
 	auto args = array_make<lbValue>(temporary_allocator(), 1);
 
 	LLVMSetLinkage(p->value, LLVMInternalLinkage);
 	lb_begin_procedure_body(p);
-	for (auto const &entry : m->objc_classes) {
-		String name = entry.key;
-		args[0] = lb_const_value(m, t_cstring, exact_value_string(name));
-		lbValue ptr = lb_emit_runtime_call(p, "objc_lookUpClass", args);
-		lb_addr_store(p, entry.value.local_module_addr, ptr);
+
+	auto register_thing = [&handled, &m, &args](lbProcedure *p, lbObjCGlobal const &g, char const *call) {
+		if (!string_set_update(&handled, g.name)) {
+			lbAddr addr = {};
+			lbValue *found = string_map_get(&m->members, g.global_name);
+			if (found) {
+				addr = lb_addr(*found);
+			} else {
+				lbValue v = {};
+				LLVMTypeRef t = lb_type(m, g.type);
+				v.value = LLVMAddGlobal(m->mod, t, g.global_name);
+				v.type = alloc_type_pointer(g.type);
+				addr = lb_addr(v);
+				LLVMSetInitializer(v.value, LLVMConstNull(t));
+			}
+
+			args[0] = lb_const_value(m, t_cstring, exact_value_string(g.name));
+			lbValue ptr = lb_emit_runtime_call(p, call, args);
+			lb_addr_store(p, addr, ptr);
+		}
+	};
+
+	for (lbObjCGlobal g = {}; mpsc_dequeue(&gen->objc_classes, &g); /**/) {
+		register_thing(p, g, "objc_lookUpClass");
 	}
 
-	for (auto const &entry : m->objc_selectors) {
-		String name = entry.key;
-		args[0] = lb_const_value(m, t_cstring, exact_value_string(name));
-		lbValue ptr = lb_emit_runtime_call(p, "sel_registerName", args);
-		lb_addr_store(p, entry.value.local_module_addr, ptr);
+	for (lbObjCGlobal g = {}; mpsc_dequeue(&gen->objc_selectors, &g); /**/) {
+		register_thing(p, g, "sel_registerName");
 	}
 
 	lb_end_procedure_body(p);
@@ -2637,7 +2658,7 @@ gb_internal bool lb_generate_code(lbGenerator *gen) {
 
 	if (gen->objc_names) {
 		TIME_SECTION("Finalize objc names");
-		lb_finalize_objc_names(gen->objc_names);
+		lb_finalize_objc_names(gen, gen->objc_names);
 	}
 
 	if (build_context.ODIN_DEBUG) {
diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index f9f96a906..3e01ada5f 100644
--- a/src/llvm_backend.hpp
+++ b/src/llvm_backend.hpp
@@ -143,11 +143,6 @@ struct lbPadType {
 	LLVMTypeRef type;
 };
 
-struct lbObjcRef {
-	Entity * entity;
-	lbAddr local_module_addr;
-};
-
 struct lbModule {
 	LLVMModuleRef mod;
 	LLVMContextRef ctx;
@@ -199,11 +194,8 @@ struct lbModule {
 	PtrMap<void *, LLVMMetadataRef> debug_values; 
 
 
-	RecursiveMutex objc_classes_mutex;
-	RecursiveMutex objc_selectors_mutex;
-
-	StringMap<lbObjcRef> objc_classes;
-	StringMap<lbObjcRef> objc_selectors;
+	StringMap<lbAddr> objc_classes;
+	StringMap<lbAddr> objc_selectors;
 
 	PtrMap<u64/*type hash*/, lbAddr> map_cell_info_map; // address of runtime.Map_Info
 	PtrMap<u64/*type hash*/, lbAddr> map_info_map;      // address of runtime.Map_Cell_Info
@@ -222,6 +214,13 @@ struct lbEntityCorrection {
 	char const *cname;
 };
 
+struct lbObjCGlobal {
+	lbModule *module;
+	gbString  global_name;
+	String    name;
+	Type *    type;
+};
+
 struct lbGenerator : LinkerData {
 	CheckerInfo *info;
 
@@ -239,6 +238,8 @@ struct lbGenerator : LinkerData {
 	lbProcedure *objc_names;
 
 	MPSCQueue<lbEntityCorrection> entities_to_correct_linkage;
+	MPSCQueue<lbObjCGlobal> objc_selectors;
+	MPSCQueue<lbObjCGlobal> objc_classes;
 };
 
 
diff --git a/src/llvm_backend_general.cpp b/src/llvm_backend_general.cpp
index b7f70893f..ce2c70661 100644
--- a/src/llvm_backend_general.cpp
+++ b/src/llvm_backend_general.cpp
@@ -171,6 +171,8 @@ gb_internal bool lb_init_generator(lbGenerator *gen, Checker *c) {
 	}
 
 	mpsc_init(&gen->entities_to_correct_linkage, heap_allocator());
+	mpsc_init(&gen->objc_selectors, heap_allocator());
+	mpsc_init(&gen->objc_classes, heap_allocator());
 
 	return true;
 }
diff --git a/src/llvm_backend_utility.cpp b/src/llvm_backend_utility.cpp
index efe196e58..0decbcdb8 100644
--- a/src/llvm_backend_utility.cpp
+++ b/src/llvm_backend_utility.cpp
@@ -2094,48 +2094,71 @@ gb_internal void lb_set_wasm_export_attributes(LLVMValueRef value, String export
 }
 
 
-
 gb_internal lbAddr lb_handle_objc_find_or_register_selector(lbProcedure *p, String const &name) {
-	mutex_lock(&p->module->objc_selectors_mutex);
-	defer (mutex_unlock(&p->module->objc_selectors_mutex));
-
-	lbObjcRef *found = string_map_get(&p->module->objc_selectors, name);
-
+	lbModule *m = p->module;
+	lbAddr *found = string_map_get(&m->objc_selectors, name);
 	if (found) {
-		return found->local_module_addr;
+		return *found;
 	}
 
 	lbModule *default_module = &p->module->gen->default_module;
-	Entity *entity = {};
 
-	if (default_module != p->module) {
-		found = string_map_get(&default_module->objc_selectors, name);
-		if (found) {
-			entity = found->entity;
-		}
+
+	gbString global_name = gb_string_make(permanent_allocator(), "__$objc_SEL::");
+	global_name = gb_string_append_length(global_name, name.text, name.len);
+
+	LLVMTypeRef t = lb_type(m, t_objc_SEL);
+	lbValue g = {};
+	g.value = LLVMAddGlobal(m->mod, t, global_name);
+	g.type = alloc_type_pointer(t_objc_SEL);
+
+	if (default_module == m) {
+		LLVMSetInitializer(g.value, LLVMConstNull(t));
+		lb_add_member(m, make_string_c(global_name), g);
+	} else {
+		LLVMSetLinkage(g.value, LLVMExternalLinkage);
 	}
 
-	if (!entity) {
-		gbString global_name = gb_string_make(permanent_allocator(), "__$objc_SEL::");
-		global_name = gb_string_append_length(global_name, name.text, name.len);
+	mpsc_enqueue(&m->gen->objc_selectors, lbObjCGlobal{m, global_name, name, t_objc_SEL});
+
+	lbAddr addr = lb_addr(g);
 
-		lbAddr default_addr = lb_add_global_generated_with_name(default_module, t_objc_SEL, {},
-			make_string(cast(u8 const *)global_name, gb_string_length(global_name)),
-			&entity);
+	string_map_set(&m->objc_selectors, name, addr);
 
-		mutex_lock(&default_module->objc_selectors_mutex);
-		string_map_set(&default_module->objc_selectors, name, lbObjcRef{entity, default_addr});
-		mutex_unlock(&default_module->objc_selectors_mutex);
+	return addr;
+}
+
+gb_internal lbAddr lb_handle_objc_find_or_register_class(lbProcedure *p, String const &name) {
+	lbModule *m = p->module;
+	lbAddr *found = string_map_get(&m->objc_classes, name);
+	if (found) {
+		return *found;
 	}
 
-	lbValue ptr = lb_find_value_from_entity(p->module, entity);
-	lbAddr local_addr = lb_addr(ptr);
+	lbModule *default_module = &p->module->gen->default_module;
+
+
+	gbString global_name = gb_string_make(permanent_allocator(), "__$objc_Class::");
+	global_name = gb_string_append_length(global_name, name.text, name.len);
 
-	if (default_module != p->module) {
-		string_map_set(&p->module->objc_selectors, name, lbObjcRef{entity, local_addr});
+	LLVMTypeRef t = lb_type(m, t_objc_Class);
+	lbValue g = {};
+	g.value = LLVMAddGlobal(m->mod, t, global_name);
+	g.type = alloc_type_pointer(t_objc_Class);
+
+	if (default_module == m) {
+		LLVMSetInitializer(g.value, LLVMConstNull(t));
+		lb_add_member(m, make_string_c(global_name), g);
+	} else {
+		LLVMSetLinkage(g.value, LLVMExternalLinkage);
 	}
+	mpsc_enqueue(&m->gen->objc_classes, lbObjCGlobal{m, global_name, name, t_objc_Class});
 
-	return local_addr;
+	lbAddr addr = lb_addr(g);
+
+	string_map_set(&m->objc_classes, name, addr);
+
+	return addr;
 }
 
 gb_internal lbValue lb_handle_objc_find_selector(lbProcedure *p, Ast *expr) {
@@ -2164,47 +2187,6 @@ gb_internal lbValue lb_handle_objc_register_selector(lbProcedure *p, Ast *expr)
 	return lb_addr_load(p, dst);
 }
 
-gb_internal lbAddr lb_handle_objc_find_or_register_class(lbProcedure *p, String const &name) {
-	mutex_lock(&p->module->objc_classes_mutex);
-	defer (mutex_unlock(&p->module->objc_classes_mutex));
-
-	lbObjcRef *found = string_map_get(&p->module->objc_classes, name);
-	if (found) {
-		return found->local_module_addr;
-	}
-
-	lbModule *default_module = &p->module->gen->default_module;
-	Entity *entity = {};
-
-	if (default_module != p->module) {
-		found = string_map_get(&default_module->objc_classes, name);
-		if (found) {
-			entity = found->entity;
-		}
-	}
-
-	if (!entity) {
-		gbString global_name = gb_string_make(permanent_allocator(), "__$objc_Class::");
-		global_name = gb_string_append_length(global_name, name.text, name.len);
-
-		lbAddr default_addr = lb_add_global_generated_with_name(default_module, t_objc_Class, {},
-			make_string(cast(u8 const *)global_name, gb_string_length(global_name)),
-			&entity);
-
-		mutex_lock(&default_module->objc_classes_mutex);
-		string_map_set(&default_module->objc_classes, name, lbObjcRef{entity, default_addr});
-		mutex_unlock(&default_module->objc_classes_mutex);
-	}
-
-	lbValue ptr = lb_find_value_from_entity(p->module, entity);
-	lbAddr local_addr = lb_addr(ptr);
-
-	if (default_module != p->module) {
-		string_map_set(&p->module->objc_classes, name, lbObjcRef{entity, local_addr});
-	}
-
-	return local_addr;
-}
 
 gb_internal lbValue lb_handle_objc_find_class(lbProcedure *p, Ast *expr) {
 	ast_node(ce, CallExpr, expr);