From d5176757c0bfca35f6040b825f6a3201edbdf136 Mon Sep 17 00:00:00 2001
From: Jason Morgado <jasonsanimations@gmail.com>
Date: Sun, 20 Oct 2024 09:01:33 -0400
Subject: [PATCH] Add type filters to AnimationPlayer's "Add Track"

Co-Authored-By: A Thousand Ships <96648715+AThousandShips@users.noreply.github.com>
Co-Authored-By: Tomek <kobewi4e@gmail.com>
---
 editor/animation_track_editor.cpp | 22 ++++++++++++++++++
 editor/gui/scene_tree_editor.cpp  | 37 +++++++++++++++++++++----------
 editor/gui/scene_tree_editor.h    |  2 ++
 3 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/editor/animation_track_editor.cpp b/editor/animation_track_editor.cpp
index c210ad0761b..bc79b14d4a3 100644
--- a/editor/animation_track_editor.cpp
+++ b/editor/animation_track_editor.cpp
@@ -5297,6 +5297,28 @@ void AnimationTrackEditor::_add_track(int p_type) {
 		return;
 	}
 	adding_track_type = p_type;
+	Vector<StringName> valid_types;
+	switch (adding_track_type) {
+		case Animation::TYPE_BLEND_SHAPE: {
+			// Blend Shape is a property of MeshInstance3D.
+			valid_types.push_back(SNAME("MeshInstance3D"));
+		} break;
+		case Animation::TYPE_POSITION_3D:
+		case Animation::TYPE_ROTATION_3D:
+		case Animation::TYPE_SCALE_3D: {
+			// 3D Properties come from nodes inheriting Node3D.
+			valid_types.push_back(SNAME("Node3D"));
+		} break;
+		case Animation::TYPE_AUDIO: {
+			valid_types.push_back(SNAME("AudioStreamPlayer"));
+			valid_types.push_back(SNAME("AudioStreamPlayer2D"));
+			valid_types.push_back(SNAME("AudioStreamPlayer3D"));
+		} break;
+		case Animation::TYPE_ANIMATION: {
+			valid_types.push_back(SNAME("AnimationPlayer"));
+		} break;
+	}
+	pick_track->set_valid_types(valid_types);
 	pick_track->popup_scenetree_dialog(nullptr, root_node);
 	pick_track->get_filter_line_edit()->clear();
 	pick_track->get_filter_line_edit()->grab_focus();
diff --git a/editor/gui/scene_tree_editor.cpp b/editor/gui/scene_tree_editor.cpp
index 2e36b66025a..0ada9aa8b27 100644
--- a/editor/gui/scene_tree_editor.cpp
+++ b/editor/gui/scene_tree_editor.cpp
@@ -1684,24 +1684,30 @@ void SceneTreeDialog::_show_all_nodes_changed(bool p_button_pressed) {
 }
 
 void SceneTreeDialog::set_valid_types(const Vector<StringName> &p_valid) {
-	if (p_valid.is_empty()) {
-		return;
+	if (allowed_types_hbox) {
+		allowed_types_hbox->queue_free();
+		allowed_types_hbox = nullptr;
+		valid_type_icons.clear();
 	}
 
 	tree->set_valid_types(p_valid);
 
-	HBoxContainer *hbox = memnew(HBoxContainer);
-	content->add_child(hbox);
-	content->move_child(hbox, 0);
+	if (p_valid.is_empty()) {
+		return;
+	}
+
+	allowed_types_hbox = memnew(HBoxContainer);
+	content->add_child(allowed_types_hbox);
+	content->move_child(allowed_types_hbox, 0);
 
 	{
 		Label *label = memnew(Label);
-		hbox->add_child(label);
+		allowed_types_hbox->add_child(label);
 		label->set_text(TTR("Allowed:"));
 	}
 
 	HFlowContainer *hflow = memnew(HFlowContainer);
-	hbox->add_child(hflow);
+	allowed_types_hbox->add_child(hflow);
 	hflow->set_h_size_flags(Control::SIZE_EXPAND_FILL);
 
 	for (const StringName &type : p_valid) {
@@ -1735,6 +1741,9 @@ void SceneTreeDialog::set_valid_types(const Vector<StringName> &p_valid) {
 	}
 
 	show_all_nodes->show();
+	if (is_inside_tree()) {
+		_update_valid_type_icons();
+	}
 }
 
 void SceneTreeDialog::_notification(int p_what) {
@@ -1753,11 +1762,7 @@ void SceneTreeDialog::_notification(int p_what) {
 		} break;
 
 		case NOTIFICATION_THEME_CHANGED: {
-			filter->set_right_icon(get_editor_theme_icon(SNAME("Search")));
-			for (TextureRect *trect : valid_type_icons) {
-				trect->set_custom_minimum_size(Vector2(get_theme_constant(SNAME("class_icon_size"), EditorStringName(Editor)), 0));
-				trect->set_texture(trect->get_meta("icon"));
-			}
+			_update_valid_type_icons();
 		} break;
 
 		case NOTIFICATION_EXIT_TREE: {
@@ -1766,6 +1771,14 @@ void SceneTreeDialog::_notification(int p_what) {
 	}
 }
 
+void SceneTreeDialog::_update_valid_type_icons() {
+	filter->set_right_icon(get_editor_theme_icon(SNAME("Search")));
+	for (TextureRect *trect : valid_type_icons) {
+		trect->set_custom_minimum_size(Vector2(get_theme_constant(SNAME("class_icon_size"), EditorStringName(Editor)), 0));
+		trect->set_texture(trect->get_meta("icon"));
+	}
+}
+
 void SceneTreeDialog::_cancel() {
 	hide();
 }
diff --git a/editor/gui/scene_tree_editor.h b/editor/gui/scene_tree_editor.h
index e623c8405d3..eed6d4b9547 100644
--- a/editor/gui/scene_tree_editor.h
+++ b/editor/gui/scene_tree_editor.h
@@ -199,6 +199,7 @@ class SceneTreeDialog : public ConfirmationDialog {
 	LineEdit *filter = nullptr;
 	CheckButton *show_all_nodes = nullptr;
 	LocalVector<TextureRect *> valid_type_icons;
+	HBoxContainer *allowed_types_hbox = nullptr;
 
 	void _select();
 	void _cancel();
@@ -208,6 +209,7 @@ class SceneTreeDialog : public ConfirmationDialog {
 	void _show_all_nodes_changed(bool p_button_pressed);
 
 protected:
+	void _update_valid_type_icons();
 	void _notification(int p_what);
 	static void _bind_methods();