From da431a3879122cda2b56df4310b33c426a12be67 Mon Sep 17 00:00:00 2001
From: Martin Felis <martin@fysx.org>
Date: Wed, 1 May 2024 13:03:57 +0200
Subject: [PATCH] Fixed setting location of new nodes and minor connection
 rejection refactor.

---
 src/AnimGraph/AnimGraphEditor.cc | 119 +++++++++++++++----------------
 1 file changed, 56 insertions(+), 63 deletions(-)

diff --git a/src/AnimGraph/AnimGraphEditor.cc b/src/AnimGraph/AnimGraphEditor.cc
index 83b338f..dd34997 100644
--- a/src/AnimGraph/AnimGraphEditor.cc
+++ b/src/AnimGraph/AnimGraphEditor.cc
@@ -454,70 +454,64 @@ void AnimGraphEditorUpdate(ax::NodeEditor::EditorContext* context) {
     ax::NodeEditor::PinId input_pin_id, output_pin_id;
     if (ax::NodeEditor::QueryNewLink(&input_pin_id, &output_pin_id)) {
       if (input_pin_id && output_pin_id) {
-        if (ax::NodeEditor::AcceptNewItem()) {
-          int source_node_index;
-          int source_node_socket_index;
+        int source_node_index;
+        int source_node_socket_index;
 
-          OutputPinIdToNodeIndexAndSocketIndex(
-              input_pin_id.Get(),
-              &source_node_index,
-              &source_node_socket_index);
-
-          const AnimNodeResource* source_node =
-              sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
-                  ->m_blend_tree_resource.GetNode(source_node_index);
-          if (source_node->m_socket_accessor->m_outputs.size()
-              < source_node_socket_index) {
-            source_node_socket_index = -1;
-          }
-
-          int target_node_index;
-          int target_node_socket_index;
-
-          InputPinIdToNodeIndexAndSocketIndex(
-              output_pin_id.Get(),
-              &target_node_index,
-              &target_node_socket_index);
-
-          const AnimNodeResource* target_node =
-              sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
-                  ->m_blend_tree_resource.GetNode(target_node_index);
-          if (target_node->m_socket_accessor->m_inputs.size()
-              < target_node_socket_index) {
-            target_node_socket_index = -1;
-          }
-
-          const Socket* source_socket =
-              sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
-                  ->m_blend_tree_resource.GetNodeOutputSocketByIndex(
-                      source_node,
-                      source_node_socket_index);
-          const Socket* target_socket =
-              sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
-                  ->m_blend_tree_resource.GetNodeInputSocketByIndex(
-                      target_node,
-                      target_node_socket_index);
-
-          if (source_socket == nullptr || target_socket == nullptr) {
-            ax::NodeEditor::RejectNewItem();
-          } else {
-            const std::string& source_socket_name =
-                source_node->m_socket_accessor
-                    ->m_outputs[source_node_socket_index]
-                    .m_name;
-
-            const std::string& target_socket_name =
-                target_node->m_socket_accessor
-                    ->m_inputs[target_node_socket_index]
-                    .m_name;
+        OutputPinIdToNodeIndexAndSocketIndex(
+            input_pin_id.Get(),
+            &source_node_index,
+            &source_node_socket_index);
 
+        const AnimNodeResource* source_node =
             sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
-                ->m_blend_tree_resource.ConnectSockets(
+                ->m_blend_tree_resource.GetNode(source_node_index);
+        if (source_node->m_socket_accessor->m_outputs.size()
+            < source_node_socket_index) {
+          source_node_socket_index = -1;
+        }
+
+        int target_node_index;
+        int target_node_socket_index;
+
+        InputPinIdToNodeIndexAndSocketIndex(
+            output_pin_id.Get(),
+            &target_node_index,
+            &target_node_socket_index);
+
+        const AnimNodeResource* target_node =
+            sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
+                ->m_blend_tree_resource.GetNode(target_node_index);
+        if (target_node->m_socket_accessor->m_inputs.size()
+            < target_node_socket_index) {
+          target_node_socket_index = -1;
+        }
+
+        const Socket* source_socket =
+            sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
+                ->m_blend_tree_resource.GetNodeOutputSocketByIndex(
                     source_node,
-                    source_socket_name,
+                    source_node_socket_index);
+        const Socket* target_socket =
+            sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
+                ->m_blend_tree_resource.GetNodeInputSocketByIndex(
                     target_node,
-                    target_socket_name);
-          }
+                    target_node_socket_index);
+
+        if (source_socket == nullptr || target_socket == nullptr
+            || !sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
+                    ->m_blend_tree_resource.IsConnectionValid(
+                        source_node,
+                        source_socket->m_name,
+                        target_node,
+                        target_socket->m_name)) {
+          ax::NodeEditor::RejectNewItem();
+        } else if (ax::NodeEditor::AcceptNewItem()) {
+          sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
+              ->m_blend_tree_resource.ConnectSockets(
+                  source_node,
+                  source_socket->m_name,
+                  target_node,
+                  target_socket->m_name);
         }
       }
     }
@@ -574,10 +568,9 @@ void AnimGraphEditorUpdate(ax::NodeEditor::EditorContext* context) {
       if (!node_type_name.empty()) {
         AnimNodeResource* node_resource =
             AnimNodeResourceFactory(node_type_name);
-        size_t node_id =
-            sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
-                ->m_blend_tree_resource.GetNumNodes();
-        ax::NodeEditor::SetNodePosition(node_id, sEditorState.mousePopupStart);
+        ax::NodeEditor::SetNodePosition(
+            ax::NodeEditor::NodeId(node_resource),
+            sEditorState.mousePopupStart);
         sEditorState.hierarchyStack[sEditorState.hierarchyStackIndex]
             ->m_blend_tree_resource.AddNode(node_resource);
       }