From 4d4ec47690e424953c93c1410d0054c79e8d9c49 Mon Sep 17 00:00:00 2001
From: rune-scape <allie.smith.epic@gmail.com>
Date: Wed, 27 Dec 2023 12:05:20 -0800
Subject: [PATCH] GDScript: Lambda hotswap fixes

---
 modules/gdscript/gdscript.cpp                 | 129 ++++++------------
 modules/gdscript/gdscript.h                   |  42 +++---
 modules/gdscript/gdscript_compiler.cpp        |  27 ++--
 modules/gdscript/gdscript_compiler.h          |  20 +--
 modules/gdscript/gdscript_lambda_callable.cpp |  32 +----
 modules/gdscript/gdscript_lambda_callable.h   |  14 +-
 6 files changed, 88 insertions(+), 176 deletions(-)

diff --git a/modules/gdscript/gdscript.cpp b/modules/gdscript/gdscript.cpp
index ec1682d47081..2de24b5413aa 100644
--- a/modules/gdscript/gdscript.cpp
+++ b/modules/gdscript/gdscript.cpp
@@ -1391,51 +1391,43 @@ String GDScript::debug_get_script_name(const Ref<Script> &p_script) {
 }
 #endif
 
-GDScript::UpdatableFuncPtr GDScript::func_ptrs_to_update_main_thread;
-thread_local GDScript::UpdatableFuncPtr *GDScript::func_ptrs_to_update_thread_local = nullptr;
-
-GDScript::UpdatableFuncPtrElement GDScript::_add_func_ptr_to_update(GDScriptFunction **p_func_ptr_ptr) {
-	UpdatableFuncPtrElement result = {};
-
-	{
-		MutexLock lock(func_ptrs_to_update_thread_local->mutex);
-		result.element = func_ptrs_to_update_thread_local->ptrs.push_back(p_func_ptr_ptr);
-		result.func_ptr = func_ptrs_to_update_thread_local;
-
-		if (likely(func_ptrs_to_update_thread_local->initialized)) {
-			return result;
-		}
-
-		func_ptrs_to_update_thread_local->initialized = true;
+GDScript::UpdatableFuncPtr::UpdatableFuncPtr(GDScriptFunction *p_function) {
+	if (p_function == nullptr) {
+		return;
 	}
 
-	MutexLock lock(func_ptrs_to_update_mutex);
-	func_ptrs_to_update.push_back(func_ptrs_to_update_thread_local);
-	func_ptrs_to_update_thread_local->rc++;
-
-	return result;
-}
+	ptr = p_function;
+	script = ptr->get_script();
+	ERR_FAIL_NULL(script);
 
-void GDScript::_remove_func_ptr_to_update(const UpdatableFuncPtrElement &p_func_ptr_element) {
-	ERR_FAIL_NULL(p_func_ptr_element.element);
-	ERR_FAIL_NULL(p_func_ptr_element.func_ptr);
-	MutexLock lock(p_func_ptr_element.func_ptr->mutex);
-	p_func_ptr_element.element->erase();
+	MutexLock script_lock(script->func_ptrs_to_update_mutex);
+	list_element = script->func_ptrs_to_update.push_back(this);
 }
 
-void GDScript::_fixup_thread_function_bookkeeping() {
-	// Transfer the ownership of these update items to the main thread,
-	// because the current one is dying, leaving theirs orphan, dangling.
+GDScript::UpdatableFuncPtr::~UpdatableFuncPtr() {
+	ERR_FAIL_NULL(script);
 
-	DEV_ASSERT(!Thread::is_main_thread());
+	if (list_element) {
+		MutexLock script_lock(script->func_ptrs_to_update_mutex);
+		list_element->erase();
+		list_element = nullptr;
+	}
+}
 
-	MutexLock lock(func_ptrs_to_update_main_thread.mutex);
-	MutexLock lock2(func_ptrs_to_update_thread_local->mutex);
+void GDScript::_recurse_replace_function_ptrs(const HashMap<GDScriptFunction *, GDScriptFunction *> &p_replacements) const {
+	MutexLock lock(func_ptrs_to_update_mutex);
+	for (UpdatableFuncPtr *updatable : func_ptrs_to_update) {
+		HashMap<GDScriptFunction *, GDScriptFunction *>::ConstIterator replacement = p_replacements.find(updatable->ptr);
+		if (replacement) {
+			updatable->ptr = replacement->value;
+		} else {
+			// Probably a lambda from another reload, ignore.
+			updatable->ptr = nullptr;
+		}
+	}
 
-	while (!func_ptrs_to_update_thread_local->ptrs.is_empty()) {
-		List<GDScriptFunction **>::Element *E = func_ptrs_to_update_thread_local->ptrs.front();
-		E->transfer_to_back(&func_ptrs_to_update_main_thread.ptrs);
-		func_ptrs_to_update_thread_local->transferred = true;
+	for (HashMap<StringName, Ref<GDScript>>::ConstIterator subscript = subclasses.begin(); subscript; ++subscript) {
+		subscript->value->_recurse_replace_function_ptrs(p_replacements);
 	}
 }
 
@@ -1457,30 +1449,9 @@ void GDScript::clear(ClearData *p_clear_data) {
 	}
 
 	{
-		MutexLock outer_lock(func_ptrs_to_update_mutex);
+		MutexLock lock(func_ptrs_to_update_mutex);
 		for (UpdatableFuncPtr *updatable : func_ptrs_to_update) {
-			bool destroy = false;
-			{
-				MutexLock inner_lock(updatable->mutex);
-				if (updatable->transferred) {
-					func_ptrs_to_update_main_thread.mutex.lock();
-				}
-				for (GDScriptFunction **func_ptr_ptr : updatable->ptrs) {
-					*func_ptr_ptr = nullptr;
-				}
-				DEV_ASSERT(updatable->rc != 0);
-				updatable->rc--;
-				if (updatable->rc == 0) {
-					destroy = true;
-				}
-				if (updatable->transferred) {
-					func_ptrs_to_update_main_thread.mutex.unlock();
-				}
-			}
-			if (destroy) {
-				DEV_ASSERT(updatable != &func_ptrs_to_update_main_thread);
-				memdelete(updatable);
-			}
+			updatable->ptr = nullptr;
 		}
 	}
 
@@ -1553,6 +1524,13 @@ GDScript::~GDScript() {
 	}
 	destructing = true;
 
+	if (is_print_verbose_enabled()) {
+		MutexLock lock(func_ptrs_to_update_mutex);
+		if (!func_ptrs_to_update.is_empty()) {
+			print_line(vformat("GDScript: %d orphaned lambdas becoming invalid at destruction of script '%s'.", func_ptrs_to_update.size(), fully_qualified_name));
+		}
+	}
+
 	clear();
 
 	{
@@ -2101,33 +2079,6 @@ void GDScriptLanguage::remove_named_global_constant(const StringName &p_name) {
 	named_globals.erase(p_name);
 }
 
-void GDScriptLanguage::thread_enter() {
-	GDScript::func_ptrs_to_update_thread_local = memnew(GDScript::UpdatableFuncPtr);
-}
-
-void GDScriptLanguage::thread_exit() {
-	// This thread may have been created before GDScript was up
-	// (which also means it can't have run any GDScript code at all).
-	if (!GDScript::func_ptrs_to_update_thread_local) {
-		return;
-	}
-
-	GDScript::_fixup_thread_function_bookkeeping();
-
-	bool destroy = false;
-	{
-		MutexLock lock(GDScript::func_ptrs_to_update_thread_local->mutex);
-		DEV_ASSERT(GDScript::func_ptrs_to_update_thread_local->rc != 0);
-		GDScript::func_ptrs_to_update_thread_local->rc--;
-		if (GDScript::func_ptrs_to_update_thread_local->rc == 0) {
-			destroy = true;
-		}
-	}
-	if (destroy) {
-		memdelete(GDScript::func_ptrs_to_update_thread_local);
-	}
-}
-
 void GDScriptLanguage::init() {
 	//populate global constants
 	int gcc = CoreConstants::get_global_constant_count();
@@ -2160,8 +2111,6 @@ void GDScriptLanguage::init() {
 		_add_global(E.name, E.ptr);
 	}
 
-	GDScript::func_ptrs_to_update_thread_local = &GDScript::func_ptrs_to_update_main_thread;
-
 #ifdef TESTS_ENABLED
 	GDScriptTests::GDScriptTestRunner::handle_cmdline();
 #endif
@@ -2211,8 +2160,6 @@ void GDScriptLanguage::finish() {
 	}
 	script_list.clear();
 	function_list.clear();
-
-	DEV_ASSERT(GDScript::func_ptrs_to_update_main_thread.rc == 1);
 }
 
 void GDScriptLanguage::profiling_start() {
diff --git a/modules/gdscript/gdscript.h b/modules/gdscript/gdscript.h
index aba4d7e72177..682fe9b04cf0 100644
--- a/modules/gdscript/gdscript.h
+++ b/modules/gdscript/gdscript.h
@@ -117,29 +117,28 @@ class GDScript : public Script {
 
 	HashMap<GDScriptFunction *, LambdaInfo> lambda_info;
 
-	// List is used here because a ptr to elements are stored, so the memory locations need to be stable
-	struct UpdatableFuncPtr {
-		List<GDScriptFunction **> ptrs;
-		Mutex mutex;
-		bool initialized : 1;
-		bool transferred : 1;
-		uint32_t rc = 1;
-		UpdatableFuncPtr() :
-				initialized(false), transferred(false) {}
-	};
-	struct UpdatableFuncPtrElement {
-		List<GDScriptFunction **>::Element *element = nullptr;
-		UpdatableFuncPtr *func_ptr = nullptr;
+public:
+	class UpdatableFuncPtr {
+		friend class GDScript;
+
+		GDScriptFunction *ptr = nullptr;
+		GDScript *script = nullptr;
+		List<UpdatableFuncPtr *>::Element *list_element = nullptr;
+
+	public:
+		GDScriptFunction *operator->() const { return ptr; }
+		operator GDScriptFunction *() const { return ptr; }
+
+		UpdatableFuncPtr(GDScriptFunction *p_function);
+		~UpdatableFuncPtr();
 	};
-	static UpdatableFuncPtr func_ptrs_to_update_main_thread;
-	static thread_local UpdatableFuncPtr *func_ptrs_to_update_thread_local;
+
+private:
+	// List is used here because a ptr to elements are stored, so the memory locations need to be stable
 	List<UpdatableFuncPtr *> func_ptrs_to_update;
 	Mutex func_ptrs_to_update_mutex;
 
-	UpdatableFuncPtrElement _add_func_ptr_to_update(GDScriptFunction **p_func_ptr_ptr);
-	static void _remove_func_ptr_to_update(const UpdatableFuncPtrElement &p_func_ptr_element);
-
-	static void _fixup_thread_function_bookkeeping();
+	void _recurse_replace_function_ptrs(const HashMap<GDScriptFunction *, GDScriptFunction *> &p_replacements) const;
 
 #ifdef TOOLS_ENABLED
 	// For static data storage during hot-reloading.
@@ -561,11 +560,6 @@ class GDScriptLanguage : public ScriptLanguage {
 	virtual void add_named_global_constant(const StringName &p_name, const Variant &p_value) override;
 	virtual void remove_named_global_constant(const StringName &p_name) override;
 
-	/* MULTITHREAD FUNCTIONS */
-
-	virtual void thread_enter() override;
-	virtual void thread_exit() override;
-
 	/* DEBUGGER FUNCTIONS */
 
 	virtual String debug_get_error() const override;
diff --git a/modules/gdscript/gdscript_compiler.cpp b/modules/gdscript/gdscript_compiler.cpp
index 7980f020b839..5cf6c9ba408b 100644
--- a/modules/gdscript/gdscript_compiler.cpp
+++ b/modules/gdscript/gdscript_compiler.cpp
@@ -1371,7 +1371,7 @@ GDScriptCodeGenerator::Address GDScriptCompiler::_parse_expression(CodeGen &code
 				return GDScriptCodeGenerator::Address();
 			}
 
-			main_script->lambda_info.insert(function, { lambda->captures.size(), lambda->use_self });
+			codegen.script->lambda_info.insert(function, { lambda->captures.size(), lambda->use_self });
 			gen->write_lambda(result, function, captures, lambda->use_self);
 
 			for (int i = 0; i < captures.size(); i++) {
@@ -3046,7 +3046,7 @@ GDScriptCompiler::FunctionLambdaInfo GDScriptCompiler::_get_function_replacement
 	FunctionLambdaInfo info;
 	info.function = p_func;
 	info.parent = p_parent_func;
-	info.script = p_parent_func;
+	info.script = p_func->get_script();
 	info.name = p_func->get_name();
 	info.line = p_func->_initial_line;
 	info.index = p_index;
@@ -3057,10 +3057,14 @@ GDScriptCompiler::FunctionLambdaInfo GDScriptCompiler::_get_function_replacement
 	info.default_arg_count = p_func->_default_arg_count;
 	info.sublambdas = _get_function_lambda_replacement_info(p_func, p_depth, p_parent_func);
 
-	GDScript::LambdaInfo *extra_info = main_script->lambda_info.getptr(p_func);
+	ERR_FAIL_NULL_V(info.script, info);
+	GDScript::LambdaInfo *extra_info = info.script->lambda_info.getptr(p_func);
 	if (extra_info != nullptr) {
 		info.capture_count = extra_info->capture_count;
 		info.use_self = extra_info->use_self;
+	} else {
+		info.capture_count = 0;
+		info.use_self = false;
 	}
 
 	return info;
@@ -3195,22 +3199,7 @@ Error GDScriptCompiler::compile(const GDScriptParser *p_parser, GDScript *p_scri
 
 	HashMap<GDScriptFunction *, GDScriptFunction *> func_ptr_replacements;
 	_get_function_ptr_replacements(func_ptr_replacements, old_lambda_info, &new_lambda_info);
-
-	{
-		MutexLock outer_lock(main_script->func_ptrs_to_update_mutex);
-		for (GDScript::UpdatableFuncPtr *updatable : main_script->func_ptrs_to_update) {
-			MutexLock inner_lock(updatable->mutex);
-			for (GDScriptFunction **func_ptr_ptr : updatable->ptrs) {
-				GDScriptFunction **replacement = func_ptr_replacements.getptr(*func_ptr_ptr);
-				if (replacement != nullptr) {
-					*func_ptr_ptr = *replacement;
-				} else {
-					// Probably a lambda from another reload, ignore.
-					*func_ptr_ptr = nullptr;
-				}
-			}
-		}
-	}
+	main_script->_recurse_replace_function_ptrs(func_ptr_replacements);
 
 	if (has_static_data && !root->annotated_static_unload) {
 		GDScriptCache::add_static_script(p_script);
diff --git a/modules/gdscript/gdscript_compiler.h b/modules/gdscript/gdscript_compiler.h
index fd6b22f52746..0adbe1ed8e20 100644
--- a/modules/gdscript/gdscript_compiler.h
+++ b/modules/gdscript/gdscript_compiler.h
@@ -45,19 +45,19 @@ class GDScriptCompiler {
 	GDScript *main_script = nullptr;
 
 	struct FunctionLambdaInfo {
-		GDScriptFunction *function;
-		GDScriptFunction *parent;
-		Ref<GDScript> script;
+		GDScriptFunction *function = nullptr;
+		GDScriptFunction *parent = nullptr;
+		GDScript *script = nullptr;
 		StringName name;
-		int line;
-		int index;
-		int depth;
+		int line = 0;
+		int index = 0;
+		int depth = 0;
 		//uint64_t code_hash;
 		//int code_size;
-		int capture_count;
-		int use_self;
-		int arg_count;
-		int default_arg_count;
+		int capture_count = 0;
+		bool use_self = false;
+		int arg_count = 0;
+		int default_arg_count = 0;
 		//Vector<GDScriptDataType> argument_types;
 		//GDScriptDataType return_type;
 		Vector<FunctionLambdaInfo> sublambdas;
diff --git a/modules/gdscript/gdscript_lambda_callable.cpp b/modules/gdscript/gdscript_lambda_callable.cpp
index 547f5607d33d..f6fa17c84f90 100644
--- a/modules/gdscript/gdscript_lambda_callable.cpp
+++ b/modules/gdscript/gdscript_lambda_callable.cpp
@@ -139,20 +139,14 @@ void GDScriptLambdaCallable::call(const Variant **p_arguments, int p_argcount, V
 	}
 }
 
-GDScriptLambdaCallable::GDScriptLambdaCallable(Ref<GDScript> p_script, GDScriptFunction *p_function, const Vector<Variant> &p_captures) {
+GDScriptLambdaCallable::GDScriptLambdaCallable(Ref<GDScript> p_script, GDScriptFunction *p_function, const Vector<Variant> &p_captures) :
+		function(p_function) {
 	ERR_FAIL_NULL(p_script.ptr());
 	ERR_FAIL_NULL(p_function);
 	script = p_script;
-	function = p_function;
 	captures = p_captures;
 
 	h = (uint32_t)hash_murmur3_one_64((uint64_t)this);
-
-	updatable_func_ptr_element = p_script->_add_func_ptr_to_update(&function);
-}
-
-GDScriptLambdaCallable::~GDScriptLambdaCallable() {
-	GDScript::_remove_func_ptr_to_update(updatable_func_ptr_element);
 }
 
 bool GDScriptLambdaSelfCallable::compare_equal(const CallableCustom *p_a, const CallableCustom *p_b) {
@@ -264,37 +258,23 @@ void GDScriptLambdaSelfCallable::call(const Variant **p_arguments, int p_argcoun
 	}
 }
 
-GDScriptLambdaSelfCallable::GDScriptLambdaSelfCallable(Ref<RefCounted> p_self, GDScriptFunction *p_function, const Vector<Variant> &p_captures) {
+GDScriptLambdaSelfCallable::GDScriptLambdaSelfCallable(Ref<RefCounted> p_self, GDScriptFunction *p_function, const Vector<Variant> &p_captures) :
+		function(p_function) {
 	ERR_FAIL_NULL(p_self.ptr());
 	ERR_FAIL_NULL(p_function);
 	reference = p_self;
 	object = p_self.ptr();
-	function = p_function;
 	captures = p_captures;
 
 	h = (uint32_t)hash_murmur3_one_64((uint64_t)this);
-
-	GDScript *gds = p_function->get_script();
-	if (gds != nullptr) {
-		updatable_func_ptr_element = gds->_add_func_ptr_to_update(&function);
-	}
 }
 
-GDScriptLambdaSelfCallable::GDScriptLambdaSelfCallable(Object *p_self, GDScriptFunction *p_function, const Vector<Variant> &p_captures) {
+GDScriptLambdaSelfCallable::GDScriptLambdaSelfCallable(Object *p_self, GDScriptFunction *p_function, const Vector<Variant> &p_captures) :
+		function(p_function) {
 	ERR_FAIL_NULL(p_self);
 	ERR_FAIL_NULL(p_function);
 	object = p_self;
-	function = p_function;
 	captures = p_captures;
 
 	h = (uint32_t)hash_murmur3_one_64((uint64_t)this);
-
-	GDScript *gds = p_function->get_script();
-	if (gds != nullptr) {
-		updatable_func_ptr_element = gds->_add_func_ptr_to_update(&function);
-	}
-}
-
-GDScriptLambdaSelfCallable::~GDScriptLambdaSelfCallable() {
-	GDScript::_remove_func_ptr_to_update(updatable_func_ptr_element);
 }
diff --git a/modules/gdscript/gdscript_lambda_callable.h b/modules/gdscript/gdscript_lambda_callable.h
index ee7d547544c7..2c5d01aa1670 100644
--- a/modules/gdscript/gdscript_lambda_callable.h
+++ b/modules/gdscript/gdscript_lambda_callable.h
@@ -42,10 +42,9 @@ class GDScriptFunction;
 class GDScriptInstance;
 
 class GDScriptLambdaCallable : public CallableCustom {
-	GDScriptFunction *function = nullptr;
+	GDScript::UpdatableFuncPtr function;
 	Ref<GDScript> script;
 	uint32_t h;
-	GDScript::UpdatableFuncPtrElement updatable_func_ptr_element;
 
 	Vector<Variant> captures;
 
@@ -62,17 +61,18 @@ class GDScriptLambdaCallable : public CallableCustom {
 	StringName get_method() const override;
 	void call(const Variant **p_arguments, int p_argcount, Variant &r_return_value, Callable::CallError &r_call_error) const override;
 
+	GDScriptLambdaCallable(GDScriptLambdaCallable &) = delete;
+	GDScriptLambdaCallable(const GDScriptLambdaCallable &) = delete;
 	GDScriptLambdaCallable(Ref<GDScript> p_script, GDScriptFunction *p_function, const Vector<Variant> &p_captures);
-	virtual ~GDScriptLambdaCallable();
+	virtual ~GDScriptLambdaCallable() = default;
 };
 
 // Lambda callable that references a particular object, so it can use `self` in the body.
 class GDScriptLambdaSelfCallable : public CallableCustom {
-	GDScriptFunction *function = nullptr;
+	GDScript::UpdatableFuncPtr function;
 	Ref<RefCounted> reference; // For objects that are RefCounted, keep a reference.
 	Object *object = nullptr; // For non RefCounted objects, use a direct pointer.
 	uint32_t h;
-	GDScript::UpdatableFuncPtrElement updatable_func_ptr_element;
 
 	Vector<Variant> captures;
 
@@ -88,9 +88,11 @@ class GDScriptLambdaSelfCallable : public CallableCustom {
 	ObjectID get_object() const override;
 	void call(const Variant **p_arguments, int p_argcount, Variant &r_return_value, Callable::CallError &r_call_error) const override;
 
+	GDScriptLambdaSelfCallable(GDScriptLambdaSelfCallable &) = delete;
+	GDScriptLambdaSelfCallable(const GDScriptLambdaSelfCallable &) = delete;
 	GDScriptLambdaSelfCallable(Ref<RefCounted> p_self, GDScriptFunction *p_function, const Vector<Variant> &p_captures);
 	GDScriptLambdaSelfCallable(Object *p_self, GDScriptFunction *p_function, const Vector<Variant> &p_captures);
-	virtual ~GDScriptLambdaSelfCallable();
+	virtual ~GDScriptLambdaSelfCallable() = default;
 };
 
 #endif // GDSCRIPT_LAMBDA_CALLABLE_H