From c593db5a2ffc2f0ad993e8d2a839f87404d2f4a6 Mon Sep 17 00:00:00 2001
From: Hans Goudey <h.goudey@me.com>
Date: Wed, 22 Dec 2021 17:22:10 -0600
Subject: Nodes: Add link drag search support for map range node

Previously only the float version of the node was connected to.
This adds connection operations for vector sockets, and exposes
the "Steps" socket properly when it's selected.
---
 .../nodes/shader/nodes/node_shader_map_range.cc    | 74 ++++++++++++++++++++--
 1 file changed, 68 insertions(+), 6 deletions(-)

(limited to 'source/blender/nodes')

diff --git a/source/blender/nodes/shader/nodes/node_shader_map_range.cc b/source/blender/nodes/shader/nodes/node_shader_map_range.cc
index 615ae276eb4..d0eab8d9727 100644
--- a/source/blender/nodes/shader/nodes/node_shader_map_range.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_map_range.cc
@@ -27,6 +27,8 @@
 
 #include "BLI_math_base_safe.h"
 
+#include "NOD_socket_search_link.hh"
+
 NODE_STORAGE_FUNCS(NodeMapRange)
 
 namespace blender::nodes {
@@ -50,8 +52,6 @@ static void sh_node_map_range_declare(NodeDeclarationBuilder &b)
   b.add_output<decl::Vector>(N_("Vector"));
 };
 
-}  // namespace blender::nodes
-
 static void node_shader_update_map_range(bNodeTree *ntree, bNode *node)
 {
   const NodeMapRange &storage = node_storage(*node);
@@ -80,7 +80,7 @@ static void node_shader_update_map_range(bNodeTree *ntree, bNode *node)
 
 static void node_shader_init_map_range(bNodeTree *UNUSED(ntree), bNode *node)
 {
-  NodeMapRange *data = (NodeMapRange *)MEM_callocN(sizeof(NodeMapRange), __func__);
+  NodeMapRange *data = MEM_cnew<NodeMapRange>(__func__);
   data->clamp = 1;
   data->data_type = CD_PROP_FLOAT;
   data->interpolation_type = NODE_MAP_RANGE_LINEAR;
@@ -89,6 +89,68 @@ static void node_shader_init_map_range(bNodeTree *UNUSED(ntree), bNode *node)
   node->storage = data;
 }
 
+class SocketSearchOp {
+ public:
+  std::string socket_name;
+  CustomDataType data_type;
+  int interpolation_type = NODE_MAP_RANGE_LINEAR;
+
+  void operator()(LinkSearchOpParams &params)
+  {
+    bNode &node = params.add_node("ShaderNodeMapRange");
+    node_storage(node).data_type = data_type;
+    node_storage(node).interpolation_type = interpolation_type;
+    params.update_and_connect_available_socket(node, socket_name);
+  }
+};
+
+static std::optional<CustomDataType> node_type_from_other_socket(const bNodeSocket &socket)
+{
+  switch (socket.type) {
+    case SOCK_FLOAT:
+    case SOCK_BOOLEAN:
+    case SOCK_INT:
+      return CD_PROP_FLOAT;
+    case SOCK_VECTOR:
+    case SOCK_RGBA:
+      return CD_PROP_FLOAT3;
+    default:
+      return {};
+  }
+}
+
+static void node_map_range_gather_link_searches(GatherLinkSearchOpParams &params)
+{
+  const std::optional<CustomDataType> type = node_type_from_other_socket(params.other_socket());
+  if (!type) {
+    return;
+  }
+
+  if (params.in_out() == SOCK_IN) {
+    if (*type == CD_PROP_FLOAT3) {
+      params.add_item(IFACE_("Vector"), SocketSearchOp{"Vector", *type}, 0);
+    }
+    else {
+      params.add_item(IFACE_("Value"), SocketSearchOp{"Value", *type}, 0);
+    }
+    params.add_item(IFACE_("From Min"), SocketSearchOp{"From Min", *type}, -1);
+    params.add_item(IFACE_("From Max"), SocketSearchOp{"From Max", *type}, -1);
+    params.add_item(IFACE_("To Min"), SocketSearchOp{"To Min", *type}, -2);
+    params.add_item(IFACE_("To Max"), SocketSearchOp{"To Max", *type}, -2);
+    params.add_item(IFACE_("Steps"), SocketSearchOp{"Steps", *type, NODE_MAP_RANGE_STEPPED}, -3);
+  }
+  else {
+    if (*type == CD_PROP_FLOAT3) {
+      params.add_item(IFACE_("Vector"), SocketSearchOp{"Vector", *type});
+    }
+    else {
+      params.add_item(IFACE_("Result"), SocketSearchOp{"Result", *type});
+    }
+  }
+}
+
+}  // namespace blender::nodes
+
 static const char *gpu_shader_get_name(int mode, bool use_vector)
 {
   if (use_vector) {
@@ -590,12 +652,12 @@ void register_node_type_sh_map_range()
 
   sh_fn_node_type_base(&ntype, SH_NODE_MAP_RANGE, "Map Range", NODE_CLASS_CONVERTER, 0);
   ntype.declare = blender::nodes::sh_node_map_range_declare;
-  node_type_init(&ntype, node_shader_init_map_range);
+  node_type_init(&ntype, blender::nodes::node_shader_init_map_range);
   node_type_storage(
       &ntype, "NodeMapRange", node_free_standard_storage, node_copy_standard_storage);
-  node_type_update(&ntype, node_shader_update_map_range);
+  node_type_update(&ntype, blender::nodes::node_shader_update_map_range);
   node_type_gpu(&ntype, gpu_shader_map_range);
   ntype.build_multi_function = blender::nodes::sh_node_map_range_build_multi_function;
-
+  ntype.gather_link_search_ops = blender::nodes::node_map_range_gather_link_searches;
   nodeRegisterType(&ntype);
 }
-- 
cgit v1.2.3