From 06998a3927e1c36212ba98615c2aefe92fe5b5e0 Mon Sep 17 00:00:00 2001
From: bjornmp <bjorn.mp@gmail.com>
Date: Mon, 25 Sep 2023 00:32:43 +0200
Subject: [PATCH] Enforce that custom nodes keep their original type

Enforce that custom nodes and resources created via the "Create New Node" dialog, should permanently retain their original type (script). This means:

- Type continuity: It should be impossible for the user to (accidentally) clear the original script of a custom node that was created via the "Create New Node" dialog.

- Extensibility: The user should be able to extend custom types as usual (create a script that inherits the original type and replace the original script of that node with his own). However, if he then clears his extension-script from that node later on, the custom type should revert to its original script instead of becoming a non-scripted type.
---
 editor/editor_data.cpp            |   2 +
 editor/editor_node.cpp            |   5 ++
 editor/editor_properties.cpp      |   1 +
 editor/editor_resource_picker.cpp |  23 +++++-
 editor/editor_resource_picker.h   |   5 ++
 editor/scene_tree_dock.cpp        | 121 ++++++++++++++++++++++++------
 editor/scene_tree_dock.h          |   1 +
 scene/property_utils.cpp          |  10 +++
 scene/scene_string_names.cpp      |   2 +
 scene/scene_string_names.h        |   2 +
 10 files changed, 147 insertions(+), 25 deletions(-)

diff --git a/editor/editor_data.cpp b/editor/editor_data.cpp
index ee16c61c897..bb02172b1a8 100644
--- a/editor/editor_data.cpp
+++ b/editor/editor_data.cpp
@@ -547,6 +547,7 @@ Variant EditorData::instantiate_custom_type(const String &p_type, const String &
 				if (n) {
 					n->set_name(p_type);
 				}
+				n->set_meta(SceneStringName(_custom_type_script), script);
 				((Object *)ob)->set_script(script);
 				return ob;
 			}
@@ -1008,6 +1009,7 @@ Variant EditorData::script_class_instance(const String &p_class) {
 			// Store in a variant to initialize the refcount if needed.
 			Variant obj = ClassDB::instantiate(script->get_instance_base_type());
 			if (obj) {
+				Object::cast_to<Object>(obj)->set_meta(SceneStringName(_custom_type_script), script);
 				obj.operator Object *()->set_script(script);
 			}
 			return obj;
diff --git a/editor/editor_node.cpp b/editor/editor_node.cpp
index d88fb134f16..e57328911dd 100644
--- a/editor/editor_node.cpp
+++ b/editor/editor_node.cpp
@@ -4672,6 +4672,11 @@ void EditorNode::stop_child_process(OS::ProcessID p_pid) {
 Ref<Script> EditorNode::get_object_custom_type_base(const Object *p_object) const {
 	ERR_FAIL_NULL_V(p_object, nullptr);
 
+	const Node *node = Object::cast_to<const Node>(p_object);
+	if (node && node->has_meta(SceneStringName(_custom_type_script))) {
+		return node->get_meta(SceneStringName(_custom_type_script));
+	}
+
 	Ref<Script> scr = p_object->get_script();
 
 	if (scr.is_valid()) {
diff --git a/editor/editor_properties.cpp b/editor/editor_properties.cpp
index c5a35e466c9..2e2a18b9d6c 100644
--- a/editor/editor_properties.cpp
+++ b/editor/editor_properties.cpp
@@ -3221,6 +3221,7 @@ void EditorPropertyResource::setup(Object *p_object, const String &p_path, const
 	}
 
 	resource_picker->set_base_type(p_base_type);
+	resource_picker->set_resource_owner(p_object);
 	resource_picker->set_editable(true);
 	resource_picker->set_h_size_flags(SIZE_EXPAND_FILL);
 	add_child(resource_picker);
diff --git a/editor/editor_resource_picker.cpp b/editor/editor_resource_picker.cpp
index 0f0287718c6..025d019f45e 100644
--- a/editor/editor_resource_picker.cpp
+++ b/editor/editor_resource_picker.cpp
@@ -224,7 +224,9 @@ void EditorResourcePicker::_update_menu_items() {
 		}
 
 		if (is_editable()) {
-			edit_menu->add_icon_item(get_editor_theme_icon(SNAME("Clear")), TTR("Clear"), OBJ_MENU_CLEAR);
+			if (!_is_custom_type_script()) {
+				edit_menu->add_icon_item(get_editor_theme_icon(SNAME("Clear")), TTR("Clear"), OBJ_MENU_CLEAR);
+			}
 			edit_menu->add_icon_item(get_editor_theme_icon(SNAME("Duplicate")), TTR("Make Unique"), OBJ_MENU_MAKE_UNIQUE);
 
 			// Check whether the resource has subresources.
@@ -694,6 +696,16 @@ bool EditorResourcePicker::_is_type_valid(const String &p_type_name, const HashS
 	return false;
 }
 
+bool EditorResourcePicker::_is_custom_type_script() const {
+	Ref<Script> resource_as_script = edited_resource;
+
+	if (resource_as_script.is_valid() && resource_owner && resource_owner->has_meta(SceneStringName(_custom_type_script))) {
+		return true;
+	}
+
+	return false;
+}
+
 Variant EditorResourcePicker::get_drag_data_fw(const Point2 &p_point, Control *p_from) {
 	if (edited_resource.is_valid()) {
 		Dictionary drag_data = EditorNode::get_singleton()->drag_resource(edited_resource, p_from);
@@ -953,6 +965,10 @@ bool EditorResourcePicker::is_toggle_pressed() const {
 	return assign_button->is_pressed();
 }
 
+void EditorResourcePicker::set_resource_owner(Object *p_object) {
+	resource_owner = p_object;
+}
+
 void EditorResourcePicker::set_editable(bool p_editable) {
 	editable = p_editable;
 	assign_button->set_disabled(!editable && !edited_resource.is_valid());
@@ -1098,7 +1114,10 @@ void EditorScriptPicker::set_create_options(Object *p_menu_node) {
 		return;
 	}
 
-	menu_node->add_icon_item(get_editor_theme_icon(SNAME("ScriptCreate")), TTR("New Script..."), OBJ_MENU_NEW_SCRIPT);
+	if (!(script_owner && script_owner->has_meta(SceneStringName(_custom_type_script)))) {
+		menu_node->add_icon_item(get_editor_theme_icon(SNAME("ScriptCreate")), TTR("New Script..."), OBJ_MENU_NEW_SCRIPT);
+	}
+
 	if (script_owner) {
 		Ref<Script> scr = script_owner->get_script();
 		if (scr.is_valid()) {
diff --git a/editor/editor_resource_picker.h b/editor/editor_resource_picker.h
index 0a32dea3ed2..8fb774a2cb9 100644
--- a/editor/editor_resource_picker.h
+++ b/editor/editor_resource_picker.h
@@ -81,6 +81,8 @@ class EditorResourcePicker : public HBoxContainer {
 		CONVERT_BASE_ID = 1000,
 	};
 
+	Object *resource_owner = nullptr;
+
 	PopupMenu *edit_menu = nullptr;
 
 	void _update_resource_preview(const String &p_path, const Ref<Texture2D> &p_preview, const Ref<Texture2D> &p_small_preview, ObjectID p_obj);
@@ -102,6 +104,7 @@ class EditorResourcePicker : public HBoxContainer {
 	void _ensure_allowed_types() const;
 	bool _is_drop_valid(const Dictionary &p_drag_data) const;
 	bool _is_type_valid(const String &p_type_name, const HashSet<StringName> &p_allowed_types) const;
+	bool _is_custom_type_script() const;
 
 	Variant get_drag_data_fw(const Point2 &p_point, Control *p_from);
 	bool can_drop_data_fw(const Point2 &p_point, const Variant &p_data, Control *p_from) const;
@@ -137,6 +140,8 @@ public:
 	void set_toggle_pressed(bool p_pressed);
 	bool is_toggle_pressed() const;
 
+	void set_resource_owner(Object *p_object);
+
 	void set_editable(bool p_editable);
 	bool is_editable() const;
 
diff --git a/editor/scene_tree_dock.cpp b/editor/scene_tree_dock.cpp
index bcab0c2883c..11aa2828338 100644
--- a/editor/scene_tree_dock.cpp
+++ b/editor/scene_tree_dock.cpp
@@ -1667,6 +1667,7 @@ void SceneTreeDock::_notification(int p_what) {
 			button_instance->set_icon(get_editor_theme_icon(SNAME("Instance")));
 			button_create_script->set_icon(get_editor_theme_icon(SNAME("ScriptCreate")));
 			button_detach_script->set_icon(get_editor_theme_icon(SNAME("ScriptRemove")));
+			button_extend_script->set_icon(get_editor_theme_icon(SNAME("ScriptExtend")));
 			button_tree_menu->set_icon(get_editor_theme_icon(SNAME("GuiTabMenuHl")));
 
 			filter->set_right_icon(get_editor_theme_icon(SNAME("Search")));
@@ -2784,33 +2785,49 @@ void SceneTreeDock::_delete_confirm(bool p_cut) {
 }
 
 void SceneTreeDock::_update_script_button() {
-	if (!profile_allow_script_editing) {
-		button_create_script->hide();
-		button_detach_script->hide();
-	} else if (editor_selection->get_selection().size() == 0) {
-		button_create_script->hide();
-		button_detach_script->hide();
-	} else if (editor_selection->get_selection().size() == 1) {
-		Node *n = editor_selection->get_selected_node_list().front()->get();
-		if (n->get_script().is_null()) {
-			button_create_script->show();
-			button_detach_script->hide();
-		} else {
-			button_create_script->hide();
-			button_detach_script->show();
-		}
-	} else {
-		button_create_script->hide();
+	bool can_create_script = false;
+	bool can_detach_script = false;
+	bool can_extend_script = false;
+
+	if (profile_allow_script_editing) {
 		Array selection = editor_selection->get_selected_nodes();
+
 		for (int i = 0; i < selection.size(); i++) {
 			Node *n = Object::cast_to<Node>(selection[i]);
-			if (!n->get_script().is_null()) {
-				button_detach_script->show();
-				return;
+			Ref<Script> s = n->get_script();
+			Ref<Script> cts;
+
+			if (n->has_meta(SceneStringName(_custom_type_script))) {
+				cts = n->get_meta(SceneStringName(_custom_type_script));
+			}
+
+			if (selection.size() == 1) {
+				if (s.is_valid()) {
+					if (cts.is_valid() && s == cts) {
+						can_extend_script = true;
+					}
+				} else {
+					can_create_script = true;
+				}
+			}
+
+			if (s.is_valid()) {
+				if (cts.is_valid()) {
+					if (s != cts) {
+						can_detach_script = true;
+						break;
+					}
+				} else {
+					can_detach_script = true;
+					break;
+				}
 			}
 		}
-		button_detach_script->hide();
 	}
+
+	button_create_script->set_visible(can_create_script);
+	button_detach_script->set_visible(can_detach_script);
+	button_extend_script->set_visible(can_extend_script);
 }
 
 void SceneTreeDock::_selection_changed() {
@@ -3057,7 +3074,28 @@ void SceneTreeDock::_replace_node(Node *p_node, Node *p_by_node, bool p_keep_pro
 	Node *newnode = p_by_node;
 
 	if (p_keep_properties) {
-		Node *default_oldnode = Object::cast_to<Node>(ClassDB::instantiate(oldnode->get_class()));
+		Node *default_oldnode = nullptr;
+
+		// If we're dealing with a custom node type, we need to create a default instance of the custom type instead of the native type for property comparison.
+		if (oldnode->has_meta(SceneStringName(_custom_type_script))) {
+			Ref<Script> cts = oldnode->get_meta(SceneStringName(_custom_type_script));
+			default_oldnode = Object::cast_to<Node>(get_editor_data()->script_class_instance(cts->get_global_name()));
+			if (default_oldnode) {
+				default_oldnode->set_name(cts->get_global_name());
+				get_editor_data()->instantiate_object_properties(default_oldnode);
+			} else {
+				// Legacy custom type, registered with "add_custom_type()".
+				// TODO: Should probably be deprecated in 4.x.
+				const EditorData::CustomType *custom_type = get_editor_data()->get_custom_type_by_path(cts->get_path());
+				if (custom_type) {
+					default_oldnode = Object::cast_to<Node>(get_editor_data()->instantiate_custom_type(custom_type->name, cts->get_instance_base_type()));
+				}
+			}
+		}
+
+		if (!default_oldnode) {
+			default_oldnode = Object::cast_to<Node>(ClassDB::instantiate(oldnode->get_class()));
+		}
 
 		List<PropertyInfo> pinfo;
 		oldnode->get_property_list(&pinfo);
@@ -3542,6 +3580,27 @@ void SceneTreeDock::_script_dropped(const String &p_file, NodePath p_to) {
 		undo_redo->add_undo_method(ed, "live_debug_remove_node", NodePath(String(edited_scene->get_path_to(n)).path_join(new_node->get_name())));
 		undo_redo->commit_action();
 	} else {
+		// Check if dropped script is compatible.
+		if (n->has_meta(SceneStringName(_custom_type_script))) {
+			Ref<Script> ct_scr = n->get_meta(SceneStringName(_custom_type_script));
+			if (!scr->inherits_script(ct_scr)) {
+				String custom_type_name = ct_scr->get_global_name();
+
+				// Legacy custom type, registered with "add_custom_type()".
+				if (custom_type_name.is_empty()) {
+					const EditorData::CustomType *custom_type = get_editor_data()->get_custom_type_by_path(ct_scr->get_path());
+					if (custom_type) {
+						custom_type_name = custom_type->name;
+					} else {
+						custom_type_name = TTR("<unknown>");
+					}
+				}
+
+				WARN_PRINT_ED(vformat("Script does not extend type: '%s'.", custom_type_name));
+				return;
+			}
+		}
+
 		undo_redo->create_action(TTR("Attach Script"), UndoRedo::MERGE_DISABLE, n);
 		undo_redo->add_do_method(InspectorDock::get_singleton(), "store_script_properties", n);
 		undo_redo->add_undo_method(InspectorDock::get_singleton(), "store_script_properties", n);
@@ -3649,6 +3708,7 @@ void SceneTreeDock::_tree_rmb(const Vector2 &p_menu_pos) {
 
 	Ref<Script> existing_script;
 	bool existing_script_removable = true;
+	bool allow_attach_new_script = true;
 	if (selection.size() == 1) {
 		Node *selected = selection.front()->get();
 
@@ -3672,6 +3732,10 @@ void SceneTreeDock::_tree_rmb(const Vector2 &p_menu_pos) {
 		if (EditorNode::get_singleton()->get_object_custom_type_base(selected) == existing_script) {
 			existing_script_removable = false;
 		}
+
+		if (selected->has_meta(SceneStringName(_custom_type_script))) {
+			allow_attach_new_script = false;
+		}
 	}
 
 	if (profile_allow_editing) {
@@ -3692,7 +3756,10 @@ void SceneTreeDock::_tree_rmb(const Vector2 &p_menu_pos) {
 
 		if (full_selection.size() == 1) {
 			add_separator = true;
-			menu->add_icon_shortcut(get_editor_theme_icon(SNAME("ScriptCreate")), ED_GET_SHORTCUT("scene_tree/attach_script"), TOOL_ATTACH_SCRIPT);
+			if (allow_attach_new_script) {
+				menu->add_icon_shortcut(get_editor_theme_icon(SNAME("ScriptCreate")), ED_GET_SHORTCUT("scene_tree/attach_script"), TOOL_ATTACH_SCRIPT);
+			}
+
 			if (existing_script.is_valid()) {
 				menu->add_icon_shortcut(get_editor_theme_icon(SNAME("ScriptExtend")), ED_GET_SHORTCUT("scene_tree/extend_script"), TOOL_EXTEND_SCRIPT);
 			}
@@ -4601,6 +4668,14 @@ SceneTreeDock::SceneTreeDock(Node *p_scene_root, EditorSelection *p_editor_selec
 	filter_hbc->add_child(button_detach_script);
 	button_detach_script->hide();
 
+	button_extend_script = memnew(Button);
+	button_extend_script->set_flat(true);
+	button_extend_script->connect(SceneStringName(pressed), callable_mp(this, &SceneTreeDock::_tool_selected).bind(TOOL_EXTEND_SCRIPT, false));
+	button_extend_script->set_tooltip_text(TTR("Extend the script of the selected node."));
+	button_extend_script->set_shortcut(ED_GET_SHORTCUT("scene_tree/extend_script"));
+	filter_hbc->add_child(button_extend_script);
+	button_extend_script->hide();
+
 	button_tree_menu = memnew(MenuButton);
 	button_tree_menu->set_flat(false);
 	button_tree_menu->set_theme_type_variation("FlatMenuButton");
diff --git a/editor/scene_tree_dock.h b/editor/scene_tree_dock.h
index 05ad0f36e48..8cee2870f6e 100644
--- a/editor/scene_tree_dock.h
+++ b/editor/scene_tree_dock.h
@@ -115,6 +115,7 @@ class SceneTreeDock : public VBoxContainer {
 	Button *button_instance = nullptr;
 	Button *button_create_script = nullptr;
 	Button *button_detach_script = nullptr;
+	Button *button_extend_script = nullptr;
 	MenuButton *button_tree_menu = nullptr;
 
 	Button *node_shortcuts_toggle = nullptr;
diff --git a/scene/property_utils.cpp b/scene/property_utils.cpp
index 94a037bd9b8..f068e34bebc 100644
--- a/scene/property_utils.cpp
+++ b/scene/property_utils.cpp
@@ -89,6 +89,16 @@ Variant PropertyUtils::get_property_default_value(const Object *p_object, const
 		*r_is_valid = false;
 	}
 
+	// Handle special case "script" property, where the default value is either null or the custom type script.
+	// Do this only if there's no states stack cache to trace for default values.
+	if (!p_states_stack_cache && p_property == CoreStringName(script) && p_object->has_meta(SceneStringName(_custom_type_script))) {
+		Ref<Script> ct_scr = p_object->get_meta(SceneStringName(_custom_type_script));
+		if (r_is_valid) {
+			*r_is_valid = true;
+		}
+		return ct_scr;
+	}
+
 	Ref<Script> topmost_script;
 
 	if (const Node *node = Object::cast_to<Node>(p_object)) {
diff --git a/scene/scene_string_names.cpp b/scene/scene_string_names.cpp
index 140e5882910..31daeb3ae3c 100644
--- a/scene/scene_string_names.cpp
+++ b/scene/scene_string_names.cpp
@@ -130,6 +130,8 @@ SceneStringNames::SceneStringNames() {
 	shader_overrides_group = StaticCString::create("_shader_overrides_group_");
 	shader_overrides_group_active = StaticCString::create("_shader_overrides_group_active_");
 
+	_custom_type_script = StaticCString::create("_custom_type_script");
+
 	pressed = StaticCString::create("pressed");
 	id_pressed = StaticCString::create("id_pressed");
 	toggled = StaticCString::create("toggled");
diff --git a/scene/scene_string_names.h b/scene/scene_string_names.h
index fc22be33b23..0a2ebeda7af 100644
--- a/scene/scene_string_names.h
+++ b/scene/scene_string_names.h
@@ -143,6 +143,8 @@ public:
 	StringName shader_overrides_group;
 	StringName shader_overrides_group_active;
 
+	StringName _custom_type_script;
+
 	StringName pressed;
 	StringName id_pressed;
 	StringName toggled;