4 files changed, 85 insertions, 56 deletions

diff --git a/modules/mono/mono_gd/gd_mono.cpp b/modules/mono/mono_gd/gd_mono.cpp
index d0b5b12d4b..2c88832998 100644
--- a/modules/mono/mono_gd/gd_mono.cpp
+++ b/modules/mono/mono_gd/gd_mono.cpp
@@ -266,6 +266,13 @@ void GDMono::add_assembly(uint32_t p_domain_id, GDMonoAssembly *p_assembly) {
 	assemblies[p_domain_id][p_assembly->get_name()] = p_assembly;
 }
 
+GDMonoAssembly **GDMono::get_loaded_assembly(const String &p_name) {
+
+	MonoDomain *domain = mono_domain_get();
+	uint32_t domain_id = domain ? mono_domain_get_id(domain) : 0;
+	return assemblies[domain_id].getptr(p_name);
+}
+
 bool GDMono::_load_assembly(const String &p_name, GDMonoAssembly **r_assembly) {
 
 	CRASH_COND(!r_assembly);
@@ -278,34 +285,17 @@ bool GDMono::_load_assembly(const String &p_name, GDMonoAssembly **r_assembly) {
 	MonoAssembly *assembly = mono_assembly_load_full(aname, NULL, &status, false);
 	mono_assembly_name_free(aname);
 
-	if (!assembly)
-		return false;
+	ERR_FAIL_NULL_V(assembly, false);
 
 	uint32_t domain_id = mono_domain_get_id(mono_domain_get());
 
 	GDMonoAssembly **stored_assembly = assemblies[domain_id].getptr(p_name);
 
-	if (stored_assembly) {
-		// Loaded by our preload hook (status is not initialized when returning from a preload hook)
-		ERR_FAIL_COND_V((*stored_assembly)->get_assembly() != assembly, false);
-		*r_assembly = *stored_assembly;
-	} else {
-		ERR_FAIL_COND_V(status != MONO_IMAGE_OK, false);
-
-		MonoImage *assembly_image = mono_assembly_get_image(assembly);
-		ERR_FAIL_NULL_V(assembly_image, false);
-
-		const char *path = mono_image_get_filename(assembly_image);
-
-		*r_assembly = memnew(GDMonoAssembly(p_name, path));
-		Error error = (*r_assembly)->wrapper_for_image(assembly_image);
+	ERR_FAIL_COND_V(status != MONO_IMAGE_OK, false);
+	ERR_FAIL_COND_V(stored_assembly == NULL, false);
 
-		if (error != OK) {
-			memdelete(*r_assembly);
-			*r_assembly = NULL;
-			ERR_FAIL_V(false);
-		}
-	}
+	ERR_FAIL_COND_V((*stored_assembly)->get_assembly() != assembly, false);
+	*r_assembly = *stored_assembly;
 
 	if (OS::get_singleton()->is_stdout_verbose())
 		OS::get_singleton()->print(String("Mono: Assembly " + p_name + " loaded from path: " + (*r_assembly)->get_path() + "\n").utf8());
diff --git a/modules/mono/mono_gd/gd_mono.h b/modules/mono/mono_gd/gd_mono.h
index ab96d575e6..b188c0730a 100644
--- a/modules/mono/mono_gd/gd_mono.h
+++ b/modules/mono/mono_gd/gd_mono.h
@@ -122,7 +122,9 @@ public:
 
 	static GDMono *get_singleton() { return singleton; }
 
+	// Do not use these, unless you know what you're doing
 	void add_assembly(uint32_t p_domain_id, GDMonoAssembly *p_assembly);
+	GDMonoAssembly **get_loaded_assembly(const String &p_name);
 
 	_FORCE_INLINE_ bool is_runtime_initialized() const { return runtime_initialized; }
 	_FORCE_INLINE_ bool is_finalizing_scripts_domain() const { return finalizing_scripts_domain; }
diff --git a/modules/mono/mono_gd/gd_mono_assembly.cpp b/modules/mono/mono_gd/gd_mono_assembly.cpp
index a98537b9e1..a623b34b68 100644
--- a/modules/mono/mono_gd/gd_mono_assembly.cpp
+++ b/modules/mono/mono_gd/gd_mono_assembly.cpp
@@ -39,28 +39,60 @@
 #include "../godotsharp_dirs.h"
 #include "gd_mono_class.h"
 
-MonoAssembly *gdmono_load_assembly_from(const String &p_name, const String &p_path) {
+bool GDMonoAssembly::no_search = false;
+Vector<String> GDMonoAssembly::search_dirs;
 
-	MonoDomain *domain = mono_domain_get();
+MonoAssembly *GDMonoAssembly::_search_hook(MonoAssemblyName *aname, void *user_data) {
 
-	GDMonoAssembly *assembly = memnew(GDMonoAssembly(p_name, p_path));
-	Error err = assembly->load(domain);
-	ERR_FAIL_COND_V(err != OK, NULL);
+	(void)user_data; // UNUSED
 
-	GDMono::get_singleton()->add_assembly(mono_domain_get_id(domain), assembly);
+	String name = mono_assembly_name_get_name(aname);
+	bool has_extension = name.ends_with(".dll") || name.ends_with(".exe");
 
-	return assembly->get_assembly();
-}
+	if (no_search)
+		return NULL;
 
-MonoAssembly *gdmono_MonoAssemblyPreLoad(MonoAssemblyName *aname, char **assemblies_path, void *user_data) {
+	no_search = true; // Avoid the recursion madness
 
-	(void)user_data; // UNUSED
+	GDMonoAssembly **loaded_asm = GDMono::get_singleton()->get_loaded_assembly(has_extension ? name.get_basename() : name);
+	if (loaded_asm)
+		return (*loaded_asm)->get_assembly();
+
+	String path;
+	MonoAssembly *res = NULL;
+
+	for (int i = 0; i < search_dirs.size(); i++) {
+		const String &search_dir = search_dirs[i];
 
-	MonoAssembly *assembly_loaded = mono_assembly_loaded(aname);
-	if (assembly_loaded) // Already loaded
-		return assembly_loaded;
+		if (has_extension) {
+			path = search_dir.plus_file(name);
+			if (FileAccess::exists(path)) {
+				res = _load_assembly_from(name.get_basename(), path);
+				break;
+			}
+		} else {
+			path = search_dir.plus_file(name + ".dll");
+			if (FileAccess::exists(path)) {
+				res = _load_assembly_from(name, path);
+				break;
+			}
+
+			path = search_dir.plus_file(name + ".exe");
+			if (FileAccess::exists(path)) {
+				res = _load_assembly_from(name, path);
+				break;
+			}
+		}
+	}
 
-	static Vector<String> search_dirs;
+	no_search = false;
+
+	return res;
+}
+
+MonoAssembly *GDMonoAssembly::_preload_hook(MonoAssemblyName *aname, char **assemblies_path, void *user_data) {
+
+	(void)user_data; // UNUSED
 
 	if (search_dirs.empty()) {
 		search_dirs.push_back(GodotSharpDirs::get_res_temp_assemblies_dir());
@@ -80,35 +112,32 @@ MonoAssembly *gdmono_MonoAssemblyPreLoad(MonoAssemblyName *aname, char **assembl
 		}
 	}
 
-	String name = mono_assembly_name_get_name(aname);
-	bool has_extension = name.ends_with(".dll") || name.ends_with(".exe");
+	return NULL;
+}
 
-	String path;
+MonoAssembly *GDMonoAssembly::_load_assembly_from(const String &p_name, const String &p_path) {
 
-	for (int i = 0; i < search_dirs.size(); i++) {
-		const String &search_dir = search_dirs[i];
+	GDMonoAssembly *assembly = memnew(GDMonoAssembly(p_name, p_path));
 
-		if (has_extension) {
-			path = search_dir.plus_file(name);
-			if (FileAccess::exists(path))
-				return gdmono_load_assembly_from(name.get_basename(), path);
-		} else {
-			path = search_dir.plus_file(name + ".dll");
-			if (FileAccess::exists(path))
-				return gdmono_load_assembly_from(name, path);
+	MonoDomain *domain = mono_domain_get();
 
-			path = search_dir.plus_file(name + ".exe");
-			if (FileAccess::exists(path))
-				return gdmono_load_assembly_from(name, path);
-		}
+	Error err = assembly->load(domain);
+
+	if (err != OK) {
+		memdelete(assembly);
+		ERR_FAIL_V(NULL);
 	}
 
-	return NULL;
+	GDMono::get_singleton()->add_assembly(domain ? mono_domain_get_id(domain) : 0, assembly);
+
+	return assembly->get_assembly();
 }
 
 void GDMonoAssembly::initialize() {
 
-	mono_install_assembly_preload_hook(&gdmono_MonoAssemblyPreLoad, NULL);
+	// TODO refonly as well?
+	mono_install_assembly_preload_hook(&GDMonoAssembly::_preload_hook, NULL);
+	mono_install_assembly_search_hook(&GDMonoAssembly::_search_hook, NULL);
 }
 
 Error GDMonoAssembly::load(MonoDomain *p_domain) {
@@ -153,7 +182,7 @@ no_pdb:
 
 	ERR_FAIL_COND_V(status != MONO_IMAGE_OK || assembly == NULL, ERR_FILE_CANT_OPEN);
 
-	if (mono_image_get_entry_point(image)) {
+	if (p_domain && mono_image_get_entry_point(image)) {
 		// TODO should this be removed? do we want to call main? what other effects does this have?
 		mono_jit_exec(p_domain, assembly, 0, NULL);
 	}
diff --git a/modules/mono/mono_gd/gd_mono_assembly.h b/modules/mono/mono_gd/gd_mono_assembly.h
index 89e091549c..710b674622 100644
--- a/modules/mono/mono_gd/gd_mono_assembly.h
+++ b/modules/mono/mono_gd/gd_mono_assembly.h
@@ -86,6 +86,14 @@ class GDMonoAssembly {
 	Vector<uint8_t> pdb_data;
 #endif
 
+	static bool no_search;
+	static Vector<String> search_dirs;
+
+	static MonoAssembly *_search_hook(MonoAssemblyName *aname, void *user_data);
+	static MonoAssembly *_preload_hook(MonoAssemblyName *aname, char **assemblies_path, void *user_data);
+
+	static MonoAssembly *_load_assembly_from(const String &p_name, const String &p_path);
+
 	friend class GDMono;
 	static void initialize();