Geometry Nodes: Add search to named attribute input node

This commit adds attribute search the the attribute input field node.
Because it's a field node, finding which attribute to display without
increasing the complexity a lot isn't obvious. In this commit, all
attributes used by nodes in the current group are included.

When an attribute is chosen from the list, the node's data type is
updated, and links connected to the output socket are reconnected.

Ref T96271

Differential Revision: https://developer.blender.org/D14516

source/blender/editors/space_node/node_geometry_attribute_search.cc

@@ -14,11 +14,13 @@
 #include "DNA_space_types.h"
 
 #include "BKE_context.h"
+#include "BKE_node_tree_update.h"
 #include "BKE_object.h"
 
 #include "RNA_access.h"
 #include "RNA_enum_types.h"
 
+#include "ED_node.h"
 #include "ED_screen.h"
 #include "ED_undo.h"
 
@@ -45,6 +47,59 @@ struct AttributeSearchData {
 /* This class must not have a destructor, since it is used by buttons and freed with #MEM_freeN. */
 BLI_STATIC_ASSERT(std::is_trivially_destructible_v<AttributeSearchData>, "");
 
+static Vector<const GeometryAttributeInfo *> get_attribute_info_from_context(
+    const bContext &C, AttributeSearchData &data)
+{
+  SpaceNode *snode = CTX_wm_space_node(&C);
+  if (!snode) {
+    BLI_assert_unreachable();
+    return {};
+  }
+  bNodeTree *node_tree = snode->edittree;
+  if (node_tree == nullptr) {
+    BLI_assert_unreachable();
+    return {};
+  }
+  bNode *node = nodeFindNodebyName(node_tree, data.node_name);
+  if (node == nullptr) {
+    BLI_assert_unreachable();
+    return {};
+  }
+
+  /* For the attribute input node, collect attribute information from all nodes in the group. */
+  if (node->type == GEO_NODE_INPUT_NAMED_ATTRIBUTE) {
+    const geo_log::TreeLog *tree_log = geo_log::ModifierLog::find_tree_by_node_editor_context(
+        *snode);
+    if (tree_log == nullptr) {
+      return {};
+    }
+
+    Vector<const GeometryAttributeInfo *> attributes;
+    Set<StringRef> names;
+    tree_log->foreach_node_log([&](const geo_log::NodeLog &node_log) {
+      for (const geo_log::SocketLog &socket_log : node_log.input_logs()) {
+        const geo_log::ValueLog *value_log = socket_log.value();
+        if (const geo_log::GeometryValueLog *geo_value_log =
+                dynamic_cast<const geo_log::GeometryValueLog *>(value_log)) {
+          for (const GeometryAttributeInfo &attribute : geo_value_log->attributes()) {
+            if (names.add(attribute.name)) {
+              attributes.append(&attribute);
+            }
+          }
+        }
+      }
+    });
+    return attributes;
+  }
+
+  const geo_log::NodeLog *node_log = geo_log::ModifierLog::find_node_by_node_editor_context(
+      *snode, data.node_name);
+  if (node_log == nullptr) {
+    return {};
+  }
+  return node_log->lookup_available_attributes();
+}
+
 static void attribute_search_update_fn(
     const bContext *C, void *arg, const char *str, uiSearchItems *items, const bool is_first)
 {
@@ -54,17 +109,47 @@ static void attribute_search_update_fn(
 
   AttributeSearchData *data = static_cast<AttributeSearchData *>(arg);
 
-  SpaceNode *snode = CTX_wm_space_node(C);
-  const geo_log::NodeLog *node_log = geo_log::ModifierLog::find_node_by_node_editor_context(
-      *snode, data->node_name);
-  if (node_log == nullptr) {
-    return;
+  Vector<const GeometryAttributeInfo *> infos = get_attribute_info_from_context(*C, *data);
+
+  /* Remove the deprecated normal attribute from the search. */
+  for (const int i : infos.index_range()) {
+    if (infos[i]->domain == ATTR_DOMAIN_FACE && infos[i]->name == "normal") {
+      infos.remove(i);
+      break;
+    }
   }
-  Vector<const GeometryAttributeInfo *> infos = node_log->lookup_available_attributes();
 
   ui::attribute_search_add_items(str, true, infos, items, is_first);
 }
 
+/**
+ * Some custom data types don't correspond to node types and therefore can't be
+ * used by the named attribute input node. Find the best option or fallback to float.
+ */
+static CustomDataType data_type_in_attribute_input_node(const CustomDataType type)
+{
+  switch (type) {
+    case CD_PROP_FLOAT:
+    case CD_PROP_INT32:
+    case CD_PROP_FLOAT3:
+    case CD_PROP_COLOR:
+    case CD_PROP_BOOL:
+      return type;
+    case CD_MLOOPCOL:
+      return CD_PROP_COLOR;
+    case CD_PROP_STRING:
+      /* Unsupported currently. */
+      return CD_PROP_FLOAT;
+    case CD_PROP_FLOAT2:
+      /* No 2D vector sockets currently. */
+      return CD_PROP_FLOAT3;
+    case CD_PROP_INT8:
+      return CD_PROP_INT32;
+    default:
+      return CD_PROP_FLOAT;
+  }
+}
+
 static void attribute_search_exec_fn(bContext *C, void *data_v, void *item_v)
 {
   if (ED_screen_animation_playing(CTX_wm_manager(C))) {
@@ -97,6 +182,28 @@ static void attribute_search_exec_fn(bContext *C, void *data_v, void *item_v)
   }
   BLI_assert(socket->type == SOCK_STRING);
 
+  /* For the attribute input node, also adjust the type and links connected to the output. */
+  if (node->type == GEO_NODE_INPUT_NAMED_ATTRIBUTE) {
+    NodeGeometryInputNamedAttribute &storage = *(NodeGeometryInputNamedAttribute *)node->storage;
+    const CustomDataType new_type = data_type_in_attribute_input_node(item->data_type);
+    if (new_type != storage.data_type) {
+      storage.data_type = new_type;
+      /* Make the output socket with the new type on the attribute input node active. */
+      node->typeinfo->updatefunc(node_tree, node);
+
+      /* Relink all node links to the newly active output socket. */
+      bNodeSocket *output_socket = bke::node_find_enabled_output_socket(*node, "Attribute");
+      LISTBASE_FOREACH (bNodeLink *, link, &node_tree->links) {
+        if (link->fromnode == node) {
+          link->fromsock = output_socket;
+          BKE_ntree_update_tag_link_changed(node_tree);
+        }
+      }
+    }
+    BKE_ntree_update_tag_node_property(node_tree, node);
+    ED_node_tree_propagate_change(C, CTX_data_main(C), node_tree);
+  }
+
   bNodeSocketValueString *value = static_cast<bNodeSocketValueString *>(socket->default_value);
   BLI_strncpy(value->value, item->name.c_str(), MAX_NAME);