Geometry Nodes: support collection sockets

Part of D9739.

3 files changed, 43 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_persistent_data_handle.hh b/source/blender/blenkernel/BKE_persistent_data_handle.hh
index 42a853886d7..bbee09c7bf4 100644
--- a/source/blender/blenkernel/BKE_persistent_data_handle.hh
+++ b/source/blender/blenkernel/BKE_persistent_data_handle.hh
@@ -27,6 +27,7 @@
 
 #include "DNA_ID.h"
 
+struct Collection;
 struct Object;
 
 namespace blender::bke {
@@ -82,6 +83,11 @@ class PersistentObjectHandle : public PersistentIDHandle {
   using PersistentIDHandle::PersistentIDHandle;
 };
 
+class PersistentCollectionHandle : public PersistentIDHandle {
+  friend PersistentDataHandleMap;
+  using PersistentIDHandle::PersistentIDHandle;
+};
+
 class PersistentDataHandleMap {
  private:
   Map<int32_t, ID *> id_by_handle_;
@@ -107,6 +113,12 @@ class PersistentDataHandleMap {
     return PersistentObjectHandle(handle);
   }
 
+  PersistentCollectionHandle lookup(Collection *collection) const
+  {
+    const int handle = handle_by_id_.lookup_default((ID *)collection, -1);
+    return PersistentCollectionHandle(handle);
+  }
+
   ID *lookup(const PersistentIDHandle &handle) const
   {
     ID *id = id_by_handle_.lookup_default(handle.handle_, nullptr);
@@ -124,6 +136,18 @@ class PersistentDataHandleMap {
     }
     return (Object *)id;
   }
+
+  Collection *lookup(const PersistentCollectionHandle &handle) const
+  {
+    ID *id = this->lookup((const PersistentIDHandle &)handle);
+    if (id == nullptr) {
+      return nullptr;
+    }
+    if (GS(id->name) != ID_GR) {
+      return nullptr;
+    }
+    return (Collection *)id;
+  }
 };
 
 }  // namespace blender::bke
diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index 29ed08f0367..513e16e1a89 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -104,6 +104,12 @@ static void addIdsUsedBySocket(const ListBase *sockets, Set<ID *> &ids)
         ids.add(&object->id);
       }
     }
+    else if (socket->type == SOCK_COLLECTION) {
+      Collection *collection = ((bNodeSocketValueCollection *)socket->default_value)->value;
+      if (collection != nullptr) {
+        ids.add(&collection->id);
+      }
+    }
   }
 }
 
@@ -399,6 +405,11 @@ class GeometryNodesEvaluator {
       blender::bke::PersistentObjectHandle object_handle = handle_map_.lookup(object);
       new (buffer) blender::bke::PersistentObjectHandle(object_handle);
     }
+    else if (bsocket->type == SOCK_COLLECTION) {
+      Collection *collection = ((bNodeSocketValueCollection *)bsocket->default_value)->value;
+      blender::bke::PersistentCollectionHandle collection_handle = handle_map_.lookup(collection);
+      new (buffer) blender::bke::PersistentCollectionHandle(collection_handle);
+    }
     else {
       blender::nodes::socket_cpp_value_get(*bsocket, buffer);
     }
@@ -850,6 +861,9 @@ static void check_property_socket_sync(const Object *ob, ModifierData *md)
       else 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);
       }
diff --git a/source/blender/nodes/intern/node_socket.cc b/source/blender/nodes/intern/node_socket.cc
index fda44b3d4f2..06dd4b8c188 100644
--- a/source/blender/nodes/intern/node_socket.cc
+++ b/source/blender/nodes/intern/node_socket.cc
@@ -664,6 +664,7 @@ class ObjectSocketMultiFunction : public blender::fn::MultiFunction {
 };
 
 MAKE_CPP_TYPE(PersistentObjectHandle, blender::bke::PersistentObjectHandle);
+MAKE_CPP_TYPE(PersistentCollectionHandle, blender::bke::PersistentCollectionHandle);
 
 static bNodeSocketType *make_socket_type_object()
 {
@@ -692,6 +693,10 @@ static bNodeSocketType *make_socket_type_geometry()
 static bNodeSocketType *make_socket_type_collection()
 {
   bNodeSocketType *socktype = make_standard_socket_type(SOCK_COLLECTION, PROP_NONE);
+  socktype->get_cpp_type = []() {
+    /* Objects are not passed along as raw pointers, but as handles. */
+    return &blender::fn::CPPType::get<blender::bke::PersistentCollectionHandle>();
+  };
   return socktype;
 }