Geometry Nodes: Add domain and data type to attribute search

This patch adds domain and data type information to each row of the
attribute search menu. The data type is displayed on the right, just
like how the list is exposed for the existing point cloud and hair
attribute panels. The domain is exposed on the left like the menu
hierarchy from menu search.

For the implementation, the attribute hint information is stored as a
set instead of a multi-value map so that every item (which we need to
point to descretely in the search process) contains the necessary data
type and domain information by itself. We also need to allocate a new
struct for every button, which requires a change to allow passing a
newly allocated argument to search buttons.

Note that the search does't yet handle the case where there are two
attributes with the same name but different domains or data types in
the input geometry set. That will be handled as a separate improvement.

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

11 files changed, 109 insertions, 46 deletions

diff --git a/source/blender/blenkernel/BKE_node_ui_storage.hh b/source/blender/blenkernel/BKE_node_ui_storage.hh
index aa7e5180b03..5f9c039ef9e 100644
--- a/source/blender/blenkernel/BKE_node_ui_storage.hh
+++ b/source/blender/blenkernel/BKE_node_ui_storage.hh
@@ -20,7 +20,6 @@
 
 #include "BLI_hash.hh"
 #include "BLI_map.hh"
-#include "BLI_multi_value_map.hh"
 #include "BLI_session_uuid.h"
 #include "BLI_set.hh"
 
@@ -80,30 +79,37 @@ struct NodeWarning {
 };
 
 struct AvailableAttributeInfo {
+  std::string name;
   AttributeDomain domain;
   CustomDataType data_type;
 
   uint64_t hash() const
   {
-    uint64_t domain_hash = (uint64_t)domain;
-    uint64_t data_type_hash = (uint64_t)data_type;
-    return (domain_hash * 33) ^ (data_type_hash * 89);
+    return blender::get_default_hash(name);
   }
 
   friend bool operator==(const AvailableAttributeInfo &a, const AvailableAttributeInfo &b)
   {
-    return a.domain == b.domain && a.data_type == b.data_type;
+    return a.name == b.name;
   }
 };
 
 struct NodeUIStorage {
   blender::Vector<NodeWarning> warnings;
-  blender::MultiValueMap<std::string, AvailableAttributeInfo> attribute_hints;
+  blender::Set<AvailableAttributeInfo> attribute_hints;
 };
 
 struct NodeTreeUIStorage {
   blender::Map<NodeTreeEvaluationContext, blender::Map<std::string, NodeUIStorage>> context_map;
   std::mutex context_map_mutex;
+
+  /**
+   * Attribute search uses this to store the fake info for the string typed into a node, in order
+   * to pass the info to the execute callback that sets node socket values. This is mutable since
+   * we can count on only one attribute search being open at a time, and there is no real data
+   * stored here.
+   */
+  mutable AvailableAttributeInfo dummy_info_for_search;
 };
 
 const NodeUIStorage *BKE_node_tree_ui_storage_get_from_context(const bContext *C,
diff --git a/source/blender/blenkernel/intern/node_ui_storage.cc b/source/blender/blenkernel/intern/node_ui_storage.cc
index f2a152ac00d..cc910bab6ac 100644
--- a/source/blender/blenkernel/intern/node_ui_storage.cc
+++ b/source/blender/blenkernel/intern/node_ui_storage.cc
@@ -163,6 +163,6 @@ void BKE_nodetree_attribute_hint_add(bNodeTree &ntree,
                                      const CustomDataType data_type)
 {
   NodeUIStorage &node_ui_storage = node_ui_storage_ensure(ntree, context, node);
-  node_ui_storage.attribute_hints.add_as(attribute_name,
-                                         AvailableAttributeInfo{domain, data_type});
+  node_ui_storage.attribute_hints.add_as(
+      AvailableAttributeInfo{attribute_name, domain, data_type});
 }
diff --git a/source/blender/editors/include/UI_interface.h b/source/blender/editors/include/UI_interface.h
index dfe0898a85b..7a189139358 100644
--- a/source/blender/editors/include/UI_interface.h
+++ b/source/blender/editors/include/UI_interface.h
@@ -1600,6 +1600,7 @@ void UI_but_func_search_set(uiBut *but,
                             uiButSearchCreateFn search_create_fn,
                             uiButSearchUpdateFn search_update_fn,
                             void *arg,
+                            const bool free_arg,
                             uiButSearchArgFreeFn search_arg_free_fn,
                             uiButHandleFunc search_exec_fn,
                             void *active);
diff --git a/source/blender/editors/interface/interface.c b/source/blender/editors/interface/interface.c
index 3c0f0587741..90d604b3190 100644
--- a/source/blender/editors/interface/interface.c
+++ b/source/blender/editors/interface/interface.c
@@ -6602,6 +6602,8 @@ uiBut *uiDefSearchBut(uiBlock *block,
  * \param search_create_fn: Function to create the menu.
  * \param search_update_fn: Function to refresh search content after the search text has changed.
  * \param arg: user value.
+ * \param free_arg: Set to true if the argument is newly allocated memory for every redraw and
+ * should be freed when the button is destroyed.
  * \param search_arg_free_fn: When non-null, use this function to free \a arg.
  * \param search_exec_fn: Function that executes the action, gets \a arg as the first argument.
  * The second argument as the active item-pointer
@@ -6612,6 +6614,7 @@ void UI_but_func_search_set(uiBut *but,
                             uiButSearchCreateFn search_create_fn,
                             uiButSearchUpdateFn search_update_fn,
                             void *arg,
+                            const bool free_arg,
                             uiButSearchArgFreeFn search_arg_free_fn,
                             uiButHandleFunc search_exec_fn,
                             void *active)
@@ -6647,11 +6650,17 @@ void UI_but_func_search_set(uiBut *but,
     }
 #endif
     /* Handling will pass the active item as arg2 later, so keep it NULL here. */
-    UI_but_func_set(but, search_exec_fn, search_but->arg, NULL);
+    if (free_arg) {
+      UI_but_funcN_set(but, search_exec_fn, search_but->arg, NULL);
+    }
+    else {
+      UI_but_func_set(but, search_exec_fn, search_but->arg, NULL);
+    }
   }
 
-  /* search buttons show red-alert if item doesn't exist, not for menus */
-  if (0 == (but->block->flag & UI_BLOCK_LOOP)) {
+  /* search buttons show red-alert if item doesn't exist, not for menus. Don't do this for
+   * buttons where any result is valid anyway, since any string will be valid anyway. */
+  if (0 == (but->block->flag & UI_BLOCK_LOOP) && !search_but->results_are_suggestions) {
     /* skip empty buttons, not all buttons need input, we only show invalid */
     if (but->drawstr[0]) {
       ui_but_search_refresh(search_but);
@@ -6791,6 +6800,7 @@ uiBut *uiDefSearchButO_ptr(uiBlock *block,
                          ui_searchbox_create_generic,
                          operator_enum_search_update_fn,
                          but,
+                         false,
                          NULL,
                          operator_enum_search_exec_fn,
                          NULL);
diff --git a/source/blender/editors/interface/interface_layout.c b/source/blender/editors/interface/interface_layout.c
index 4430d00f2e3..cabd98902a6 100644
--- a/source/blender/editors/interface/interface_layout.c
+++ b/source/blender/editors/interface/interface_layout.c
@@ -2718,6 +2718,7 @@ uiBut *ui_but_add_search(
                            ui_searchbox_create_generic,
                            ui_rna_collection_search_update_fn,
                            coll_search,
+                           false,
                            ui_rna_collection_search_arg_free_fn,
                            NULL,
                            NULL);
diff --git a/source/blender/editors/interface/interface_template_search_menu.c b/source/blender/editors/interface/interface_template_search_menu.c
index df19bdb650e..91ad6619889 100644
--- a/source/blender/editors/interface/interface_template_search_menu.c
+++ b/source/blender/editors/interface/interface_template_search_menu.c
@@ -1148,6 +1148,7 @@ void UI_but_func_menu_search(uiBut *but)
                          ui_searchbox_create_menu,
                          menu_search_update_fn,
                          data,
+                         false,
                          menu_search_arg_free_fn,
                          menu_search_exec_fn,
                          NULL);
diff --git a/source/blender/editors/interface/interface_template_search_operator.c b/source/blender/editors/interface/interface_template_search_operator.c
index 2c83f184ff0..2b765a1a2f5 100644
--- a/source/blender/editors/interface/interface_template_search_operator.c
+++ b/source/blender/editors/interface/interface_template_search_operator.c
@@ -121,6 +121,7 @@ void UI_but_func_operator_search(uiBut *but)
                          operator_search_update_fn,
                          NULL,
                          false,
+                         NULL,
                          operator_search_exec_fn,
                          NULL);
 }
diff --git a/source/blender/editors/interface/interface_templates.c b/source/blender/editors/interface/interface_templates.c
index 2d1663a3ecd..760fbe75688 100644
--- a/source/blender/editors/interface/interface_templates.c
+++ b/source/blender/editors/interface/interface_templates.c
@@ -309,6 +309,7 @@ static uiBlock *template_common_search_menu(const bContext *C,
                          ui_searchbox_create_generic,
                          search_update_fn,
                          search_arg,
+                         false,
                          NULL,
                          search_exec_fn,
                          active_item);
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 6d0cd254505..87944842079 100644
--- a/source/blender/editors/space_node/node_geometry_attribute_search.cc
+++ b/source/blender/editors/space_node/node_geometry_attribute_search.cc
@@ -30,6 +30,11 @@
 #include "BKE_node_ui_storage.hh"
 #include "BKE_object.h"
 
+#include "RNA_access.h"
+#include "RNA_enum_types.h"
+
+#include "BLT_translation.h"
+
 #include "UI_interface.h"
 #include "UI_resources.h"
 
@@ -37,44 +42,77 @@
 
 using blender::IndexRange;
 using blender::Map;
-using blender::MultiValueMap;
 using blender::Set;
 using blender::StringRef;
 
 struct AttributeSearchData {
   const bNodeTree &node_tree;
   const bNode &node;
+  bNodeSocket &socket;
+};
 
-  uiBut *search_button;
+/* 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>, "");
 
-  /* Used to keep track of a button pointer over multiple redraws. Since the UI code
-   * may reallocate the button, without this we might end up with a dangling pointer. */
-  uiButStore *button_store;
-  uiBlock *button_store_block;
-};
+static StringRef attribute_data_type_string(const CustomDataType type)
+{
+  const char *name = nullptr;
+  RNA_enum_name_from_value(rna_enum_attribute_type_items, type, &name);
+  return StringRef(IFACE_(name));
+}
+
+static StringRef attribute_domain_string(const AttributeDomain domain)
+{
+  const char *name = nullptr;
+  RNA_enum_name_from_value(rna_enum_attribute_domain_items, domain, &name);
+  return StringRef(IFACE_(name));
+}
+
+/* Unicode arrow. */
+#define MENU_SEP "\xe2\x96\xb6"
+
+static bool attribute_search_item_add(uiSearchItems *items, const AvailableAttributeInfo &item)
+{
+  const StringRef data_type_name = attribute_data_type_string(item.data_type);
+  const StringRef domain_name = attribute_domain_string(item.domain);
+  std::string search_item_text = domain_name + " " + MENU_SEP + item.name + UI_SEP_CHAR +
+                                 data_type_name;
+
+  return UI_search_item_add(
+      items, search_item_text.c_str(), (void *)&item, ICON_NONE, UI_BUT_HAS_SEP_CHAR, 0);
+}
 
 static void attribute_search_update_fn(
     const bContext *C, void *arg, const char *str, uiSearchItems *items, const bool is_first)
 {
   AttributeSearchData *data = static_cast<AttributeSearchData *>(arg);
+  NodeTreeUIStorage *tree_ui_storage = data->node_tree.ui_storage;
+  if (tree_ui_storage == nullptr) {
+    return;
+  }
   const NodeUIStorage *ui_storage = BKE_node_tree_ui_storage_get_from_context(
       C, data->node_tree, data->node);
   if (ui_storage == nullptr) {
     return;
   }
 
-  const MultiValueMap<std::string, AvailableAttributeInfo> &attribute_hints =
-      ui_storage->attribute_hints;
+  const Set<AvailableAttributeInfo> &attribute_hints = ui_storage->attribute_hints;
 
-  if (str[0] != '\0' && attribute_hints.lookup_as(StringRef(str)).is_empty()) {
-    /* Any string may be valid, so add the current search string with the hints. */
-    UI_search_item_add(items, str, (void *)str, ICON_ADD, 0, 0);
+  /* Any string may be valid, so add the current search string along with the hints. */
+  if (str[0] != '\0') {
+    /* Note that the attribute domain and data type are dummies, since
+     * #AvailableAttributeInfo equality is only based on the string. */
+    if (!attribute_hints.contains(AvailableAttributeInfo{str, ATTR_DOMAIN_AUTO, CD_PROP_BOOL})) {
+      tree_ui_storage->dummy_info_for_search.name = std::string(str);
+      UI_search_item_add(items, str, &tree_ui_storage->dummy_info_for_search, ICON_ADD, 0, 0);
+    }
   }
 
   if (str[0] == '\0' && !is_first) {
     /* Allow clearing the text field when the string is empty, but not on the first pass,
      * or opening an attribute field for the first time would show this search item. */
-    UI_search_item_add(items, str, (void *)str, ICON_X, 0, 0);
+    tree_ui_storage->dummy_info_for_search.name = std::string(str);
+    UI_search_item_add(items, str, &tree_ui_storage->dummy_info_for_search, ICON_X, 0, 0);
   }
 
   /* Don't filter when the menu is first opened, but still run the search
@@ -82,16 +120,16 @@ static void attribute_search_update_fn(
   const char *string = is_first ? "" : str;
 
   StringSearch *search = BLI_string_search_new();
-  for (const std::string &attribute_name : attribute_hints.keys()) {
-    BLI_string_search_add(search, attribute_name.c_str(), (void *)&attribute_name);
+  for (const AvailableAttributeInfo &item : attribute_hints) {
+    BLI_string_search_add(search, item.name.c_str(), (void *)&item);
   }
 
-  std::string **filtered_items;
+  AvailableAttributeInfo **filtered_items;
   const int filtered_amount = BLI_string_search_query(search, string, (void ***)&filtered_items);
 
   for (const int i : IndexRange(filtered_amount)) {
-    std::string *item = filtered_items[i];
-    if (!UI_search_item_add(items, item->c_str(), item, ICON_NONE, 0, 0)) {
+    const AvailableAttributeInfo *item = filtered_items[i];
+    if (!attribute_search_item_add(items, *item)) {
       break;
     }
   }
@@ -100,12 +138,14 @@ static void attribute_search_update_fn(
   BLI_string_search_free(search);
 }
 
-static void attribute_search_free_fn(void *arg)
+static void attribute_search_exec_fn(bContext *UNUSED(C), void *data_v, void *item_v)
 {
-  AttributeSearchData *data = static_cast<AttributeSearchData *>(arg);
+  AttributeSearchData *data = static_cast<AttributeSearchData *>(data_v);
+  AvailableAttributeInfo *item = static_cast<AvailableAttributeInfo *>(item_v);
 
-  UI_butstore_free(data->button_store_block, data->button_store);
-  delete data;
+  bNodeSocket &socket = data->socket;
+  bNodeSocketValueString *value = static_cast<bNodeSocketValueString *>(socket.default_value);
+  BLI_strncpy(value->value, item->name.c_str(), MAX_NAME);
 }
 
 void node_geometry_add_attribute_search_button(const bNodeTree *node_tree,
@@ -132,22 +172,17 @@ void node_geometry_add_attribute_search_button(const bNodeTree *node_tree,
                                  0.0f,
                                  "");
 
-  AttributeSearchData *data = new AttributeSearchData{
-      *node_tree,
-      *node,
-      but,
-      UI_butstore_create(block),
-      block,
-  };
-
-  UI_butstore_register(data->button_store, &data->search_button);
+  AttributeSearchData *data = OBJECT_GUARDED_NEW(
+      AttributeSearchData, {*node_tree, *node, *static_cast<bNodeSocket *>(socket_ptr->data)});
 
   UI_but_func_search_set_results_are_suggestions(but, true);
+  UI_but_func_search_set_sep_string(but, MENU_SEP);
   UI_but_func_search_set(but,
                          nullptr,
                          attribute_search_update_fn,
                          static_cast<void *>(data),
-                         attribute_search_free_fn,
+                         true,
                          nullptr,
+                         attribute_search_exec_fn,
                          nullptr);
 }
diff --git a/source/blender/editors/space_node/node_select.c b/source/blender/editors/space_node/node_select.c
index 6548f6f5775..de63aa07acb 100644
--- a/source/blender/editors/space_node/node_select.c
+++ b/source/blender/editors/space_node/node_select.c
@@ -1260,7 +1260,8 @@ static uiBlock *node_find_menu(bContext *C, ARegion *region, void *arg_op)
                        0,
                        0,
                        "");
-  UI_but_func_search_set(but, NULL, node_find_update_fn, op->type, NULL, node_find_exec_fn, NULL);
+  UI_but_func_search_set(
+      but, NULL, node_find_update_fn, op->type, false, NULL, node_find_exec_fn, NULL);
   UI_but_flag_enable(but, UI_BUT_ACTIVATE_ON_INIT);
 
   /* fake button, it holds space for search items */
diff --git a/source/blender/editors/space_outliner/outliner_tools.c b/source/blender/editors/space_outliner/outliner_tools.c
index 4159074df9c..c7b46cd640b 100644
--- a/source/blender/editors/space_outliner/outliner_tools.c
+++ b/source/blender/editors/space_outliner/outliner_tools.c
@@ -598,8 +598,14 @@ static uiBlock *merged_element_search_menu(bContext *C, ARegion *region, void *d
   short menu_width = 10 * UI_UNIT_X;
   but = uiDefSearchBut(
       block, search, 0, ICON_VIEWZOOM, sizeof(search), 10, 10, menu_width, UI_UNIT_Y, 0, 0, "");
-  UI_but_func_search_set(
-      but, NULL, merged_element_search_update_fn, data, NULL, merged_element_search_exec_fn, NULL);
+  UI_but_func_search_set(but,
+                         NULL,
+                         merged_element_search_update_fn,
+                         data,
+                         NULL,
+                         false,
+                         merged_element_search_exec_fn,
+                         NULL);
   UI_but_flag_enable(but, UI_BUT_ACTIVATE_ON_INIT);
 
   /* Fake button to hold space for search items */