1 files changed, 113 insertions, 6 deletions

diff --git a/source/blender/editors/space_node/node_geometry_attribute_search.cc b/source/blender/editors/space_node/node_geometry_attribute_search.cc
index 680a3be1f49..2536496b50d 100644
--- a/source/blender/editors/space_node/node_geometry_attribute_search.cc
+++ b/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);