Nodes: Consistent link drag search for math and vector math nodes

Previously operations for the math node when connecting to
outputs weren't added. It also used a different method to
check whether the link would be valid.

2 files changed, 17 insertions, 29 deletions

diff --git a/source/blender/nodes/shader/nodes/node_shader_math.cc b/source/blender/nodes/shader/nodes/node_shader_math.cc
index 059c0bfc5b5..ca30b16f7ff 100644
--- a/source/blender/nodes/shader/nodes/node_shader_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_math.cc
@@ -61,32 +61,24 @@ class SocketSearchOp {
 
 static void sh_node_math_gather_link_searches(GatherLinkSearchOpParams &params)
 {
-  const NodeDeclaration &declaration = *params.node_type().fixed_declaration;
-  if (params.in_out() == SOCK_OUT) {
-    search_link_ops_for_declarations(params, declaration.outputs());
+  if (!params.node_tree().typeinfo->validate_link(
+          static_cast<eNodeSocketDatatype>(params.other_socket().type), SOCK_FLOAT)) {
     return;
   }
 
-  /* Expose first Value socket. */
-  if (params.node_tree().typeinfo->validate_link(
-          static_cast<eNodeSocketDatatype>(params.other_socket().type), SOCK_FLOAT)) {
-
-    const bool is_geometry_node_tree = params.node_tree().type == NTREE_GEOMETRY;
-    const int weight = ELEM(params.other_socket().type, SOCK_FLOAT, SOCK_BOOLEAN, SOCK_INT) ? 0 :
-                                                                                              -1;
-
-    for (const EnumPropertyItem *item = rna_enum_node_math_items; item->identifier != nullptr;
-         item++) {
-      if (item->name != nullptr && item->identifier[0] != '\0') {
-        const int gn_weight =
-            (is_geometry_node_tree &&
-             ELEM(item->value, NODE_MATH_COMPARE, NODE_MATH_GREATER_THAN, NODE_MATH_LESS_THAN)) ?
-                -1 :
-                weight;
-        params.add_item(IFACE_(item->name),
-                        SocketSearchOp{"Value", (NodeMathOperation)item->value},
-                        gn_weight);
-      }
+  const bool is_geometry_node_tree = params.node_tree().type == NTREE_GEOMETRY;
+  const int weight = ELEM(params.other_socket().type, SOCK_FLOAT, SOCK_BOOLEAN, SOCK_INT) ? 0 : -1;
+
+  for (const EnumPropertyItem *item = rna_enum_node_math_items; item->identifier != nullptr;
+       item++) {
+    if (item->name != nullptr && item->identifier[0] != '\0') {
+      const int gn_weight =
+          (is_geometry_node_tree &&
+           ELEM(item->value, NODE_MATH_COMPARE, NODE_MATH_GREATER_THAN, NODE_MATH_LESS_THAN)) ?
+              -1 :
+              weight;
+      params.add_item(
+          IFACE_(item->name), SocketSearchOp{"Value", (NodeMathOperation)item->value}, gn_weight);
     }
   }
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
index 8d55c3a38b7..1276592e2af 100644
--- a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
@@ -55,12 +55,8 @@ class SocketSearchOp {
 
 static void sh_node_vector_math_gather_link_searches(GatherLinkSearchOpParams &params)
 {
-  if (!ELEM(params.other_socket().type,
-            SOCK_FLOAT,
-            SOCK_BOOLEAN,
-            SOCK_INT,
-            SOCK_VECTOR,
-            SOCK_RGBA)) {
+  if (!params.node_tree().typeinfo->validate_link(
+          static_cast<eNodeSocketDatatype>(params.other_socket().type), SOCK_VECTOR)) {
     return;
   }