From ee2cc347c6fb9dbf6ba096961b335fb8b4319553 Mon Sep 17 00:00:00 2001
From: Yuri Sizov <yuris@humnom.net>
Date: Fri, 31 Mar 2023 21:17:59 +0200
Subject: [PATCH] Add support for icons in GDExtension classes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Rémi Verschelde <rverschelde@gmail.com>
---
 core/extension/gdextension.cpp         |  9 +++++++++
 core/extension/gdextension.h           |  2 ++
 core/extension/gdextension_manager.cpp | 23 +++++++++++++++++++++++
 core/extension/gdextension_manager.h   |  4 ++++
 editor/editor_data.cpp                 | 12 ++++++++++++
 editor/editor_data.h                   |  2 ++
 editor/editor_node.cpp                 | 26 +++++++++++++++++++-------
 7 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/core/extension/gdextension.cpp b/core/extension/gdextension.cpp
index 829e1d8e5bb..f158755a856 100644
--- a/core/extension/gdextension.cpp
+++ b/core/extension/gdextension.cpp
@@ -549,6 +549,15 @@ Ref<Resource> GDExtensionResourceLoader::load(const String &p_path, const String
 		return Ref<Resource>();
 	}
 
+	// Handle icons if any are specified.
+	if (config->has_section("icons")) {
+		List<String> keys;
+		config->get_section_keys("icons", &keys);
+		for (const String &key : keys) {
+			lib->class_icon_paths[key] = config->get_value("icons", key);
+		}
+	}
+
 	return lib;
 }
 
diff --git a/core/extension/gdextension.h b/core/extension/gdextension.h
index 39523a142f7..9d946ca7ba5 100644
--- a/core/extension/gdextension.h
+++ b/core/extension/gdextension.h
@@ -67,6 +67,8 @@ protected:
 	static void _bind_methods();
 
 public:
+	HashMap<String, String> class_icon_paths;
+
 	static String get_extension_list_config_file();
 	static String find_extension_library(const String &p_path, Ref<ConfigFile> p_config, std::function<bool(String)> p_has_feature, PackedStringArray *r_tags = nullptr);
 
diff --git a/core/extension/gdextension_manager.cpp b/core/extension/gdextension_manager.cpp
index 8701e6d77b2..63e809bc7c3 100644
--- a/core/extension/gdextension_manager.cpp
+++ b/core/extension/gdextension_manager.cpp
@@ -50,6 +50,11 @@ GDExtensionManager::LoadStatus GDExtensionManager::load_extension(const String &
 			extension->initialize_library(GDExtension::InitializationLevel(i));
 		}
 	}
+
+	for (const KeyValue<String, String> &kv : extension->class_icon_paths) {
+		gdextension_class_icon_paths[kv.key] = kv.value;
+	}
+
 	gdextension_map[p_path] = extension;
 	return LOAD_STATUS_OK;
 }
@@ -74,6 +79,11 @@ GDExtensionManager::LoadStatus GDExtensionManager::unload_extension(const String
 			extension->deinitialize_library(GDExtension::InitializationLevel(i));
 		}
 	}
+
+	for (const KeyValue<String, String> &kv : extension->class_icon_paths) {
+		gdextension_class_icon_paths.erase(kv.key);
+	}
+
 	gdextension_map.erase(p_path);
 	return LOAD_STATUS_OK;
 }
@@ -95,6 +105,19 @@ Ref<GDExtension> GDExtensionManager::get_extension(const String &p_path) {
 	return E->value;
 }
 
+bool GDExtensionManager::class_has_icon_path(const String &p_class) const {
+	// TODO: Check that the icon belongs to a registered class somehow.
+	return gdextension_class_icon_paths.has(p_class);
+}
+
+String GDExtensionManager::class_get_icon_path(const String &p_class) const {
+	// TODO: Check that the icon belongs to a registered class somehow.
+	if (gdextension_class_icon_paths.has(p_class)) {
+		return gdextension_class_icon_paths[p_class];
+	}
+	return "";
+}
+
 void GDExtensionManager::initialize_extensions(GDExtension::InitializationLevel p_level) {
 	ERR_FAIL_COND(int32_t(p_level) - 1 != level);
 	for (KeyValue<String, Ref<GDExtension>> &E : gdextension_map) {
diff --git a/core/extension/gdextension_manager.h b/core/extension/gdextension_manager.h
index 456942af0de..3643f043d80 100644
--- a/core/extension/gdextension_manager.h
+++ b/core/extension/gdextension_manager.h
@@ -38,6 +38,7 @@ class GDExtensionManager : public Object {
 
 	int32_t level = -1;
 	HashMap<String, Ref<GDExtension>> gdextension_map;
+	HashMap<String, String> gdextension_class_icon_paths;
 
 	static void _bind_methods();
 
@@ -59,6 +60,9 @@ public:
 	Vector<String> get_loaded_extensions() const;
 	Ref<GDExtension> get_extension(const String &p_path);
 
+	bool class_has_icon_path(const String &p_class) const;
+	String class_get_icon_path(const String &p_class) const;
+
 	void initialize_extensions(GDExtension::InitializationLevel p_level);
 	void deinitialize_extensions(GDExtension::InitializationLevel p_level);
 
diff --git a/editor/editor_data.cpp b/editor/editor_data.cpp
index 2d095d2dc77..c381c8c322a 100644
--- a/editor/editor_data.cpp
+++ b/editor/editor_data.cpp
@@ -31,6 +31,7 @@
 #include "editor_data.h"
 
 #include "core/config/project_settings.h"
+#include "core/extension/gdextension_manager.h"
 #include "core/io/file_access.h"
 #include "core/io/image_loader.h"
 #include "core/io/resource_loader.h"
@@ -1030,6 +1031,17 @@ void EditorData::script_class_load_icon_paths() {
 	}
 }
 
+Ref<Texture2D> EditorData::extension_class_get_icon(const String &p_class) const {
+	if (GDExtensionManager::get_singleton()->class_has_icon_path(p_class)) {
+		String icon_path = GDExtensionManager::get_singleton()->class_get_icon_path(p_class);
+		Ref<Texture2D> icon = _load_script_icon(icon_path);
+		if (icon.is_valid()) {
+			return icon;
+		}
+	}
+	return nullptr;
+}
+
 Ref<Texture2D> EditorData::_load_script_icon(const String &p_path) const {
 	if (!p_path.is_empty() && ResourceLoader::exists(p_path)) {
 		Ref<Texture2D> icon = ResourceLoader::load(p_path);
diff --git a/editor/editor_data.h b/editor/editor_data.h
index fcf43e72f09..370963074c5 100644
--- a/editor/editor_data.h
+++ b/editor/editor_data.h
@@ -243,6 +243,8 @@ public:
 	void script_class_save_icon_paths();
 	void script_class_load_icon_paths();
 
+	Ref<Texture2D> extension_class_get_icon(const String &p_class) const;
+
 	Ref<Texture2D> get_script_icon(const Ref<Script> &p_script);
 	void clear_script_icon_cache();
 
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index 74e6907215f..4fe43f38925 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -4477,21 +4477,33 @@ Ref<Texture2D> EditorNode::_get_class_or_script_icon(const String &p_class, cons
 			return script_icon;
 		}
 
-		// No custom icon was found in the inheritance chain, so check the built-in
-		// base class instead.
+		// No custom icon was found in the inheritance chain, so check the base
+		// class of the script instead.
 		String base_type;
 		p_script->get_language()->get_global_class_name(p_script->get_path(), &base_type);
-		if (gui_base) {
-			if (gui_base->has_theme_icon(base_type, "EditorIcons")) {
-				return gui_base->get_theme_icon(base_type, "EditorIcons");
-			}
-			return gui_base->get_theme_icon(p_fallback, "EditorIcons");
+
+		// Check if the base type is an extension-defined type.
+		Ref<Texture2D> ext_icon = ed.extension_class_get_icon(base_type);
+		if (ext_icon.is_valid()) {
+			return ext_icon;
+		}
+
+		// Look for the base type in the editor theme.
+		// This is only relevant for built-in classes.
+		if (gui_base && gui_base->has_theme_icon(base_type, "EditorIcons")) {
+			return gui_base->get_theme_icon(base_type, "EditorIcons");
 		}
 	}
 
 	// Script was not valid or didn't yield any useful values, try the class name
 	// directly.
 
+	// Check if the class name is an extension-defined type.
+	Ref<Texture2D> ext_icon = ed.extension_class_get_icon(p_class);
+	if (ext_icon.is_valid()) {
+		return ext_icon;
+	}
+
 	// Check if the class name is a custom type.
 	// TODO: Should probably be deprecated in 4.x
 	const EditorData::CustomType *ctype = ed.get_custom_type_by_name(p_class);