From c0b2c75c441db711ea27f23b0539371eeb159649 Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Wed, 24 Feb 2021 15:59:44 +0100
Subject: Geometry Nodes: store available attribute names in node ui storage

This information will be used by the attribute search in string sockets.

Ref T85657.

Differential Revision: https://developer.blender.org/D10462
---
 source/blender/blenkernel/BKE_node_ui_storage.hh   |  7 +++++
 .../blender/blenkernel/intern/node_ui_storage.cc   |  9 +++++++
 source/blender/modifiers/intern/MOD_nodes.cc       | 30 ++++++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/source/blender/blenkernel/BKE_node_ui_storage.hh b/source/blender/blenkernel/BKE_node_ui_storage.hh
index 951c3bdc62d..231eb11d473 100644
--- a/source/blender/blenkernel/BKE_node_ui_storage.hh
+++ b/source/blender/blenkernel/BKE_node_ui_storage.hh
@@ -19,6 +19,7 @@
 #include "BLI_hash.hh"
 #include "BLI_map.hh"
 #include "BLI_session_uuid.h"
+#include "BLI_set.hh"
 
 #include "DNA_ID.h"
 #include "DNA_modifier_types.h"
@@ -76,6 +77,7 @@ struct NodeWarning {
 
 struct NodeUIStorage {
   blender::Vector<NodeWarning> warnings;
+  blender::Set<std::string> attribute_name_hints;
 };
 
 struct NodeTreeUIStorage {
@@ -94,3 +96,8 @@ void BKE_nodetree_error_message_add(bNodeTree &ntree,
                                     const bNode &node,
                                     const NodeWarningType type,
                                     std::string message);
+
+void BKE_nodetree_attribute_hint_add(bNodeTree &ntree,
+                                     const NodeTreeEvaluationContext &context,
+                                     const bNode &node,
+                                     const blender::StringRef attribute_name);
diff --git a/source/blender/blenkernel/intern/node_ui_storage.cc b/source/blender/blenkernel/intern/node_ui_storage.cc
index e03617a6505..ea492c4d36a 100644
--- a/source/blender/blenkernel/intern/node_ui_storage.cc
+++ b/source/blender/blenkernel/intern/node_ui_storage.cc
@@ -136,3 +136,12 @@ void BKE_nodetree_error_message_add(bNodeTree &ntree,
   NodeUIStorage &node_ui_storage = find_node_ui_storage(ntree, context, node);
   node_ui_storage.warnings.append({type, std::move(message)});
 }
+
+void BKE_nodetree_attribute_hint_add(bNodeTree &ntree,
+                                     const NodeTreeEvaluationContext &context,
+                                     const bNode &node,
+                                     const StringRef attribute_name)
+{
+  NodeUIStorage &node_ui_storage = find_node_ui_storage(ntree, context, node);
+  node_ui_storage.attribute_name_hints.add_as(attribute_name);
+}
diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index fe5a8bc59d4..c7731815a2a 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -82,6 +82,7 @@ using blender::Map;
 using blender::Set;
 using blender::Span;
 using blender::StringRef;
+using blender::StringRefNull;
 using blender::Vector;
 using blender::bke::PersistentCollectionHandle;
 using blender::bke::PersistentDataHandleMap;
@@ -388,6 +389,8 @@ class GeometryNodesEvaluator {
   {
     const bNode &bnode = params.node();
 
+    this->store_ui_hints(node, params);
+
     /* Use the geometry-node-execute callback if it exists. */
     if (bnode.typeinfo->geometry_node_execute != nullptr) {
       bnode.typeinfo->geometry_node_execute(params);
@@ -405,6 +408,33 @@ class GeometryNodesEvaluator {
     this->execute_unknown_node(node, params);
   }
 
+  void store_ui_hints(const DNode &node, GeoNodeExecParams params) const
+  {
+    for (const DInputSocket *dsocket : node.inputs()) {
+      if (!dsocket->is_available()) {
+        continue;
+      }
+      if (dsocket->bsocket()->type != SOCK_GEOMETRY) {
+        continue;
+      }
+
+      bNodeTree *btree_cow = node.node_ref().tree().btree();
+      bNodeTree *btree_original = (bNodeTree *)DEG_get_original_id((ID *)btree_cow);
+      const NodeTreeEvaluationContext context(*self_object_, *modifier_);
+
+      const GeometrySet &geometry_set = params.get_input<GeometrySet>(dsocket->identifier());
+      const Vector<const GeometryComponent *> components = geometry_set.get_components_for_read();
+
+      for (const GeometryComponent *component : components) {
+        component->attribute_foreach([&](StringRefNull attribute_name,
+                                         const AttributeMetaData &UNUSED(meta_data)) {
+          BKE_nodetree_attribute_hint_add(*btree_original, context, *node.bnode(), attribute_name);
+          return true;
+        });
+      }
+    }
+  }
+
   void execute_multi_function_node(const DNode &node,
                                    GeoNodeExecParams params,
                                    const MultiFunction &fn)
-- 
cgit v1.2.3