Fix editor needs restart after adding GDExtensions

10 files changed, 196 insertions, 65 deletions

diff --git a/core/extension/gdextension_manager.cpp b/core/extension/gdextension_manager.cpp
index eeae6b1996..01efe0d96e 100644
--- a/core/extension/gdextension_manager.cpp
+++ b/core/extension/gdextension_manager.cpp
@@ -32,14 +32,18 @@
 
 #include "core/extension/gdextension_compat_hashes.h"
 #include "core/extension/gdextension_library_loader.h"
+#include "core/io/dir_access.h"
 #include "core/io/file_access.h"
 #include "core/object/script_language.h"
 
-GDExtensionManager::LoadStatus GDExtensionManager::_load_extension_internal(const Ref<GDExtension> &p_extension) {
+GDExtensionManager::LoadStatus GDExtensionManager::_load_extension_internal(const Ref<GDExtension> &p_extension, bool p_first_load) {
 	if (level >= 0) { // Already initialized up to some level.
-		int32_t minimum_level = p_extension->get_minimum_library_initialization_level();
-		if (minimum_level < MIN(level, GDExtension::INITIALIZATION_LEVEL_SCENE)) {
-			return LOAD_STATUS_NEEDS_RESTART;
+		int32_t minimum_level = 0;
+		if (!p_first_load) {
+			minimum_level = p_extension->get_minimum_library_initialization_level();
+			if (minimum_level < MIN(level, GDExtension::INITIALIZATION_LEVEL_SCENE)) {
+				return LOAD_STATUS_NEEDS_RESTART;
+			}
 		}
 		// Initialize up to current level.
 		for (int32_t i = minimum_level; i <= level; i++) {
@@ -51,10 +55,20 @@ GDExtensionManager::LoadStatus GDExtensionManager::_load_extension_internal(cons
 		gdextension_class_icon_paths[kv.key] = kv.value;
 	}
 
+#ifdef TOOLS_ENABLED
+	// Signals that a new extension is loaded so GDScript can register new class names.
+	emit_signal("extension_loaded", p_extension);
+#endif
+
 	return LOAD_STATUS_OK;
 }
 
 GDExtensionManager::LoadStatus GDExtensionManager::_unload_extension_internal(const Ref<GDExtension> &p_extension) {
+#ifdef TOOLS_ENABLED
+	// Signals that a new extension is unloading so GDScript can unregister class names.
+	emit_signal("extension_unloading", p_extension);
+#endif
+
 	if (level >= 0) { // Already initialized up to some level.
 		// Deinitialize down from current level.
 		for (int32_t i = level; i >= GDExtension::INITIALIZATION_LEVEL_CORE; i--) {
@@ -89,7 +103,7 @@ GDExtensionManager::LoadStatus GDExtensionManager::load_extension_with_loader(co
 		return LOAD_STATUS_FAILED;
 	}
 
-	LoadStatus status = _load_extension_internal(extension);
+	LoadStatus status = _load_extension_internal(extension, true);
 	if (status != LOAD_STATUS_OK) {
 		return status;
 	}
@@ -135,7 +149,7 @@ GDExtensionManager::LoadStatus GDExtensionManager::reload_extension(const String
 		return LOAD_STATUS_FAILED;
 	}
 
-	status = _load_extension_internal(extension);
+	status = _load_extension_internal(extension, false);
 	if (status != LOAD_STATUS_OK) {
 		return status;
 	}
@@ -274,6 +288,71 @@ void GDExtensionManager::reload_extensions() {
 #endif
 }
 
+bool GDExtensionManager::ensure_extensions_loaded(const HashSet<String> &p_extensions) {
+	Vector<String> extensions_added;
+	Vector<String> extensions_removed;
+
+	for (const String &E : p_extensions) {
+		if (!is_extension_loaded(E)) {
+			extensions_added.push_back(E);
+		}
+	}
+
+	Vector<String> loaded_extensions = get_loaded_extensions();
+	for (const String &loaded_extension : loaded_extensions) {
+		if (!p_extensions.has(loaded_extension)) {
+			// The extension may not have a .gdextension file.
+			if (!FileAccess::exists(loaded_extension)) {
+				extensions_removed.push_back(loaded_extension);
+			}
+		}
+	}
+
+	String extension_list_config_file = GDExtension::get_extension_list_config_file();
+	if (p_extensions.size()) {
+		if (extensions_added.size() || extensions_removed.size()) {
+			// Extensions were added or removed.
+			Ref<FileAccess> f = FileAccess::open(extension_list_config_file, FileAccess::WRITE);
+			for (const String &E : p_extensions) {
+				f->store_line(E);
+			}
+		}
+	} else {
+		if (loaded_extensions.size() || FileAccess::exists(extension_list_config_file)) {
+			// Extensions were removed.
+			Ref<DirAccess> da = DirAccess::create(DirAccess::ACCESS_RESOURCES);
+			da->remove(extension_list_config_file);
+		}
+	}
+
+	bool needs_restart = false;
+	for (const String &extension : extensions_added) {
+		GDExtensionManager::LoadStatus st = GDExtensionManager::get_singleton()->load_extension(extension);
+		if (st == GDExtensionManager::LOAD_STATUS_NEEDS_RESTART) {
+			needs_restart = true;
+		}
+	}
+
+	for (const String &extension : extensions_removed) {
+		GDExtensionManager::LoadStatus st = GDExtensionManager::get_singleton()->unload_extension(extension);
+		if (st == GDExtensionManager::LOAD_STATUS_NEEDS_RESTART) {
+			needs_restart = true;
+		}
+	}
+
+#ifdef TOOLS_ENABLED
+	if (extensions_added.size() || extensions_removed.size()) {
+		// Emitting extensions_reloaded so EditorNode can reload Inspector and regenerate documentation.
+		emit_signal("extensions_reloaded");
+
+		// Reload all scripts to clear out old references.
+		callable_mp_static(&GDExtensionManager::_reload_all_scripts).call_deferred();
+	}
+#endif
+
+	return needs_restart;
+}
+
 GDExtensionManager *GDExtensionManager::get_singleton() {
 	return singleton;
 }
@@ -294,6 +373,8 @@ void GDExtensionManager::_bind_methods() {
 	BIND_ENUM_CONSTANT(LOAD_STATUS_NEEDS_RESTART);
 
 	ADD_SIGNAL(MethodInfo("extensions_reloaded"));
+	ADD_SIGNAL(MethodInfo("extension_loaded", PropertyInfo(Variant::OBJECT, "extension", PROPERTY_HINT_RESOURCE_TYPE, "GDExtension")));
+	ADD_SIGNAL(MethodInfo("extension_unloading", PropertyInfo(Variant::OBJECT, "extension", PROPERTY_HINT_RESOURCE_TYPE, "GDExtension")));
 }
 
 GDExtensionManager *GDExtensionManager::singleton = nullptr;
diff --git a/core/extension/gdextension_manager.h b/core/extension/gdextension_manager.h
index b488189604..39a600474c 100644
--- a/core/extension/gdextension_manager.h
+++ b/core/extension/gdextension_manager.h
@@ -54,7 +54,7 @@ public:
 	};
 
 private:
-	LoadStatus _load_extension_internal(const Ref<GDExtension> &p_extension);
+	LoadStatus _load_extension_internal(const Ref<GDExtension> &p_extension, bool p_first_load);
 	LoadStatus _unload_extension_internal(const Ref<GDExtension> &p_extension);
 
 #ifdef TOOLS_ENABLED
@@ -85,6 +85,7 @@ public:
 
 	void load_extensions();
 	void reload_extensions();
+	bool ensure_extensions_loaded(const HashSet<String> &p_extensions);
 
 	GDExtensionManager();
 	~GDExtensionManager();
diff --git a/core/object/class_db.cpp b/core/object/class_db.cpp
index 5c793a676f..c929b29ee9 100644
--- a/core/object/class_db.cpp
+++ b/core/object/class_db.cpp
@@ -271,6 +271,22 @@ void ClassDB::get_extensions_class_list(List<StringName> *p_classes) {
 
 	p_classes->sort_custom<StringName::AlphCompare>();
 }
+
+void ClassDB::get_extension_class_list(const Ref<GDExtension> &p_extension, List<StringName> *p_classes) {
+	OBJTYPE_RLOCK;
+
+	for (const KeyValue<StringName, ClassInfo> &E : classes) {
+		if (E.value.api != API_EXTENSION && E.value.api != API_EDITOR_EXTENSION) {
+			continue;
+		}
+		if (!E.value.gdextension || E.value.gdextension->library != p_extension.ptr()) {
+			continue;
+		}
+		p_classes->push_back(E.key);
+	}
+
+	p_classes->sort_custom<StringName::AlphCompare>();
+}
 #endif
 
 void ClassDB::get_inheriters_from_class(const StringName &p_class, List<StringName> *p_classes) {
diff --git a/core/object/class_db.h b/core/object/class_db.h
index d6a95b58e2..620092a6c4 100644
--- a/core/object/class_db.h
+++ b/core/object/class_db.h
@@ -285,6 +285,7 @@ public:
 	static void get_class_list(List<StringName> *p_classes);
 #ifdef TOOLS_ENABLED
 	static void get_extensions_class_list(List<StringName> *p_classes);
+	static void get_extension_class_list(const Ref<GDExtension> &p_extension, List<StringName> *p_classes);
 	static ObjectGDExtension *get_placeholder_extension(const StringName &p_class);
 #endif
 	static void get_inheriters_from_class(const StringName &p_class, List<StringName> *p_classes);
diff --git a/doc/classes/GDExtensionManager.xml b/doc/classes/GDExtensionManager.xml
index 211bc023c0..97d2d08752 100644
--- a/doc/classes/GDExtensionManager.xml
+++ b/doc/classes/GDExtensionManager.xml
@@ -56,6 +56,20 @@
 		</method>
 	</methods>
 	<signals>
+		<signal name="extension_loaded">
+			<param index="0" name="extension" type="GDExtension" />
+			<description>
+				Emitted after the editor has finished loading a new extension.
+				[b]Note:[/b] This signal is only emitted in editor builds.
+			</description>
+		</signal>
+		<signal name="extension_unloading">
+			<param index="0" name="extension" type="GDExtension" />
+			<description>
+				Emitted before the editor starts unloading an extension.
+				[b]Note:[/b] This signal is only emitted in editor builds.
+			</description>
+		</signal>
 		<signal name="extensions_reloaded">
 			<description>
 				Emitted after the editor has finished reloading one or more extensions.
diff --git a/editor/editor_file_system.cpp b/editor/editor_file_system.cpp
index 474a45cf2b..4a8a352e58 100644
--- a/editor/editor_file_system.cpp
+++ b/editor/editor_file_system.cpp
@@ -243,18 +243,27 @@ void EditorFileSystem::_first_scan_filesystem() {
 	first_scan_root_dir = memnew(ScannedDirectory);
 	first_scan_root_dir->full_path = "res://";
 	HashSet<String> existing_class_names;
+	HashSet<String> extensions;
 
 	ep.step(TTR("Scanning file structure..."), 0, true);
 	nb_files_total = _scan_new_dir(first_scan_root_dir, d);
 
 	// This loads the global class names from the scripts and ensures that even if the
 	// global_script_class_cache.cfg was missing or invalid, the global class names are valid in ScriptServer.
+	// At the same time, to prevent looping multiple times in all files, it looks for extensions.
 	ep.step(TTR("Loading global class names..."), 1, true);
-	_first_scan_process_scripts(first_scan_root_dir, existing_class_names);
+	_first_scan_process_scripts(first_scan_root_dir, existing_class_names, extensions);
 
 	// Removing invalid global class to prevent having invalid paths in ScriptServer.
 	_remove_invalid_global_class_names(existing_class_names);
 
+	// Processing extensions to add new extensions or remove invalid ones.
+	// Important to do it in the first scan so custom types, new class names, custom importers, etc...
+	// from extensions are ready to go before plugins, autoloads and resources validation/importation.
+	// At this point, a restart of the editor should not be needed so we don't use the return value.
+	ep.step(TTR("Verifying GDExtensions..."), 2, true);
+	GDExtensionManager::get_singleton()->ensure_extensions_loaded(extensions);
+
 	// Now that all the global class names should be loaded, create autoloads and plugins.
 	// This is done after loading the global class names because autoloads and plugins can use
 	// global class names.
@@ -267,9 +276,9 @@ void EditorFileSystem::_first_scan_filesystem() {
 	ep.step(TTR("Starting file scan..."), 5, true);
 }
 
-void EditorFileSystem::_first_scan_process_scripts(const ScannedDirectory *p_scan_dir, HashSet<String> &p_existing_class_names) {
+void EditorFileSystem::_first_scan_process_scripts(const ScannedDirectory *p_scan_dir, HashSet<String> &p_existing_class_names, HashSet<String> &p_extensions) {
 	for (ScannedDirectory *scan_sub_dir : p_scan_dir->subdirs) {
-		_first_scan_process_scripts(scan_sub_dir, p_existing_class_names);
+		_first_scan_process_scripts(scan_sub_dir, p_existing_class_names, p_extensions);
 	}
 
 	for (const String &scan_file : p_scan_dir->files) {
@@ -285,6 +294,8 @@ void EditorFileSystem::_first_scan_process_scripts(const ScannedDirectory *p_sca
 			if (!script_class_name.is_empty()) {
 				p_existing_class_names.insert(script_class_name);
 			}
+		} else if (type == SNAME("GDExtension")) {
+			p_extensions.insert(path);
 		}
 	}
 }
@@ -3016,57 +3027,7 @@ bool EditorFileSystem::_scan_extensions() {
 
 	_scan_extensions_dir(d, extensions);
 
-	//verify against loaded extensions
-
-	Vector<String> extensions_added;
-	Vector<String> extensions_removed;
-
-	for (const String &E : extensions) {
-		if (!GDExtensionManager::get_singleton()->is_extension_loaded(E)) {
-			extensions_added.push_back(E);
-		}
-	}
-
-	Vector<String> loaded_extensions = GDExtensionManager::get_singleton()->get_loaded_extensions();
-	for (int i = 0; i < loaded_extensions.size(); i++) {
-		if (!extensions.has(loaded_extensions[i])) {
-			// The extension may not have a .gdextension file.
-			if (!FileAccess::exists(loaded_extensions[i])) {
-				extensions_removed.push_back(loaded_extensions[i]);
-			}
-		}
-	}
-
-	String extension_list_config_file = GDExtension::get_extension_list_config_file();
-	if (extensions.size()) {
-		if (extensions_added.size() || extensions_removed.size()) { //extensions were added or removed
-			Ref<FileAccess> f = FileAccess::open(extension_list_config_file, FileAccess::WRITE);
-			for (const String &E : extensions) {
-				f->store_line(E);
-			}
-		}
-	} else {
-		if (loaded_extensions.size() || FileAccess::exists(extension_list_config_file)) { //extensions were removed
-			Ref<DirAccess> da = DirAccess::create(DirAccess::ACCESS_RESOURCES);
-			da->remove(extension_list_config_file);
-		}
-	}
-
-	bool needs_restart = false;
-	for (int i = 0; i < extensions_added.size(); i++) {
-		GDExtensionManager::LoadStatus st = GDExtensionManager::get_singleton()->load_extension(extensions_added[i]);
-		if (st == GDExtensionManager::LOAD_STATUS_NEEDS_RESTART) {
-			needs_restart = true;
-		}
-	}
-	for (int i = 0; i < extensions_removed.size(); i++) {
-		GDExtensionManager::LoadStatus st = GDExtensionManager::get_singleton()->unload_extension(extensions_removed[i]);
-		if (st == GDExtensionManager::LOAD_STATUS_NEEDS_RESTART) {
-			needs_restart = true;
-		}
-	}
-
-	return needs_restart;
+	return GDExtensionManager::get_singleton()->ensure_extensions_loaded(extensions);
 }
 
 void EditorFileSystem::_bind_methods() {
diff --git a/editor/editor_file_system.h b/editor/editor_file_system.h
index 1bc24416eb..ca4a64bfac 100644
--- a/editor/editor_file_system.h
+++ b/editor/editor_file_system.h
@@ -187,7 +187,7 @@ class EditorFileSystem : public Node {
 
 	void _scan_filesystem();
 	void _first_scan_filesystem();
-	void _first_scan_process_scripts(const ScannedDirectory *p_scan_dir, HashSet<String> &p_existing_class_names);
+	void _first_scan_process_scripts(const ScannedDirectory *p_scan_dir, HashSet<String> &p_existing_class_names, HashSet<String> &p_extensions);
 
 	HashSet<String> late_update_files;
 
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index c6144a34cb..199f3f55f8 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -481,6 +481,9 @@ void EditorNode::_gdextensions_reloaded() {
 	// In case the developer is inspecting an object that will be changed by the reload.
 	InspectorDock::get_inspector_singleton()->update_tree();
 
+	// Reload script editor to revalidate GDScript if classes are added or removed.
+	ScriptEditor::get_singleton()->reload_scripts(true);
+
 	// Regenerate documentation.
 	EditorHelp::generate_doc();
 }
diff --git a/modules/gdscript/gdscript.cpp b/modules/gdscript/gdscript.cpp
index c72755642b..36c5ad0937 100644
--- a/modules/gdscript/gdscript.cpp
+++ b/modules/gdscript/gdscript.cpp
@@ -57,6 +57,7 @@
 #include "scene/scene_string_names.h"
 
 #ifdef TOOLS_ENABLED
+#include "core/extension/gdextension_manager.h"
 #include "editor/editor_paths.h"
 #endif
 
@@ -2177,9 +2178,26 @@ void GDScriptLanguage::_add_global(const StringName &p_name, const Variant &p_va
 		global_array.write[globals[p_name]] = p_value;
 		return;
 	}
-	globals[p_name] = global_array.size();
-	global_array.push_back(p_value);
-	_global_array = global_array.ptrw();
+
+	if (global_array_empty_indexes.size()) {
+		int index = global_array_empty_indexes[global_array_empty_indexes.size() - 1];
+		globals[p_name] = index;
+		global_array.write[index] = p_value;
+		global_array_empty_indexes.resize(global_array_empty_indexes.size() - 1);
+	} else {
+		globals[p_name] = global_array.size();
+		global_array.push_back(p_value);
+		_global_array = global_array.ptrw();
+	}
+}
+
+void GDScriptLanguage::_remove_global(const StringName &p_name) {
+	if (!globals.has(p_name)) {
+		return;
+	}
+	global_array_empty_indexes.push_back(globals[p_name]);
+	global_array.write[globals[p_name]] = Variant::NIL;
+	globals.erase(p_name);
 }
 
 void GDScriptLanguage::add_global_constant(const StringName &p_variable, const Variant &p_value) {
@@ -2237,11 +2255,40 @@ void GDScriptLanguage::init() {
 		_add_global(E.name, E.ptr);
 	}
 
+#ifdef TOOLS_ENABLED
+	if (Engine::get_singleton()->is_editor_hint()) {
+		GDExtensionManager::get_singleton()->connect("extension_loaded", callable_mp(this, &GDScriptLanguage::_extension_loaded));
+		GDExtensionManager::get_singleton()->connect("extension_unloading", callable_mp(this, &GDScriptLanguage::_extension_unloading));
+	}
+#endif
+
 #ifdef TESTS_ENABLED
 	GDScriptTests::GDScriptTestRunner::handle_cmdline();
 #endif
 }
 
+#ifdef TOOLS_ENABLED
+void GDScriptLanguage::_extension_loaded(const Ref<GDExtension> &p_extension) {
+	List<StringName> class_list;
+	ClassDB::get_extension_class_list(p_extension, &class_list);
+	for (const StringName &n : class_list) {
+		if (globals.has(n)) {
+			continue;
+		}
+		Ref<GDScriptNativeClass> nc = memnew(GDScriptNativeClass(n));
+		_add_global(n, nc);
+	}
+}
+
+void GDScriptLanguage::_extension_unloading(const Ref<GDExtension> &p_extension) {
+	List<StringName> class_list;
+	ClassDB::get_extension_class_list(p_extension, &class_list);
+	for (const StringName &n : class_list) {
+		_remove_global(n);
+	}
+}
+#endif
+
 String GDScriptLanguage::get_type() const {
 	return "GDScript";
 }
diff --git a/modules/gdscript/gdscript.h b/modules/gdscript/gdscript.h
index 6527a0ea4d..4d21651365 100644
--- a/modules/gdscript/gdscript.h
+++ b/modules/gdscript/gdscript.h
@@ -417,6 +417,7 @@ class GDScriptLanguage : public ScriptLanguage {
 	Vector<Variant> global_array;
 	HashMap<StringName, int> globals;
 	HashMap<StringName, Variant> named_globals;
+	Vector<int> global_array_empty_indexes;
 
 	struct CallLevel {
 		Variant *stack = nullptr;
@@ -448,6 +449,7 @@ class GDScriptLanguage : public ScriptLanguage {
 	int _debug_max_call_stack = 0;
 
 	void _add_global(const StringName &p_name, const Variant &p_value);
+	void _remove_global(const StringName &p_name);
 
 	friend class GDScriptInstance;
 
@@ -467,6 +469,11 @@ class GDScriptLanguage : public ScriptLanguage {
 
 	HashMap<String, ObjectID> orphan_subclasses;
 
+#ifdef TOOLS_ENABLED
+	void _extension_loaded(const Ref<GDExtension> &p_extension);
+	void _extension_unloading(const Ref<GDExtension> &p_extension);
+#endif
+
 public:
 	int calls;