From 0b73b2c6cf465ffa1d70eca7ecf298de73577250 Mon Sep 17 00:00:00 2001
From: Charlie Jolly <mistajolly@gmail.com>
Date: Tue, 28 Dec 2021 11:39:18 +0000
Subject: Nodes: Improve link drag search for Math and Vector Math nodes

Exposes math operations via rna emums.

This uses the rna enum to build the search list.

Differential Revision: https://developer.blender.org/D13680
---
 .../blender/nodes/shader/nodes/node_shader_math.cc | 33 +++++++++++++---
 .../nodes/shader/nodes/node_shader_vector_math.cc  | 45 ++++++++++++++++++----
 2 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/source/blender/nodes/shader/nodes/node_shader_math.cc b/source/blender/nodes/shader/nodes/node_shader_math.cc
index e683d55edc8..677e62d14b2 100644
--- a/source/blender/nodes/shader/nodes/node_shader_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_math.cc
@@ -26,6 +26,8 @@
 #include "NOD_math_functions.hh"
 #include "NOD_socket_search_link.hh"
 
+#include "RNA_enum_types.h"
+
 /* **************** SCALAR MATH ******************** */
 
 namespace blender::nodes::node_shader_math_cc {
@@ -45,15 +47,36 @@ static void sh_node_math_declare(NodeDeclarationBuilder &b)
   b.add_output<decl::Float>(N_("Value"));
 };
 
+class SocketSearchOp {
+ public:
+  std::string socket_name;
+  NodeMathOperation mode = NODE_MATH_ADD;
+  void operator()(LinkSearchOpParams &params)
+  {
+    bNode &node = params.add_node("ShaderNodeMath");
+    node.custom1 = mode;
+    params.update_and_connect_available_socket(node, socket_name);
+  }
+};
+
 static void sh_node_math_gather_link_searches(GatherLinkSearchOpParams &params)
 {
-  /* For now, do something very basic (only exposing "Add", and a single "Value" socket). */
+  const NodeDeclaration &declaration = *params.node_type().fixed_declaration;
+  if (params.in_out() == SOCK_OUT) {
+    search_link_ops_for_declarations(params, declaration.outputs());
+    return;
+  }
+
+  /* Expose first Value socket. */
   if (params.node_tree().typeinfo->validate_link(
           static_cast<eNodeSocketDatatype>(params.other_socket().type), SOCK_FLOAT)) {
-    params.add_item(IFACE_("Value"), [](LinkSearchOpParams &params) {
-      bNode &node = params.add_node("ShaderNodeMath");
-      params.update_and_connect_available_socket(node, "Value");
-    });
+    for (const EnumPropertyItem *item = rna_enum_node_math_items; item->identifier != nullptr;
+         item++) {
+      if (item->name != nullptr) {
+        params.add_item(IFACE_(item->name),
+                        SocketSearchOp{"Value", (NodeMathOperation)item->value});
+      }
+    }
   }
 }
 
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 56465f57e96..312929abd0d 100644
--- a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
@@ -26,6 +26,8 @@
 #include "NOD_math_functions.hh"
 #include "NOD_socket_search_link.hh"
 
+#include "RNA_enum_types.h"
+
 namespace blender::nodes::node_shader_vector_math_cc {
 
 static void sh_node_vector_math_declare(NodeDeclarationBuilder &b)
@@ -39,15 +41,44 @@ static void sh_node_vector_math_declare(NodeDeclarationBuilder &b)
   b.add_output<decl::Float>(N_("Value"));
 };
 
+class SocketSearchOp {
+ public:
+  std::string socket_name;
+  NodeVectorMathOperation mode = NODE_VECTOR_MATH_ADD;
+  void operator()(LinkSearchOpParams &params)
+  {
+    bNode &node = params.add_node("ShaderNodeVectorMath");
+    node.custom1 = mode;
+    params.update_and_connect_available_socket(node, socket_name);
+  }
+};
+
 static void sh_node_vector_math_gather_link_searches(GatherLinkSearchOpParams &params)
 {
-  /* For now, do something very basic (only exposing "Add", and a single "Vector" socket). */
-  if (params.node_tree().typeinfo->validate_link(
-          static_cast<eNodeSocketDatatype>(params.other_socket().type), SOCK_VECTOR)) {
-    params.add_item(IFACE_("Vector"), [](LinkSearchOpParams &params) {
-      bNode &node = params.add_node("ShaderNodeVectorMath");
-      params.update_and_connect_available_socket(node, "Vector");
-    });
+  if (!ELEM(params.other_socket().type,
+            SOCK_FLOAT,
+            SOCK_BOOLEAN,
+            SOCK_INT,
+            SOCK_VECTOR,
+            SOCK_RGBA)) {
+    return;
+  }
+
+  for (const EnumPropertyItem *item = rna_enum_node_vec_math_items; item->identifier != nullptr;
+       item++) {
+    if (item->name != nullptr) {
+      if ((params.in_out() == SOCK_OUT) && ELEM(item->value,
+                                                NODE_VECTOR_MATH_LENGTH,
+                                                NODE_VECTOR_MATH_DISTANCE,
+                                                NODE_VECTOR_MATH_DOT_PRODUCT)) {
+        params.add_item(IFACE_(item->name),
+                        SocketSearchOp{"Value", (NodeVectorMathOperation)item->value});
+      }
+      else {
+        params.add_item(IFACE_(item->name),
+                        SocketSearchOp{"Vector", (NodeVectorMathOperation)item->value});
+      }
+    }
   }
 }
 
-- 
cgit v1.2.3