From 9fae65404a223a86816685b0b4036a57b8f976b7 Mon Sep 17 00:00:00 2001
From: Yuri Sizov <yuris@humnom.net>
Date: Fri, 31 Mar 2023 21:17:59 +0200
Subject: [PATCH] Streamline class icon resolution in the editor

---
 editor/editor_data.cpp       |  55 ++++++++++++-
 editor/editor_data.h         |   6 ++
 editor/editor_help.cpp       |  18 ++--
 editor/editor_help.h         |   2 +-
 editor/editor_inspector.cpp  |  42 +++-------
 editor/editor_node.cpp       | 155 +++++++++++++----------------------
 editor/editor_node.h         |   6 +-
 editor/editor_quick_open.cpp |   5 +-
 8 files changed, 137 insertions(+), 152 deletions(-)

diff --git a/editor/editor_data.cpp b/editor/editor_data.cpp
index 3059ce445c2..d2af7879d27 100644
--- a/editor/editor_data.cpp
+++ b/editor/editor_data.cpp
@@ -32,9 +32,11 @@
 
 #include "core/config/project_settings.h"
 #include "core/io/file_access.h"
+#include "core/io/image_loader.h"
 #include "core/io/resource_loader.h"
 #include "editor/editor_node.h"
 #include "editor/editor_plugin.h"
+#include "editor/editor_scale.h"
 #include "editor/editor_undo_redo_manager.h"
 #include "editor/plugins/script_editor_plugin.h"
 #include "scene/resources/packed_scene.h"
@@ -457,10 +459,10 @@ void EditorData::add_custom_type(const String &p_type, const String &p_inherits,
 	ct.name = p_type;
 	ct.icon = p_icon;
 	ct.script = p_script;
+
 	if (!custom_types.has(p_inherits)) {
 		custom_types[p_inherits] = Vector<CustomType>();
 	}
-
 	custom_types[p_inherits].push_back(ct);
 }
 
@@ -1028,8 +1030,57 @@ void EditorData::script_class_load_icon_paths() {
 	}
 }
 
+Ref<ImageTexture> EditorData::_load_script_icon(const String &p_path) const {
+	if (p_path.length()) {
+		Ref<Image> img = memnew(Image);
+		Error err = ImageLoader::load_image(p_path, img);
+		if (err == OK) {
+			img->resize(16 * EDSCALE, 16 * EDSCALE, Image::INTERPOLATE_LANCZOS);
+			return ImageTexture::create_from_image(img);
+		}
+	}
+	return nullptr;
+}
+
+Ref<Texture2D> EditorData::get_script_icon(const Ref<Script> &p_script) {
+	// Take from the local cache, if available.
+	if (_script_icon_cache.has(p_script) && _script_icon_cache[p_script].is_valid()) {
+		return _script_icon_cache[p_script];
+	}
+
+	Ref<Script> base_scr = p_script;
+	while (base_scr.is_valid()) {
+		// Check for scripted classes.
+		StringName name = script_class_get_name(base_scr->get_path());
+		String icon_path = script_class_get_icon_path(name);
+		Ref<ImageTexture> icon = _load_script_icon(icon_path);
+		if (icon.is_valid()) {
+			_script_icon_cache[p_script] = icon;
+			return icon;
+		}
+
+		// Check for legacy custom classes defined by plugins.
+		// TODO: Should probably be deprecated in 4.x
+		const EditorData::CustomType *ctype = get_custom_type_by_path(base_scr->get_path());
+		if (ctype && ctype->icon.is_valid()) {
+			_script_icon_cache[p_script] = ctype->icon;
+			return ctype->icon;
+		}
+
+		// Move to the base class.
+		base_scr = base_scr->get_base_script();
+	}
+
+	// If no icon found, cache it as null.
+	_script_icon_cache[p_script] = Ref<Texture>();
+	return nullptr;
+}
+
+void EditorData::clear_script_icon_cache() {
+	_script_icon_cache.clear();
+}
+
 EditorData::EditorData() {
-	current_edited_scene = -1;
 	undo_redo_manager = memnew(EditorUndoRedoManager);
 	script_class_load_icon_paths();
 }
diff --git a/editor/editor_data.h b/editor/editor_data.h
index d00501280d0..63eb0a5d6e7 100644
--- a/editor/editor_data.h
+++ b/editor/editor_data.h
@@ -144,6 +144,9 @@ private:
 
 	HashMap<StringName, String> _script_class_icon_paths;
 	HashMap<String, StringName> _script_class_file_to_path;
+	HashMap<Ref<Script>, Ref<Texture>> _script_icon_cache;
+
+	Ref<ImageTexture> _load_script_icon(const String &p_path) const;
 
 public:
 	EditorPlugin *get_editor(Object *p_object);
@@ -240,6 +243,9 @@ public:
 	void script_class_save_icon_paths();
 	void script_class_load_icon_paths();
 
+	Ref<Texture2D> get_script_icon(const Ref<Script> &p_script);
+	void clear_script_icon_cache();
+
 	EditorData();
 	~EditorData();
 };
diff --git a/editor/editor_help.cpp b/editor/editor_help.cpp
index cc569e777b0..af35bda47c4 100644
--- a/editor/editor_help.cpp
+++ b/editor/editor_help.cpp
@@ -263,16 +263,8 @@ void EditorHelp::_add_type(const String &p_type, const String &p_enum) {
 	class_desc->pop();
 }
 
-void EditorHelp::_add_type_icon(const String &p_type, int p_size) {
-	Ref<Texture2D> icon;
-	if (has_theme_icon(p_type, SNAME("EditorIcons"))) {
-		icon = get_theme_icon(p_type, SNAME("EditorIcons"));
-	} else if (ClassDB::class_exists(p_type) && ClassDB::is_parent_class(p_type, "Object")) {
-		icon = get_theme_icon(SNAME("Object"), SNAME("EditorIcons"));
-	} else {
-		icon = get_theme_icon(SNAME("ArrowRight"), SNAME("EditorIcons"));
-	}
-
+void EditorHelp::_add_type_icon(const String &p_type, int p_size, const String &p_fallback) {
+	Ref<Texture2D> icon = EditorNode::get_singleton()->get_class_icon(p_type, p_fallback);
 	Vector2i size = Vector2i(icon->get_width(), icon->get_height());
 	if (p_size > 0) {
 		// Ensures icon scales proportionally on both axis, based on icon height.
@@ -644,7 +636,7 @@ void EditorHelp::_update_doc() {
 	section_line.push_back(Pair<String, int>(TTR("Top"), 0));
 	_push_title_font();
 	class_desc->add_text(TTR("Class:") + " ");
-	_add_type_icon(edited_class, theme_cache.doc_title_font_size);
+	_add_type_icon(edited_class, theme_cache.doc_title_font_size, "Object");
 	class_desc->add_text(" ");
 	class_desc->push_color(theme_cache.headline_color);
 	_add_text(edited_class);
@@ -676,7 +668,7 @@ void EditorHelp::_update_doc() {
 		String inherits = cd.inherits;
 
 		while (!inherits.is_empty()) {
-			_add_type_icon(inherits);
+			_add_type_icon(inherits, 0, "ArrowRight");
 			class_desc->add_text(non_breaking_space); // Otherwise icon borrows hyperlink from _add_type().
 			_add_type(inherits);
 
@@ -709,7 +701,7 @@ void EditorHelp::_update_doc() {
 				if (prev) {
 					class_desc->add_text(" , ");
 				}
-				_add_type_icon(E.value.name);
+				_add_type_icon(E.value.name, 0, "ArrowRight");
 				class_desc->add_text(non_breaking_space); // Otherwise icon borrows hyperlink from _add_type().
 				_add_type(E.value.name);
 				prev = true;
diff --git a/editor/editor_help.h b/editor/editor_help.h
index b2ffe3bc292..a23dbca213c 100644
--- a/editor/editor_help.h
+++ b/editor/editor_help.h
@@ -159,7 +159,7 @@ class EditorHelp : public VBoxContainer {
 
 	//void _button_pressed(int p_idx);
 	void _add_type(const String &p_type, const String &p_enum = String());
-	void _add_type_icon(const String &p_type, int p_size = 0);
+	void _add_type_icon(const String &p_type, int p_size = 0, const String &p_fallback = "");
 	void _add_method(const DocData::MethodDoc &p_method, bool p_overview = true);
 
 	void _add_bulletpoint();
diff --git a/editor/editor_inspector.cpp b/editor/editor_inspector.cpp
index 44426dc1438..a8552091630 100644
--- a/editor/editor_inspector.cpp
+++ b/editor/editor_inspector.cpp
@@ -2753,46 +2753,28 @@ void EditorInspector::update_tree() {
 			String label = p.name;
 			doc_name = p.name;
 
-			// Set the category icon.
+			// Use category's owner script to update some of its information.
 			if (!EditorNode::get_editor_data().is_type_recognized(type) && p.hint_string.length() && FileAccess::exists(p.hint_string)) {
-				// If we have a category inside a script, search for the first script with a valid icon.
+				StringName script_name;
+
 				Ref<Script> scr = ResourceLoader::load(p.hint_string, "Script");
-				StringName base_type;
-				StringName name;
 				if (scr.is_valid()) {
-					base_type = scr->get_instance_base_type();
-					name = EditorNode::get_editor_data().script_class_get_name(scr->get_path());
+					script_name = EditorNode::get_editor_data().script_class_get_name(scr->get_path());
+
+					// Update the docs reference and the label based on the script.
 					Vector<DocData::ClassDoc> docs = scr->get_documentation();
 					if (!docs.is_empty()) {
 						doc_name = docs[0].name;
 					}
-					if (name != StringName() && label != name) {
-						label = name;
+					if (script_name != StringName() && label != script_name) {
+						label = script_name;
 					}
 				}
-				while (scr.is_valid()) {
-					name = EditorNode::get_editor_data().script_class_get_name(scr->get_path());
-					String icon_path = EditorNode::get_editor_data().script_class_get_icon_path(name);
-					if (name != StringName() && !icon_path.is_empty()) {
-						category->icon = ResourceLoader::load(icon_path, "Texture");
-						break;
-					}
 
-					const EditorData::CustomType *ctype = EditorNode::get_editor_data().get_custom_type_by_path(scr->get_path());
-					if (ctype) {
-						category->icon = ctype->icon;
-						break;
-					}
-					scr = scr->get_base_script();
-				}
-				if (category->icon.is_null() && has_theme_icon(base_type, SNAME("EditorIcons"))) {
-					category->icon = get_theme_icon(base_type, SNAME("EditorIcons"));
-				}
-			}
-			if (category->icon.is_null()) {
-				if (!type.is_empty()) { // Can happen for built-in scripts.
-					category->icon = EditorNode::get_singleton()->get_class_icon(type, "Object");
-				}
+				// Find the corresponding icon.
+				category->icon = EditorNode::get_singleton()->get_class_icon(script_name, "Object");
+			} else if (!type.is_empty()) {
+				category->icon = EditorNode::get_singleton()->get_class_icon(type, "Object");
 			}
 
 			// Set the category label.
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index 3e6faf99e70..dbe85eca648 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -34,7 +34,6 @@
 #include "core/input/input.h"
 #include "core/io/config_file.h"
 #include "core/io/file_access.h"
-#include "core/io/image_loader.h"
 #include "core/io/resource_loader.h"
 #include "core/io/resource_saver.h"
 #include "core/object/class_db.h"
@@ -3682,7 +3681,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.
-	script_icon_cache.clear();
+	editor_data.clear_script_icon_cache();
 
 	if (editor_data.get_edited_scene() == index) {
 		// Scene to remove is current scene.
@@ -4446,18 +4445,6 @@ StringName EditorNode::get_object_custom_type_name(const Object *p_object) const
 	return StringName();
 }
 
-Ref<ImageTexture> EditorNode::_load_custom_class_icon(const String &p_path) const {
-	if (p_path.length()) {
-		Ref<Image> img = memnew(Image);
-		Error err = ImageLoader::load_image(p_path, img);
-		if (err == OK) {
-			img->resize(16 * EDSCALE, 16 * EDSCALE, Image::INTERPOLATE_LANCZOS);
-			return ImageTexture::create_from_image(img);
-		}
-	}
-	return nullptr;
-}
-
 void EditorNode::_pick_main_scene_custom_action(const String &p_custom_action_name) {
 	if (p_custom_action_name == "select_current") {
 		Node *scene = editor_data.get_edited_scene_root();
@@ -4480,106 +4467,74 @@ void EditorNode::_pick_main_scene_custom_action(const String &p_custom_action_na
 	}
 }
 
+Ref<Texture2D> EditorNode::_get_class_or_script_icon(const String &p_class, const Ref<Script> &p_script, const String &p_fallback) {
+	ERR_FAIL_COND_V_MSG(p_class.is_empty(), nullptr, "Class name cannot be empty.");
+	EditorData &ed = EditorNode::get_editor_data();
+
+	// Check for a script icon first.
+	if (p_script.is_valid()) {
+		Ref<Texture2D> script_icon = ed.get_script_icon(p_script);
+		if (script_icon.is_valid()) {
+			return script_icon;
+		}
+
+		// No custom icon was found in the inheritance chain, so check the built-in
+		// base class 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");
+		}
+	}
+
+	// Script was not valid or didn't yield any useful values, try the class name
+	// directly.
+
+	// 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);
+	if (ctype && ctype->icon.is_valid()) {
+		return ctype->icon;
+	}
+
+	// Look up the class name or the fallback name in the editor theme.
+	// This is only relevant for built-in classes.
+	if (gui_base) {
+		if (gui_base->has_theme_icon(p_class, SNAME("EditorIcons"))) {
+			return gui_base->get_theme_icon(p_class, SNAME("EditorIcons"));
+		}
+
+		if (p_fallback.length() && gui_base->has_theme_icon(p_fallback, SNAME("EditorIcons"))) {
+			return gui_base->get_theme_icon(p_fallback, SNAME("EditorIcons"));
+		}
+	}
+
+	return nullptr;
+}
+
 Ref<Texture2D> EditorNode::get_object_icon(const Object *p_object, const String &p_fallback) {
-	ERR_FAIL_COND_V(!p_object || !gui_base, nullptr);
+	ERR_FAIL_NULL_V_MSG(p_object, nullptr, "Object cannot be null.");
 
 	Ref<Script> scr = p_object->get_script();
 	if (scr.is_null() && p_object->is_class("Script")) {
 		scr = p_object;
 	}
 
-	if (scr.is_valid() && !script_icon_cache.has(scr)) {
-		Ref<Script> base_scr = scr;
-		while (base_scr.is_valid()) {
-			StringName name = EditorNode::get_editor_data().script_class_get_name(base_scr->get_path());
-			String icon_path = EditorNode::get_editor_data().script_class_get_icon_path(name);
-			Ref<ImageTexture> icon = _load_custom_class_icon(icon_path);
-			if (icon.is_valid()) {
-				script_icon_cache[scr] = icon;
-				return icon;
-			}
-
-			// TODO: should probably be deprecated in 4.x
-			StringName base = base_scr->get_instance_base_type();
-			if (base != StringName() && EditorNode::get_editor_data().get_custom_types().has(base)) {
-				const Vector<EditorData::CustomType> &types = EditorNode::get_editor_data().get_custom_types()[base];
-				for (int i = 0; i < types.size(); ++i) {
-					if (types[i].script == base_scr && types[i].icon.is_valid()) {
-						script_icon_cache[scr] = types[i].icon;
-						return types[i].icon;
-					}
-				}
-			}
-			base_scr = base_scr->get_base_script();
-		}
-
-		// If no icon found, cache it as null.
-		script_icon_cache[scr] = Ref<Texture>();
-	} else if (scr.is_valid() && script_icon_cache.has(scr) && script_icon_cache[scr].is_valid()) {
-		return script_icon_cache[scr];
-	}
-
-	// TODO: Should probably be deprecated in 4.x.
-	if (p_object->has_meta("_editor_icon")) {
-		return p_object->get_meta("_editor_icon");
-	}
-
-	if (gui_base->has_theme_icon(p_object->get_class(), SNAME("EditorIcons"))) {
-		return gui_base->get_theme_icon(p_object->get_class(), SNAME("EditorIcons"));
-	}
-
-	if (p_fallback.length()) {
-		return gui_base->get_theme_icon(p_fallback, SNAME("EditorIcons"));
-	}
-
-	return nullptr;
+	return _get_class_or_script_icon(p_object->get_class(), scr, p_fallback);
 }
 
-Ref<Texture2D> EditorNode::get_class_icon(const String &p_class, const String &p_fallback) const {
+Ref<Texture2D> EditorNode::get_class_icon(const String &p_class, const String &p_fallback) {
 	ERR_FAIL_COND_V_MSG(p_class.is_empty(), nullptr, "Class name cannot be empty.");
 
+	Ref<Script> scr;
 	if (ScriptServer::is_global_class(p_class)) {
-		String class_name = p_class;
-		Ref<Script> scr = EditorNode::get_editor_data().script_class_load_script(class_name);
-
-		while (true) {
-			String icon_path = EditorNode::get_editor_data().script_class_get_icon_path(class_name);
-			Ref<Texture> icon = _load_custom_class_icon(icon_path);
-			if (icon.is_valid()) {
-				return icon; // Current global class has icon.
-			}
-
-			// Find next global class along the inheritance chain.
-			do {
-				Ref<Script> base_scr = scr->get_base_script();
-				if (base_scr.is_null()) {
-					// We've reached a native class, use its icon.
-					String base_type;
-					scr->get_language()->get_global_class_name(scr->get_path(), &base_type);
-					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");
-				}
-				scr = base_scr;
-				class_name = EditorNode::get_editor_data().script_class_get_name(scr->get_path());
-			} while (class_name.is_empty());
-		}
+		scr = EditorNode::get_editor_data().script_class_load_script(p_class);
 	}
 
-	if (const EditorData::CustomType *ctype = EditorNode::get_editor_data().get_custom_type_by_name(p_class)) {
-		return ctype->icon;
-	}
-
-	if (gui_base->has_theme_icon(p_class, SNAME("EditorIcons"))) {
-		return gui_base->get_theme_icon(p_class, SNAME("EditorIcons"));
-	}
-
-	if (p_fallback.length() && gui_base->has_theme_icon(p_fallback, SNAME("EditorIcons"))) {
-		return gui_base->get_theme_icon(p_fallback, SNAME("EditorIcons"));
-	}
-
-	return nullptr;
+	return _get_class_or_script_icon(p_class, scr, p_fallback);
 }
 
 bool EditorNode::is_object_of_custom_type(const Object *p_object, const StringName &p_class) {
diff --git a/editor/editor_node.h b/editor/editor_node.h
index 66a3bf5be2d..6bf2750fd4f 100644
--- a/editor/editor_node.h
+++ b/editor/editor_node.h
@@ -513,7 +513,6 @@ private:
 	PrintHandlerList print_handler;
 
 	HashMap<String, Ref<Texture2D>> icon_type_cache;
-	HashMap<Ref<Script>, Ref<Texture>> script_icon_cache;
 
 	static EditorBuildCallback build_callbacks[MAX_BUILD_CALLBACKS];
 	static EditorPluginInitializeCallback plugin_init_callbacks[MAX_INIT_CALLBACKS];
@@ -697,7 +696,8 @@ private:
 
 	void _feature_profile_changed();
 	bool _is_class_editor_disabled_by_feature_profile(const StringName &p_class);
-	Ref<ImageTexture> _load_custom_class_icon(const String &p_path) const;
+
+	Ref<Texture2D> _get_class_or_script_icon(const String &p_class, const Ref<Script> &p_script, const String &p_fallback = "Object");
 
 	void _pick_main_scene_custom_action(const String &p_custom_action_name);
 
@@ -879,7 +879,7 @@ public:
 	Ref<Script> get_object_custom_type_base(const Object *p_object) const;
 	StringName get_object_custom_type_name(const Object *p_object) const;
 	Ref<Texture2D> get_object_icon(const Object *p_object, const String &p_fallback = "Object");
-	Ref<Texture2D> get_class_icon(const String &p_class, const String &p_fallback = "Object") const;
+	Ref<Texture2D> get_class_icon(const String &p_class, const String &p_fallback = "Object");
 
 	bool is_object_of_custom_type(const Object *p_object, const StringName &p_class);
 
diff --git a/editor/editor_quick_open.cpp b/editor/editor_quick_open.cpp
index b90edb8f906..c7c41a6ed0b 100644
--- a/editor/editor_quick_open.cpp
+++ b/editor/editor_quick_open.cpp
@@ -69,10 +69,9 @@ void EditorQuickOpen::_build_search_cache(EditorFileSystemDirectory *p_efsd) {
 	for (int i = 0; i < p_efsd->get_file_count(); i++) {
 		String file = p_efsd->get_file_path(i);
 		String engine_type = p_efsd->get_file_type(i);
-
 		String script_type = p_efsd->get_file_resource_script_class(i);
-
 		String actual_type = script_type.is_empty() ? engine_type : script_type;
+
 		// Iterate all possible base types.
 		for (String &parent_type : base_types) {
 			if (ClassDB::is_parent_class(engine_type, parent_type) || EditorNode::get_editor_data().script_class_is_parent(script_type, parent_type)) {
@@ -81,7 +80,7 @@ void EditorQuickOpen::_build_search_cache(EditorFileSystemDirectory *p_efsd) {
 				// Store refs to used icons.
 				String ext = file.get_extension();
 				if (!icons.has(ext)) {
-					icons.insert(ext, get_theme_icon((has_theme_icon(actual_type, SNAME("EditorIcons")) ? actual_type : "Object"), SNAME("EditorIcons")));
+					icons.insert(ext, EditorNode::get_singleton()->get_class_icon(actual_type, "Object"));
 				}
 
 				// Stop testing base types as soon as we got a match.