From 0b9b479ba7b08bba774f16202d882d9745dd74e8 Mon Sep 17 00:00:00 2001
From: Hendrik Brucker <hendrik.brucker@mail.de>
Date: Mon, 8 Apr 2024 15:31:27 +0200
Subject: [PATCH] [VisualShader] Remove invalid graph connections when ports
 are removed

Co-authored-by: EddieBreeg <eddiebreeg0@protonmail.com>
---
 .../plugins/visual_shader_editor_plugin.cpp   | 35 +++++++++++++++++--
 editor/plugins/visual_shader_editor_plugin.h  |  3 ++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/editor/plugins/visual_shader_editor_plugin.cpp b/editor/plugins/visual_shader_editor_plugin.cpp
index b6a4a1411735..67ed28f31c45 100644
--- a/editor/plugins/visual_shader_editor_plugin.cpp
+++ b/editor/plugins/visual_shader_editor_plugin.cpp
@@ -6747,8 +6747,35 @@ class VisualShaderNodePluginDefaultEditor : public VBoxContainer {
 		undo_redo->add_do_property(node.ptr(), p_property, p_value);
 		undo_redo->add_undo_property(node.ptr(), p_property, node->get(p_property));
 
+		Ref<VisualShaderNode> vsnode = editor->get_visual_shader()->get_node(shader_type, node_id);
+		ERR_FAIL_COND(vsnode.is_null());
+
+		// Check for invalid connections due to removed ports.
+		// We need to know the new state of the node to generate the proper undo/redo instructions.
+		// Quite hacky but the best way I could come up with for now.
+		Ref<VisualShaderNode> vsnode_new = vsnode->duplicate();
+		vsnode_new->set(p_property, p_value);
+		const int input_port_count = vsnode_new->get_input_port_count();
+		const int output_port_count = vsnode_new->get_output_port_count();
+
+		List<VisualShader::Connection> conns;
+		editor->get_visual_shader()->get_node_connections(shader_type, &conns);
+		VisualShaderGraphPlugin *graph_plugin = editor->get_graph_plugin();
+		bool undo_node_already_updated = false;
+		for (const VisualShader::Connection &c : conns) {
+			if ((c.from_node == node_id && c.from_port >= output_port_count) || (c.to_node == node_id && c.to_port >= input_port_count)) {
+				undo_redo->add_do_method(editor->get_visual_shader().ptr(), "disconnect_nodes", shader_type, c.from_node, c.from_port, c.to_node, c.to_port);
+				undo_redo->add_do_method(graph_plugin, "disconnect_nodes", shader_type, c.from_node, c.from_port, c.to_node, c.to_port);
+				// We need to update the node before reconnecting to avoid accessing a non-existing port.
+				undo_redo->add_undo_method(graph_plugin, "update_node_deferred", shader_type, node_id);
+				undo_node_already_updated = true;
+				undo_redo->add_undo_method(editor->get_visual_shader().ptr(), "connect_nodes", shader_type, c.from_node, c.from_port, c.to_node, c.to_port);
+				undo_redo->add_undo_method(graph_plugin, "connect_nodes", shader_type, c.from_node, c.from_port, c.to_node, c.to_port);
+			}
+		}
+
 		if (p_value.get_type() == Variant::OBJECT) {
-			Ref<Resource> prev_res = node->get(p_property);
+			Ref<Resource> prev_res = vsnode->get(p_property);
 			Ref<Resource> curr_res = p_value;
 
 			if (curr_res.is_null()) {
@@ -6763,14 +6790,16 @@ class VisualShaderNodePluginDefaultEditor : public VBoxContainer {
 			}
 		}
 		if (p_property != "constant") {
-			VisualShaderGraphPlugin *graph_plugin = editor->get_graph_plugin();
 			if (graph_plugin) {
 				undo_redo->add_do_method(editor, "_update_next_previews", node_id);
 				undo_redo->add_undo_method(editor, "_update_next_previews", node_id);
 				undo_redo->add_do_method(graph_plugin, "update_node_deferred", shader_type, node_id);
-				undo_redo->add_undo_method(graph_plugin, "update_node_deferred", shader_type, node_id);
+				if (!undo_node_already_updated) {
+					undo_redo->add_undo_method(graph_plugin, "update_node_deferred", shader_type, node_id);
+				}
 			}
 		}
+
 		undo_redo->commit_action();
 
 		updating = false;
diff --git a/editor/plugins/visual_shader_editor_plugin.h b/editor/plugins/visual_shader_editor_plugin.h
index 683a6bc88388..eea649c7bafb 100644
--- a/editor/plugins/visual_shader_editor_plugin.h
+++ b/editor/plugins/visual_shader_editor_plugin.h
@@ -555,7 +555,10 @@ class VisualShaderEditor : public VBoxContainer {
 	void update_custom_type(const Ref<Resource> &p_resource);
 
 	virtual Size2 get_minimum_size() const override;
+
 	void edit(VisualShader *p_visual_shader);
+	Ref<VisualShader> get_visual_shader() const { return visual_shader; }
+
 	VisualShaderEditor();
 };