From 2c77f07aaa4c04c6f439a4636156ba234f6d2e13 Mon Sep 17 00:00:00 2001
From: Yuri Sizov <yuris@humnom.net>
Date: Thu, 24 Aug 2023 12:49:20 +0200
Subject: [PATCH] Add a script method to get its class icon

Co-authored-by: Danil Alexeev <danil@alexeev.xyz>
---
 core/object/script_language.h              |  1 +
 core/object/script_language_extension.cpp  |  1 +
 core/object/script_language_extension.h    |  7 +++++++
 doc/classes/ScriptExtension.xml            |  5 +++++
 editor/editor_data.cpp                     |  8 +++++++-
 editor/editor_node.cpp                     |  2 +-
 modules/gdscript/doc_classes/@GDScript.xml |  2 +-
 modules/gdscript/gdscript.cpp              | 14 +++++++-------
 modules/gdscript/gdscript.h                |  2 ++
 modules/gdscript/gdscript_compiler.cpp     |  1 +
 modules/gdscript/gdscript_parser.cpp       | 17 +++++++++++++++--
 modules/gdscript/gdscript_parser.h         |  1 +
 modules/mono/csharp_script.h               |  3 +++
 13 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/core/object/script_language.h b/core/object/script_language.h
index 3ea6a6e4c3c0..215e5880fa4a 100644
--- a/core/object/script_language.h
+++ b/core/object/script_language.h
@@ -141,6 +141,7 @@ public:
 
 #ifdef TOOLS_ENABLED
 	virtual Vector<DocData::ClassDoc> get_documentation() const = 0;
+	virtual String get_class_icon_path() const = 0;
 	virtual PropertyInfo get_class_category() const;
 #endif // TOOLS_ENABLED
 
diff --git a/core/object/script_language_extension.cpp b/core/object/script_language_extension.cpp
index 0df9d5833440..bf8bac476a3e 100644
--- a/core/object/script_language_extension.cpp
+++ b/core/object/script_language_extension.cpp
@@ -52,6 +52,7 @@ void ScriptExtension::_bind_methods() {
 	GDVIRTUAL_BIND(_reload, "keep_state");
 
 	GDVIRTUAL_BIND(_get_documentation);
+	GDVIRTUAL_BIND(_get_class_icon_path);
 
 	GDVIRTUAL_BIND(_has_method, "method");
 	GDVIRTUAL_BIND(_get_method_info, "method");
diff --git a/core/object/script_language_extension.h b/core/object/script_language_extension.h
index 1a0ec294798f..7fa1fe681093 100644
--- a/core/object/script_language_extension.h
+++ b/core/object/script_language_extension.h
@@ -77,6 +77,7 @@ public:
 	EXBIND1R(Error, reload, bool)
 
 	GDVIRTUAL0RC(TypedArray<Dictionary>, _get_documentation)
+	GDVIRTUAL0RC(String, _get_class_icon_path)
 #ifdef TOOLS_ENABLED
 	virtual Vector<DocData::ClassDoc> get_documentation() const override {
 		TypedArray<Dictionary> doc;
@@ -89,6 +90,12 @@ public:
 
 		return class_doc;
 	}
+
+	virtual String get_class_icon_path() const override {
+		String ret;
+		GDVIRTUAL_CALL(_get_class_icon_path, ret);
+		return ret;
+	}
 #endif // TOOLS_ENABLED
 
 	EXBIND1RC(bool, has_method, const StringName &)
diff --git a/doc/classes/ScriptExtension.xml b/doc/classes/ScriptExtension.xml
index 934cbf5a269a..f7c8ecb3fbe3 100644
--- a/doc/classes/ScriptExtension.xml
+++ b/doc/classes/ScriptExtension.xml
@@ -22,6 +22,11 @@
 			<description>
 			</description>
 		</method>
+		<method name="_get_class_icon_path" qualifiers="virtual const">
+			<return type="String" />
+			<description>
+			</description>
+		</method>
 		<method name="_get_constants" qualifiers="virtual const">
 			<return type="Dictionary" />
 			<description>
diff --git a/editor/editor_data.cpp b/editor/editor_data.cpp
index 0b2c2bea154d..ae90ae6a420d 100644
--- a/editor/editor_data.cpp
+++ b/editor/editor_data.cpp
@@ -1122,8 +1122,14 @@ Ref<Texture2D> EditorData::get_script_icon(const Ref<Script> &p_script) {
 	Ref<Script> base_scr = p_script;
 	while (base_scr.is_valid()) {
 		// Check for scripted classes.
+		String icon_path;
 		StringName class_name = script_class_get_name(base_scr->get_path());
-		String icon_path = script_class_get_icon_path(class_name);
+		if (base_scr->is_built_in() || class_name == StringName()) {
+			icon_path = base_scr->get_class_icon_path();
+		} else {
+			icon_path = script_class_get_icon_path(class_name);
+		}
+
 		Ref<Texture2D> icon = _load_script_icon(icon_path);
 		if (icon.is_valid()) {
 			_script_icon_cache[p_script] = icon;
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index c6f5a6082bd9..9e29d9ddf3b1 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -3457,7 +3457,7 @@ void EditorNode::_remove_edited_scene(bool p_change_tab) {
 
 void EditorNode::_remove_scene(int index, bool p_change_tab) {
 	// Clear icon cache in case some scripts are no longer needed.
-	// FIXME: Perfectly the cache should never be cleared and only updated on per-script basis, when an icon changes.
+	// FIXME: Ideally the cache should never be cleared and only updated on per-script basis, when an icon changes.
 	editor_data.clear_script_icon_cache();
 
 	if (editor_data.get_edited_scene() == index) {
diff --git a/modules/gdscript/doc_classes/@GDScript.xml b/modules/gdscript/doc_classes/@GDScript.xml
index dd6b668c45f8..4f1a256ec994 100644
--- a/modules/gdscript/doc_classes/@GDScript.xml
+++ b/modules/gdscript/doc_classes/@GDScript.xml
@@ -595,7 +595,7 @@
 			<return type="void" />
 			<param index="0" name="icon_path" type="String" />
 			<description>
-				Add a custom icon to the current script. The script must be registered as a global class using the [code]class_name[/code] keyword for this to have a visible effect. The icon specified at [param icon_path] is displayed in the Scene dock for every node of that class, as well as in various editor dialogs.
+				Add a custom icon to the current script. The icon specified at [param icon_path] is displayed in the Scene dock for every node of that class, as well as in various editor dialogs.
 				[codeblock]
 				@icon("res://path/to/class/icon.svg")
 				[/codeblock]
diff --git a/modules/gdscript/gdscript.cpp b/modules/gdscript/gdscript.cpp
index 42b08f8a68df..fe7ccc31a5a3 100644
--- a/modules/gdscript/gdscript.cpp
+++ b/modules/gdscript/gdscript.cpp
@@ -482,6 +482,10 @@ void GDScript::_clear_doc() {
 	docs.clear();
 	doc = DocData::ClassDoc();
 }
+
+String GDScript::get_class_icon_path() const {
+	return simplified_icon_path;
+}
 #endif
 
 bool GDScript::_update_exports(bool *r_err, bool p_recursive_call, PlaceHolderScriptInstance *p_instance_to_update) {
@@ -2527,13 +2531,6 @@ String GDScriptLanguage::get_global_class_name(const String &p_path, String *r_b
 	 * Before changing this function, please ask the current maintainer of EditorFileSystem.
 	 */
 
-	if (r_icon_path) {
-		if (c->icon_path.is_empty() || c->icon_path.is_absolute_path()) {
-			*r_icon_path = c->icon_path.simplify_path();
-		} else if (c->icon_path.is_relative_path()) {
-			*r_icon_path = p_path.get_base_dir().path_join(c->icon_path).simplify_path();
-		}
-	}
 	if (r_base_type) {
 		const GDScriptParser::ClassNode *subclass = c;
 		String path = p_path;
@@ -2601,6 +2598,9 @@ String GDScriptLanguage::get_global_class_name(const String &p_path, String *r_b
 			}
 		}
 	}
+	if (r_icon_path) {
+		*r_icon_path = c->simplified_icon_path;
+	}
 	return c->identifier != nullptr ? String(c->identifier->name) : String();
 }
 
diff --git a/modules/gdscript/gdscript.h b/modules/gdscript/gdscript.h
index c41b1a0defd2..1fb0f01c1dc7 100644
--- a/modules/gdscript/gdscript.h
+++ b/modules/gdscript/gdscript.h
@@ -144,6 +144,7 @@ class GDScript : public Script {
 	String path;
 	String name;
 	String fully_qualified_name;
+	String simplified_icon_path;
 	SelfList<GDScript> script_list;
 
 	SelfList<GDScriptFunctionState>::List pending_func_states;
@@ -250,6 +251,7 @@ public:
 	virtual Vector<DocData::ClassDoc> get_documentation() const override {
 		return docs;
 	}
+	virtual String get_class_icon_path() const override;
 #endif // TOOLS_ENABLED
 
 	virtual Error reload(bool p_keep_state = false) override;
diff --git a/modules/gdscript/gdscript_compiler.cpp b/modules/gdscript/gdscript_compiler.cpp
index f964db231a6d..985eb97b29e2 100644
--- a/modules/gdscript/gdscript_compiler.cpp
+++ b/modules/gdscript/gdscript_compiler.cpp
@@ -2928,6 +2928,7 @@ void GDScriptCompiler::convert_to_initializer_type(Variant &p_variant, const GDS
 void GDScriptCompiler::make_scripts(GDScript *p_script, const GDScriptParser::ClassNode *p_class, bool p_keep_state) {
 	p_script->fully_qualified_name = p_class->fqcn;
 	p_script->name = p_class->identifier ? p_class->identifier->name : "";
+	p_script->simplified_icon_path = p_class->simplified_icon_path;
 
 	HashMap<StringName, Ref<GDScript>> old_subclasses;
 
diff --git a/modules/gdscript/gdscript_parser.cpp b/modules/gdscript/gdscript_parser.cpp
index 1dde67d2d127..f953e74932b8 100644
--- a/modules/gdscript/gdscript_parser.cpp
+++ b/modules/gdscript/gdscript_parser.cpp
@@ -3836,18 +3836,31 @@ bool GDScriptParser::tool_annotation(const AnnotationNode *p_annotation, Node *p
 bool GDScriptParser::icon_annotation(const AnnotationNode *p_annotation, Node *p_node) {
 	ERR_FAIL_COND_V_MSG(p_node->type != Node::CLASS, false, R"("@icon" annotation can only be applied to classes.)");
 	ERR_FAIL_COND_V(p_annotation->resolved_arguments.is_empty(), false);
+
 	ClassNode *p_class = static_cast<ClassNode *>(p_node);
+	String path = p_annotation->resolved_arguments[0];
+
 #ifdef DEBUG_ENABLED
 	if (!p_class->icon_path.is_empty()) {
 		push_error(R"("@icon" annotation can only be used once.)", p_annotation);
 		return false;
 	}
-	if (String(p_annotation->resolved_arguments[0]).is_empty()) {
+	if (path.is_empty()) {
 		push_error(R"("@icon" annotation argument must contain the path to the icon.)", p_annotation->arguments[0]);
 		return false;
 	}
 #endif // DEBUG_ENABLED
-	p_class->icon_path = p_annotation->resolved_arguments[0];
+
+	p_class->icon_path = path;
+
+	if (path.is_empty() || path.is_absolute_path()) {
+		p_class->simplified_icon_path = path.simplify_path();
+	} else if (path.is_relative_path()) {
+		p_class->simplified_icon_path = script_path.get_base_dir().path_join(path).simplify_path();
+	} else {
+		p_class->simplified_icon_path = path;
+	}
+
 	return true;
 }
 
diff --git a/modules/gdscript/gdscript_parser.h b/modules/gdscript/gdscript_parser.h
index 652faaebc344..b97e1fda33f8 100644
--- a/modules/gdscript/gdscript_parser.h
+++ b/modules/gdscript/gdscript_parser.h
@@ -724,6 +724,7 @@ public:
 
 		IdentifierNode *identifier = nullptr;
 		String icon_path;
+		String simplified_icon_path;
 		Vector<Member> members;
 		HashMap<StringName, int> members_indices;
 		ClassNode *outer = nullptr;
diff --git a/modules/mono/csharp_script.h b/modules/mono/csharp_script.h
index 9802067b4635..22115667a7bb 100644
--- a/modules/mono/csharp_script.h
+++ b/modules/mono/csharp_script.h
@@ -164,6 +164,9 @@ public:
 		Vector<DocData::ClassDoc> docs;
 		return docs;
 	}
+	virtual String get_class_icon_path() const override {
+		return icon_path;
+	}
 #endif // TOOLS_ENABLED
 
 	Error reload(bool p_keep_state = false) override;