Properly handle node remove via undo

4 files changed, 36 insertions, 0 deletions

diff --git a/editor/plugins/collision_shape_2d_editor_plugin.cpp b/editor/plugins/collision_shape_2d_editor_plugin.cpp
index 709ff91fc0..2d1d6de574 100644
--- a/editor/plugins/collision_shape_2d_editor_plugin.cpp
+++ b/editor/plugins/collision_shape_2d_editor_plugin.cpp
@@ -39,6 +39,13 @@
 #include "scene/resources/rectangle_shape_2d.h"
 #include "scene/resources/segment_shape_2d.h"
 
+void CollisionShape2DEditor::_node_removed(Node *p_node) {
+
+	if (p_node == node) {
+		node = NULL;
+	}
+}
+
 Variant CollisionShape2DEditor::get_handle_value(int idx) const {
 
 	switch (shape_type) {
@@ -525,6 +532,20 @@ void CollisionShape2DEditor::forward_canvas_draw_over_viewport(Control *p_overla
 	}
 }
 
+void CollisionShape2DEditor::_notification(int p_what) {
+
+	switch (p_what) {
+
+		case NOTIFICATION_ENTER_TREE: {
+			get_tree()->connect("node_removed", callable_mp(this, &CollisionShape2DEditor::_node_removed));
+		} break;
+
+		case NOTIFICATION_EXIT_TREE: {
+			get_tree()->disconnect("node_removed", callable_mp(this, &CollisionShape2DEditor::_node_removed));
+		} break;
+	}
+}
+
 void CollisionShape2DEditor::edit(Node *p_node) {
 
 	if (!canvas_item_editor) {
diff --git a/editor/plugins/collision_shape_2d_editor_plugin.h b/editor/plugins/collision_shape_2d_editor_plugin.h
index 025420a886..6bb621c476 100644
--- a/editor/plugins/collision_shape_2d_editor_plugin.h
+++ b/editor/plugins/collision_shape_2d_editor_plugin.h
@@ -71,6 +71,8 @@ class CollisionShape2DEditor : public Control {
 	void _get_current_shape_type();
 
 protected:
+	void _notification(int p_what);
+	void _node_removed(Node *p_node);
 	static void _bind_methods();
 
 public:
diff --git a/editor/plugins/tile_map_editor_plugin.cpp b/editor/plugins/tile_map_editor_plugin.cpp
index 83dc84a0c2..6fa7955ee5 100644
--- a/editor/plugins/tile_map_editor_plugin.cpp
+++ b/editor/plugins/tile_map_editor_plugin.cpp
@@ -38,6 +38,13 @@
 #include "editor/editor_settings.h"
 #include "scene/gui/split_container.h"
 
+void TileMapEditor::_node_removed(Node *p_node) {
+
+	if (p_node == node) {
+		node = NULL;
+	}
+}
+
 void TileMapEditor::_notification(int p_what) {
 
 	switch (p_what) {
@@ -60,6 +67,7 @@ void TileMapEditor::_notification(int p_what) {
 
 		case NOTIFICATION_ENTER_TREE: {
 
+			get_tree()->connect("node_removed", callable_mp(this, &TileMapEditor::_node_removed));
 			paint_button->set_icon(get_theme_icon("Edit", "EditorIcons"));
 			bucket_fill_button->set_icon(get_theme_icon("Bucket", "EditorIcons"));
 			picker_button->set_icon(get_theme_icon("ColorPick", "EditorIcons"));
@@ -80,6 +88,10 @@ void TileMapEditor::_notification(int p_what) {
 			p->set_item_icon(p->get_item_index(OPTION_ERASE_SELECTION), get_theme_icon("Remove", "EditorIcons"));
 
 		} break;
+
+		case NOTIFICATION_EXIT_TREE: {
+			get_tree()->disconnect("node_removed", callable_mp(this, &TileMapEditor::_node_removed));
+		} break;
 	}
 }
 
diff --git a/editor/plugins/tile_map_editor_plugin.h b/editor/plugins/tile_map_editor_plugin.h
index 74969d3e64..f43e5bb5cb 100644
--- a/editor/plugins/tile_map_editor_plugin.h
+++ b/editor/plugins/tile_map_editor_plugin.h
@@ -211,6 +211,7 @@ class TileMapEditor : public VBoxContainer {
 
 protected:
 	void _notification(int p_what);
+	void _node_removed(Node *p_node);
 	static void _bind_methods();
 	CellOp _get_op_from_cell(const Point2i &p_pos);