Geometry Nodes: Enable exposing object and collection sockets

This patch allows connecting wires for object and collection socket
types to the "Group Input" node, which exposes them to be adjusted
in the modifier.

Thanks to @angavrilov's recent work in rB8964c02348f6, it is now
possible to edit pointer IDProperties in the interface when they are
drawn with `uiItemPointerR`.

This patch is composed of a few changes:
  - Add code to create pointer properties in the modifier settings for
    object and collection sockets, and also to draw them in the UI.
  - Also search through the modifier's `IDProperty` settings to find IDs
    used by the modifier.
  - Change the setting's UI layout to support the change.

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

1 files changed, 125 insertions, 37 deletions

diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index b8729a51df8..3ba88c0e18f 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -45,6 +45,7 @@
 #include "DNA_screen_types.h"
 
 #include "BKE_customdata.h"
+#include "BKE_global.h"
 #include "BKE_idprop.h"
 #include "BKE_lib_query.h"
 #include "BKE_mesh.h"
@@ -81,6 +82,9 @@ using blender::Set;
 using blender::Span;
 using blender::StringRef;
 using blender::Vector;
+using blender::bke::PersistentCollectionHandle;
+using blender::bke::PersistentDataHandleMap;
+using blender::bke::PersistentObjectHandle;
 using blender::fn::GMutablePointer;
 using blender::fn::GValueMap;
 using blender::nodes::GeoNodeExecParams;
@@ -114,7 +118,7 @@ static void addIdsUsedBySocket(const ListBase *sockets, Set<ID *> &ids)
   }
 }
 
-static void findUsedIds(const bNodeTree &tree, Set<ID *> &ids)
+static void find_used_ids_from_nodes(const bNodeTree &tree, Set<ID *> &ids)
 {
   Set<const bNodeTree *> handled_groups;
 
@@ -125,12 +129,27 @@ static void findUsedIds(const bNodeTree &tree, Set<ID *> &ids)
     if (node->type == NODE_GROUP) {
       const bNodeTree *group = (bNodeTree *)node->id;
       if (group != nullptr && handled_groups.add(group)) {
-        findUsedIds(*group, ids);
+        find_used_ids_from_nodes(*group, ids);
       }
     }
   }
 }
 
+static void find_used_ids_from_settings(const NodesModifierSettings &settings, Set<ID *> &ids)
+{
+  IDP_foreach_property(
+      settings.properties,
+      IDP_TYPE_FILTER_ID,
+      [](IDProperty *property, void *user_data) {
+        Set<ID *> *ids = (Set<ID *> *)user_data;
+        ID *id = IDP_Id(property);
+        if (id != nullptr) {
+          ids->add(id);
+        }
+      },
+      &ids);
+}
+
 static void updateDepsgraph(ModifierData *md, const ModifierUpdateDepsgraphContext *ctx)
 {
   NodesModifierData *nmd = reinterpret_cast<NodesModifierData *>(md);
@@ -139,7 +158,8 @@ static void updateDepsgraph(ModifierData *md, const ModifierUpdateDepsgraphConte
     DEG_add_node_tree_relation(ctx->node, nmd->node_group, "Nodes Modifier");
 
     Set<ID *> used_ids;
-    findUsedIds(*nmd->node_group, used_ids);
+    find_used_ids_from_settings(nmd->settings, used_ids);
+    find_used_ids_from_nodes(*nmd->node_group, used_ids);
     for (ID *id : used_ids) {
       if (GS(id->name) == ID_OB) {
         Object *object = reinterpret_cast<Object *>(id);
@@ -154,7 +174,6 @@ static void updateDepsgraph(ModifierData *md, const ModifierUpdateDepsgraphConte
     }
   }
 
-  /* TODO: Add relations for IDs in settings. */
   /* TODO: Add dependency for collection changes. */
 }
 
@@ -465,7 +484,9 @@ struct SocketPropertyType {
   IDProperty *(*create_default_ui_prop)(const bNodeSocket &socket, const char *name);
   PropertyType (*rna_subtype_get)(const bNodeSocket &socket);
   bool (*is_correct_type)(const IDProperty &property);
-  void (*init_cpp_value)(const IDProperty &property, void *r_value);
+  void (*init_cpp_value)(const IDProperty &property,
+                         const PersistentDataHandleMap &handles,
+                         void *r_value);
 };
 
 static IDProperty *socket_add_property(IDProperty *settings_prop_group,
@@ -550,9 +571,9 @@ static const SocketPropertyType *get_socket_property_type(const bNodeSocket &bso
             return (PropertyType)((bNodeSocketValueFloat *)socket.default_value)->subtype;
           },
           [](const IDProperty &property) { return property.type == IDP_FLOAT; },
-          [](const IDProperty &property, void *r_value) {
-            *(float *)r_value = IDP_Float(&property);
-          },
+          [](const IDProperty &property,
+             const PersistentDataHandleMap &UNUSED(handles),
+             void *r_value) { *(float *)r_value = IDP_Float(&property); },
       };
       return &float_type;
     }
@@ -586,7 +607,9 @@ static const SocketPropertyType *get_socket_property_type(const bNodeSocket &bso
             return (PropertyType)((bNodeSocketValueInt *)socket.default_value)->subtype;
           },
           [](const IDProperty &property) { return property.type == IDP_INT; },
-          [](const IDProperty &property, void *r_value) { *(int *)r_value = IDP_Int(&property); },
+          [](const IDProperty &property,
+             const PersistentDataHandleMap &UNUSED(handles),
+             void *r_value) { *(int *)r_value = IDP_Int(&property); },
       };
       return &int_type;
     }
@@ -629,9 +652,9 @@ static const SocketPropertyType *get_socket_property_type(const bNodeSocket &bso
             return property.type == IDP_ARRAY && property.subtype == IDP_FLOAT &&
                    property.len == 3;
           },
-          [](const IDProperty &property, void *r_value) {
-            copy_v3_v3((float *)r_value, (const float *)IDP_Array(&property));
-          },
+          [](const IDProperty &property,
+             const PersistentDataHandleMap &UNUSED(handles),
+             void *r_value) { copy_v3_v3((float *)r_value, (const float *)IDP_Array(&property)); },
       };
       return &vector_type;
     }
@@ -661,9 +684,9 @@ static const SocketPropertyType *get_socket_property_type(const bNodeSocket &bso
           },
           nullptr,
           [](const IDProperty &property) { return property.type == IDP_INT; },
-          [](const IDProperty &property, void *r_value) {
-            *(bool *)r_value = IDP_Int(&property) != 0;
-          },
+          [](const IDProperty &property,
+             const PersistentDataHandleMap &UNUSED(handles),
+             void *r_value) { *(bool *)r_value = IDP_Int(&property) != 0; },
       };
       return &boolean_type;
     }
@@ -683,12 +706,54 @@ static const SocketPropertyType *get_socket_property_type(const bNodeSocket &bso
           },
           nullptr,
           [](const IDProperty &property) { return property.type == IDP_STRING; },
-          [](const IDProperty &property, void *r_value) {
-            new (r_value) std::string(IDP_String(&property));
-          },
+          [](const IDProperty &property,
+             const PersistentDataHandleMap &UNUSED(handles),
+             void *r_value) { new (r_value) std::string(IDP_String(&property)); },
       };
       return &string_type;
     }
+    case SOCK_OBJECT: {
+      static const SocketPropertyType object_type = {
+          [](const bNodeSocket &socket, const char *name) {
+            bNodeSocketValueObject *value = (bNodeSocketValueObject *)socket.default_value;
+            IDPropertyTemplate idprop = {0};
+            idprop.id = (ID *)value->value;
+            return IDP_New(IDP_ID, &idprop, name);
+          },
+          nullptr,
+          nullptr,
+          nullptr,
+          nullptr,
+          [](const IDProperty &property) { return property.type == IDP_ID; },
+          [](const IDProperty &property, const PersistentDataHandleMap &handles, void *r_value) {
+            ID *id = IDP_Id(&property);
+            Object *object = (id && GS(id->name) == ID_OB) ? (Object *)id : nullptr;
+            new (r_value) PersistentObjectHandle(handles.lookup(object));
+          },
+      };
+      return &object_type;
+    }
+    case SOCK_COLLECTION: {
+      static const SocketPropertyType collection_type = {
+          [](const bNodeSocket &socket, const char *name) {
+            bNodeSocketValueCollection *value = (bNodeSocketValueCollection *)socket.default_value;
+            IDPropertyTemplate idprop = {0};
+            idprop.id = (ID *)value->value;
+            return IDP_New(IDP_ID, &idprop, name);
+          },
+          nullptr,
+          nullptr,
+          nullptr,
+          nullptr,
+          [](const IDProperty &property) { return property.type == IDP_ID; },
+          [](const IDProperty &property, const PersistentDataHandleMap &handles, void *r_value) {
+            ID *id = IDP_Id(&property);
+            Collection *collection = (id && GS(id->name) == ID_GR) ? (Collection *)id : nullptr;
+            new (r_value) PersistentCollectionHandle(handles.lookup(collection));
+          },
+      };
+      return &collection_type;
+    }
     default: {
       return nullptr;
     }
@@ -772,6 +837,7 @@ void MOD_nodes_init(Main *bmain, NodesModifierData *nmd)
 }
 
 static void initialize_group_input(NodesModifierData &nmd,
+                                   const PersistentDataHandleMap &handle_map,
                                    const bNodeSocket &socket,
                                    const CPPType &cpp_type,
                                    void *r_value)
@@ -793,15 +859,18 @@ static void initialize_group_input(NodesModifierData &nmd,
   }
   if (!property_type->is_correct_type(*property)) {
     blender::nodes::socket_cpp_value_get(socket, r_value);
+    return;
   }
-  property_type->init_cpp_value(*property, r_value);
+  property_type->init_cpp_value(*property, handle_map, r_value);
 }
 
-static void fill_data_handle_map(const DerivedNodeTree &tree,
+static void fill_data_handle_map(const NodesModifierSettings &settings,
+                                 const DerivedNodeTree &tree,
                                  blender::bke::PersistentDataHandleMap &handle_map)
 {
   Set<ID *> used_ids;
-  findUsedIds(*tree.btree(), used_ids);
+  find_used_ids_from_settings(settings, used_ids);
+  find_used_ids_from_nodes(*tree.btree(), used_ids);
 
   int current_handle = 0;
   for (ID *id : used_ids) {
@@ -826,6 +895,9 @@ static GeometrySet compute_geometry(const DerivedNodeTree &tree,
   blender::LinearAllocator<> &allocator = resources.linear_allocator();
   blender::nodes::MultiFunctionByNode mf_by_node = get_multi_function_per_node(tree, resources);
 
+  PersistentDataHandleMap handle_map;
+  fill_data_handle_map(nmd->settings, tree, handle_map);
+
   Map<const DOutputSocket *, GMutablePointer> group_inputs;
 
   if (group_input_sockets.size() > 0) {
@@ -845,7 +917,7 @@ static GeometrySet compute_geometry(const DerivedNodeTree &tree,
     for (const DOutputSocket *socket : remaining_input_sockets) {
       const CPPType &cpp_type = *blender::nodes::socket_cpp_type_get(*socket->typeinfo());
       void *value_in = allocator.allocate(cpp_type.size(), cpp_type.alignment());
-      initialize_group_input(*nmd, *socket->bsocket(), cpp_type, value_in);
+      initialize_group_input(*nmd, handle_map, *socket->bsocket(), cpp_type, value_in);
       group_inputs.add_new(socket, {cpp_type, value_in});
     }
   }
@@ -853,9 +925,6 @@ static GeometrySet compute_geometry(const DerivedNodeTree &tree,
   Vector<const DInputSocket *> group_outputs;
   group_outputs.append(&socket_to_compute);
 
-  blender::bke::PersistentDataHandleMap handle_map;
-  fill_data_handle_map(tree, handle_map);
-
   GeometryNodesEvaluator evaluator{
       group_inputs, group_outputs, mf_by_node, handle_map, ctx->object};
   Vector<GMutablePointer> results = evaluator.execute();
@@ -883,18 +952,9 @@ static void check_property_socket_sync(const Object *ob, ModifierData *md)
 
     IDProperty *property = IDP_GetPropertyFromGroup(nmd->settings.properties, socket->identifier);
     if (property == nullptr) {
-      if (socket->type == SOCK_STRING) {
-        BKE_modifier_set_error(ob, md, "String socket can not be exposed in the modifier");
-      }
-      else if (socket->type == SOCK_OBJECT) {
-        BKE_modifier_set_error(ob, md, "Object socket can not be exposed in the modifier");
-      }
-      else if (socket->type == SOCK_GEOMETRY) {
+      if (socket->type == SOCK_GEOMETRY) {
         BKE_modifier_set_error(ob, md, "Node group can only have one geometry input");
       }
-      else if (socket->type == SOCK_COLLECTION) {
-        BKE_modifier_set_error(ob, md, "Collection socket can not be exposed in the modifier");
-      }
       else {
         BKE_modifier_set_error(ob, md, "Missing property for input socket \"%s\"", socket->name);
       }
@@ -992,6 +1052,7 @@ static void modifyGeometrySet(ModifierData *md,
  * the node socket identifier for the property names, since they are unique, but also having
  * the correct label displayed in the UI.  */
 static void draw_property_for_socket(uiLayout *layout,
+                                     PointerRNA *bmain_ptr,
                                      PointerRNA *settings_ptr,
                                      const IDProperty *modifier_props,
                                      const bNodeSocket &socket)
@@ -1013,13 +1074,36 @@ static void draw_property_for_socket(uiLayout *layout,
 
     char rna_path[sizeof(socket_id_esc) + 4];
     BLI_snprintf(rna_path, ARRAY_SIZE(rna_path), "[\"%s\"]", socket_id_esc);
-    uiItemR(layout, settings_ptr, rna_path, 0, socket.name, ICON_NONE);
+
+    /* Use #uiItemPointerR to draw pointer properties because #uiItemR would not have enough
+     * information about what type of ID to select for editing the values. This is because
+     * pointer IDProperties contain no information about their type. */
+    switch (socket.type) {
+      case SOCK_OBJECT: {
+        uiItemPointerR(
+            layout, settings_ptr, rna_path, bmain_ptr, "objects", socket.name, ICON_OBJECT_DATA);
+        break;
+      }
+      case SOCK_COLLECTION: {
+        uiItemPointerR(layout,
+                       settings_ptr,
+                       rna_path,
+                       bmain_ptr,
+                       "collections",
+                       socket.name,
+                       ICON_OUTLINER_COLLECTION);
+        break;
+      }
+      default:
+        uiItemR(layout, settings_ptr, rna_path, 0, socket.name, ICON_NONE);
+    }
   }
 }
 
 static void panel_draw(const bContext *C, Panel *panel)
 {
   uiLayout *layout = panel->layout;
+  Main *bmain = CTX_data_main(C);
 
   PointerRNA *ptr = modifier_panel_get_property_pointers(panel, nullptr);
   NodesModifierData *nmd = static_cast<NodesModifierData *>(ptr->data);
@@ -1044,8 +1128,12 @@ static void panel_draw(const bContext *C, Panel *panel)
     PointerRNA settings_ptr;
     RNA_pointer_create(ptr->owner_id, &RNA_NodesModifierSettings, &nmd->settings, &settings_ptr);
 
+    PointerRNA bmain_ptr;
+    RNA_main_pointer_create(bmain, &bmain_ptr);
+
     LISTBASE_FOREACH (bNodeSocket *, socket, &nmd->node_group->inputs) {
-      draw_property_for_socket(layout, &settings_ptr, nmd->settings.properties, *socket);
+      draw_property_for_socket(
+          layout, &bmain_ptr, &settings_ptr, nmd->settings.properties, *socket);
     }
   }