From be0d7522ea470c10d1c84c57e4e63a96f4dcd176 Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Tue, 4 Jun 2024 16:58:14 +0200
Subject: [PATCH] Geometry Nodes: support link drag search for capture
 attribute node

Now it's possible to use link-drag-search with the extend socket of the Capture Attribute node again.

Implementation wise, the main unexpected things I noticed are that
`update_and_connect_available_socket` did not update the node declaration and currently uses
socket names instead of identifiers. This works fine right now, but should eventually be changed
to use identifiers (separate from this commit though).

Pull Request: https://projects.blender.org/blender/blender/pulls/122716
---
 .../nodes/node_geo_attribute_capture.cc       | 22 +++++++++++++++++++
 .../nodes/intern/socket_search_link.cc        |  2 ++
 2 files changed, 24 insertions(+)

diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_capture.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_capture.cc
index 5771d67fd14..b40729c2585 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_capture.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_capture.cc
@@ -300,6 +300,27 @@ static void node_copy_storage(bNodeTree * /*dst_tree*/, bNode *dst_node, const b
   socket_items::copy_array<CaptureAttributeItemsAccessor>(*src_node, *dst_node);
 }
 
+static void node_gather_link_searches(GatherLinkSearchOpParams &params)
+{
+  const eNodeSocketDatatype type = eNodeSocketDatatype(params.other_socket().type);
+  if (type == SOCK_GEOMETRY) {
+    params.add_item(IFACE_("Geometry"), [](LinkSearchOpParams &params) {
+      bNode &node = params.add_node("GeometryNodeCaptureAttribute");
+      params.connect_available_socket(node, "Geometry");
+    });
+  }
+  if (!CaptureAttributeItemsAccessor::supports_socket_type(type)) {
+    return;
+  }
+
+  params.add_item(IFACE_("Value"), [type](LinkSearchOpParams &params) {
+    bNode &node = params.add_node("GeometryNodeCaptureAttribute");
+    socket_items::add_item_with_socket_type_and_name<CaptureAttributeItemsAccessor>(
+        node, type, params.socket.name);
+    params.update_and_connect_available_socket(node, params.socket.name);
+  });
+}
+
 static void node_register()
 {
   static blender::bke::bNodeType ntype;
@@ -315,6 +336,7 @@ static void node_register()
   ntype.draw_buttons = node_layout;
   ntype.draw_buttons_ex = node_layout_ex;
   ntype.register_operators = node_operators;
+  ntype.gather_link_search_ops = node_gather_link_searches;
   blender::bke::nodeRegisterType(&ntype);
 }
 NOD_REGISTER_NODE(node_register)
diff --git a/source/blender/nodes/intern/socket_search_link.cc b/source/blender/nodes/intern/socket_search_link.cc
index 197b8866e32..e8cf8738842 100644
--- a/source/blender/nodes/intern/socket_search_link.cc
+++ b/source/blender/nodes/intern/socket_search_link.cc
@@ -14,6 +14,7 @@
 #include "BLT_translation.hh"
 
 #include "NOD_node_declaration.hh"
+#include "NOD_socket.hh"
 #include "NOD_socket_search_link.hh"
 
 namespace blender::nodes {
@@ -91,6 +92,7 @@ bNode &LinkSearchOpParams::add_node(const bke::bNodeType &node_type)
 void LinkSearchOpParams::update_and_connect_available_socket(bNode &new_node,
                                                              StringRef socket_name)
 {
+  update_node_declaration_and_sockets(this->node_tree, new_node);
   if (new_node.typeinfo->updatefunc) {
     new_node.typeinfo->updatefunc(&node_tree, &new_node);
   }