From 8e6fdd5227c34df120a68f0cdeb2179f30f90918 Mon Sep 17 00:00:00 2001
From: Yuri Roubinsky <chaosus89@gmail.com>
Date: Sat, 18 Sep 2021 18:08:22 +0300
Subject: [PATCH] Draggin in/out from ports to create nodes in Animation Blend
 Tree editor

---
 .../animation_blend_tree_editor_plugin.cpp    | 75 +++++++++++++++----
 .../animation_blend_tree_editor_plugin.h      | 18 ++++-
 2 files changed, 73 insertions(+), 20 deletions(-)

diff --git a/editor/plugins/animation_blend_tree_editor_plugin.cpp b/editor/plugins/animation_blend_tree_editor_plugin.cpp
index 030d90eeca..24cb660f7a 100644
--- a/editor/plugins/animation_blend_tree_editor_plugin.cpp
+++ b/editor/plugins/animation_blend_tree_editor_plugin.cpp
@@ -66,9 +66,13 @@ void AnimationNodeBlendTreeEditor::remove_custom_type(const Ref<Script> &p_scrip
 	_update_options_menu();
 }
 
-void AnimationNodeBlendTreeEditor::_update_options_menu() {
+void AnimationNodeBlendTreeEditor::_update_options_menu(bool p_has_input_ports) {
 	add_node->get_popup()->clear();
+	add_node->get_popup()->set_size(Size2i(-1, -1));
 	for (int i = 0; i < add_options.size(); i++) {
+		if (p_has_input_ports && add_options[i].input_port_count == 0) {
+			continue;
+		}
 		add_node->get_popup()->add_item(add_options[i].name, i);
 	}
 
@@ -309,6 +313,11 @@ void AnimationNodeBlendTreeEditor::_add_node(int p_idx) {
 		return;
 	}
 
+	if (!from_node.is_empty() && anode->get_input_count() == 0) {
+		from_node = "";
+		return;
+	}
+
 	Point2 instance_pos = graph->get_scroll_ofs();
 	if (use_popup_menu_position) {
 		instance_pos += popup_menu_position;
@@ -328,11 +337,51 @@ void AnimationNodeBlendTreeEditor::_add_node(int p_idx) {
 	undo_redo->create_action(TTR("Add Node to BlendTree"));
 	undo_redo->add_do_method(blend_tree.ptr(), "add_node", name, anode, instance_pos / EDSCALE);
 	undo_redo->add_undo_method(blend_tree.ptr(), "remove_node", name);
+
+	if (!from_node.is_empty()) {
+		undo_redo->add_do_method(blend_tree.ptr(), "connect_node", name, 0, from_node);
+		from_node = "";
+	}
+	if (!to_node.is_empty() && to_slot != -1) {
+		undo_redo->add_do_method(blend_tree.ptr(), "connect_node", to_node, to_slot, name);
+		to_node = "";
+		to_slot = -1;
+	}
+
 	undo_redo->add_do_method(this, "_update_graph");
 	undo_redo->add_undo_method(this, "_update_graph");
 	undo_redo->commit_action();
 }
 
+void AnimationNodeBlendTreeEditor::_popup(bool p_has_input_ports, const Vector2 &p_popup_position, const Vector2 &p_node_position) {
+	_update_options_menu(p_has_input_ports);
+	use_popup_menu_position = true;
+	popup_menu_position = p_popup_position;
+	add_node->get_popup()->set_position(p_node_position);
+	add_node->get_popup()->popup();
+}
+
+void AnimationNodeBlendTreeEditor::_popup_request(const Vector2 &p_position) {
+	_popup(false, graph->get_local_mouse_position(), p_position);
+}
+
+void AnimationNodeBlendTreeEditor::_connection_to_empty(const String &p_from, int p_from_slot, const Vector2 &p_release_position) {
+	Ref<AnimationNode> node = blend_tree->get_node(p_from);
+	if (node.is_valid()) {
+		from_node = p_from;
+		_popup(true, p_release_position, graph->get_global_mouse_position());
+	}
+}
+
+void AnimationNodeBlendTreeEditor::_connection_from_empty(const String &p_to, int p_to_slot, const Vector2 &p_release_position) {
+	Ref<AnimationNode> node = blend_tree->get_node(p_to);
+	if (node.is_valid()) {
+		to_node = p_to;
+		to_slot = p_to_slot;
+		_popup(false, p_release_position, graph->get_global_mouse_position());
+	}
+}
+
 void AnimationNodeBlendTreeEditor::_node_dragged(const Vector2 &p_from, const Vector2 &p_to, const StringName &p_which) {
 	updating = true;
 	undo_redo->create_action(TTR("Node Moved"));
@@ -431,14 +480,6 @@ void AnimationNodeBlendTreeEditor::_delete_nodes_request() {
 	undo_redo->commit_action();
 }
 
-void AnimationNodeBlendTreeEditor::_popup_request(const Vector2 &p_position) {
-	_update_options_menu();
-	use_popup_menu_position = true;
-	popup_menu_position = graph->get_local_mouse_position();
-	add_node->get_popup()->set_position(p_position);
-	add_node->get_popup()->popup();
-}
-
 void AnimationNodeBlendTreeEditor::_node_selected(Object *p_node) {
 	GraphNode *gn = Object::cast_to<GraphNode>(p_node);
 	ERR_FAIL_COND(!gn);
@@ -890,6 +931,8 @@ AnimationNodeBlendTreeEditor::AnimationNodeBlendTreeEditor() {
 	graph->connect("scroll_offset_changed", callable_mp(this, &AnimationNodeBlendTreeEditor::_scroll_changed));
 	graph->connect("delete_nodes_request", callable_mp(this, &AnimationNodeBlendTreeEditor::_delete_nodes_request));
 	graph->connect("popup_request", callable_mp(this, &AnimationNodeBlendTreeEditor::_popup_request));
+	graph->connect("connection_to_empty", callable_mp(this, &AnimationNodeBlendTreeEditor::_connection_to_empty));
+	graph->connect("connection_from_empty", callable_mp(this, &AnimationNodeBlendTreeEditor::_connection_from_empty));
 	float graph_minimap_opacity = EditorSettings::get_singleton()->get("editors/visual_editors/minimap_opacity");
 	graph->set_minimap_opacity(graph_minimap_opacity);
 
@@ -905,13 +948,13 @@ AnimationNodeBlendTreeEditor::AnimationNodeBlendTreeEditor() {
 	add_node->connect("about_to_popup", callable_mp(this, &AnimationNodeBlendTreeEditor::_update_options_menu));
 
 	add_options.push_back(AddOption("Animation", "AnimationNodeAnimation"));
-	add_options.push_back(AddOption("OneShot", "AnimationNodeOneShot"));
-	add_options.push_back(AddOption("Add2", "AnimationNodeAdd2"));
-	add_options.push_back(AddOption("Add3", "AnimationNodeAdd3"));
-	add_options.push_back(AddOption("Blend2", "AnimationNodeBlend2"));
-	add_options.push_back(AddOption("Blend3", "AnimationNodeBlend3"));
-	add_options.push_back(AddOption("Seek", "AnimationNodeTimeSeek"));
-	add_options.push_back(AddOption("TimeScale", "AnimationNodeTimeScale"));
+	add_options.push_back(AddOption("OneShot", "AnimationNodeOneShot", 2));
+	add_options.push_back(AddOption("Add2", "AnimationNodeAdd2", 2));
+	add_options.push_back(AddOption("Add3", "AnimationNodeAdd3", 3));
+	add_options.push_back(AddOption("Blend2", "AnimationNodeBlend2", 2));
+	add_options.push_back(AddOption("Blend3", "AnimationNodeBlend3", 3));
+	add_options.push_back(AddOption("Seek", "AnimationNodeTimeSeek", 1));
+	add_options.push_back(AddOption("TimeScale", "AnimationNodeTimeScale", 1));
 	add_options.push_back(AddOption("Transition", "AnimationNodeTransition"));
 	add_options.push_back(AddOption("BlendTree", "AnimationNodeBlendTree"));
 	add_options.push_back(AddOption("BlendSpace1D", "AnimationNodeBlendSpace1D"));
diff --git a/editor/plugins/animation_blend_tree_editor_plugin.h b/editor/plugins/animation_blend_tree_editor_plugin.h
index 9f09069719..0fcafad40e 100644
--- a/editor/plugins/animation_blend_tree_editor_plugin.h
+++ b/editor/plugins/animation_blend_tree_editor_plugin.h
@@ -64,22 +64,28 @@ class AnimationNodeBlendTreeEditor : public AnimationTreeNodeEditorPlugin {
 	Map<StringName, ProgressBar *> animations;
 	Vector<EditorProperty *> visible_properties;
 
+	String to_node = "";
+	int to_slot = -1;
+	String from_node = "";
+
 	void _update_graph();
 
 	struct AddOption {
 		String name;
 		String type;
 		Ref<Script> script;
-		AddOption(const String &p_name = String(), const String &p_type = String()) :
+		int input_port_count;
+		AddOption(const String &p_name = String(), const String &p_type = String(), bool p_input_port_count = 0) :
 				name(p_name),
-				type(p_type) {
+				type(p_type),
+				input_port_count(p_input_port_count) {
 		}
 	};
 
 	Vector<AddOption> add_options;
 
 	void _add_node(int p_idx);
-	void _update_options_menu();
+	void _update_options_menu(bool p_has_input_ports = false);
 
 	static AnimationNodeBlendTreeEditor *singleton;
 
@@ -98,7 +104,6 @@ class AnimationNodeBlendTreeEditor : public AnimationTreeNodeEditorPlugin {
 	void _anim_selected(int p_index, Array p_options, const String &p_node);
 	void _delete_request(const String &p_which);
 	void _delete_nodes_request();
-	void _popup_request(const Vector2 &p_position);
 
 	bool _update_filters(const Ref<AnimationNode> &anode);
 	void _edit_filters(const String &p_which);
@@ -106,6 +111,11 @@ class AnimationNodeBlendTreeEditor : public AnimationTreeNodeEditorPlugin {
 	void _filter_toggled();
 	Ref<AnimationNode> _filter_edit;
 
+	void _popup(bool p_has_input_ports, const Vector2 &p_popup_position, const Vector2 &p_node_position);
+	void _popup_request(const Vector2 &p_position);
+	void _connection_to_empty(const String &p_from, int p_from_slot, const Vector2 &p_release_position);
+	void _connection_from_empty(const String &p_to, int p_to_slot, const Vector2 &p_release_position);
+
 	void _property_changed(const StringName &p_property, const Variant &p_value, const String &p_field, bool p_changing);
 	void _removed_from_graph();