Geometry Nodes: Implement link drag search for two nodes

It was never added for the field on domain and field at index nodes.
They need special handling because they have many what should be
a multi-type socket declaration.

2 files changed, 37 insertions, 0 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc b/source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc
index 7d44ac34f15..bde4af12d84 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_field_at_index.cc
@@ -9,6 +9,8 @@
 
 #include "BLI_task.hh"
 
+#include "NOD_socket_search_link.hh"
+
 namespace blender::nodes::node_geo_field_at_index_cc {
 
 static void node_declare(NodeDeclarationBuilder &b)
@@ -70,6 +72,23 @@ static void node_update(bNodeTree *ntree, bNode *node)
   nodeSetSocketAvailability(ntree, sock_out_bool, data_type == CD_PROP_BOOL);
 }
 
+static void node_gather_link_searches(GatherLinkSearchOpParams &params)
+{
+  const NodeDeclaration &declaration = *params.node_type().fixed_declaration;
+  search_link_ops_for_declarations(params, declaration.inputs().take_front(1));
+
+  const bNodeType &node_type = params.node_type();
+  const std::optional<eCustomDataType> type = node_data_type_to_custom_data_type(
+      (eNodeSocketDatatype)params.other_socket().type);
+  if (type && *type != CD_PROP_STRING) {
+    params.add_item(IFACE_("Value"), [node_type, type](LinkSearchOpParams &params) {
+      bNode &node = params.add_node(node_type);
+      node.custom2 = *type;
+      params.update_and_connect_available_socket(node, "Value");
+    });
+  }
+}
+
 class FieldAtIndex final : public GeometryFieldInput {
  private:
   Field<int> index_field_;
@@ -175,5 +194,6 @@ void register_node_type_geo_field_at_index()
   ntype.draw_buttons = file_ns::node_layout;
   ntype.initfunc = file_ns::node_init;
   ntype.updatefunc = file_ns::node_update;
+  ntype.gather_link_search_ops = file_ns::node_gather_link_searches;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc b/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc
index 59e243db4a2..abebe5de893 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_field_on_domain.cc
@@ -9,6 +9,8 @@
 
 #include "BLI_task.hh"
 
+#include "NOD_socket_search_link.hh"
+
 namespace blender::nodes::node_geo_field_on_domain_cc {
 
 static void node_declare(NodeDeclarationBuilder &b)
@@ -67,6 +69,20 @@ static void node_update(bNodeTree *ntree, bNode *node)
   nodeSetSocketAvailability(ntree, sock_out_bool, data_type == CD_PROP_BOOL);
 }
 
+static void node_gather_link_searches(GatherLinkSearchOpParams &params)
+{
+  const bNodeType &node_type = params.node_type();
+  const std::optional<eCustomDataType> type = node_data_type_to_custom_data_type(
+      (eNodeSocketDatatype)params.other_socket().type);
+  if (type && *type != CD_PROP_STRING) {
+    params.add_item(IFACE_("Value"), [node_type, type](LinkSearchOpParams &params) {
+      bNode &node = params.add_node(node_type);
+      node.custom2 = *type;
+      params.update_and_connect_available_socket(node, "Value");
+    });
+  }
+}
+
 class FieldOnDomain final : public GeometryFieldInput {
  private:
   GField src_field_;
@@ -143,5 +159,6 @@ void register_node_type_geo_field_on_domain()
   ntype.draw_buttons = file_ns::node_layout;
   ntype.initfunc = file_ns::node_init;
   ntype.updatefunc = file_ns::node_update;
+  ntype.gather_link_search_ops = file_ns::node_gather_link_searches;
   nodeRegisterType(&ntype);
 }