BKE_collection: Add a util returning a gset with all objects in given scene's collections.

This is internaly using the code of `BKE_scene_objects_iterator` and
steals its gset. More efficient than using that iterator directly to
rebuild another GSet...

2 files changed, 42 insertions, 4 deletions

diff --git a/source/blender/blenkernel/BKE_collection.h b/source/blender/blenkernel/BKE_collection.h
index 4fa285a4e85..3412be92a3a 100644
--- a/source/blender/blenkernel/BKE_collection.h
+++ b/source/blender/blenkernel/BKE_collection.h
@@ -227,6 +227,8 @@ void BKE_scene_objects_iterator_begin(struct BLI_Iterator *iter, void *data_in);
 void BKE_scene_objects_iterator_next(struct BLI_Iterator *iter);
 void BKE_scene_objects_iterator_end(struct BLI_Iterator *iter);
 
+struct GSet *BKE_scene_objects_as_gset(struct Scene *scene, struct GSet *objects_gset);
+
 #define FOREACH_SCENE_COLLECTION_BEGIN(scene, _instance) \
   ITER_BEGIN (BKE_scene_collections_iterator_begin, \
               BKE_scene_collections_iterator_next, \
diff --git a/source/blender/blenkernel/intern/collection.c b/source/blender/blenkernel/intern/collection.c
index 08ece42e6bb..426f8a44aff 100644
--- a/source/blender/blenkernel/intern/collection.c
+++ b/source/blender/blenkernel/intern/collection.c
@@ -2078,14 +2078,18 @@ typedef struct SceneObjectsIteratorData {
   BLI_Iterator scene_collection_iter;
 } SceneObjectsIteratorData;
 
-void BKE_scene_objects_iterator_begin(BLI_Iterator *iter, void *data_in)
+static void scene_objects_iterator_begin(BLI_Iterator *iter, Scene *scene, GSet *visited_objects)
 {
-  Scene *scene = data_in;
   SceneObjectsIteratorData *data = MEM_callocN(sizeof(SceneObjectsIteratorData), __func__);
   iter->data = data;
 
   /* Lookup list to make sure that each object is only processed once. */
-  data->visited = BLI_gset_ptr_new(__func__);
+  if (visited_objects != NULL) {
+    data->visited = visited_objects;
+  }
+  else {
+    data->visited = BLI_gset_ptr_new(__func__);
+  }
 
   /* We wrap the scenecollection iterator here to go over the scene collections. */
   BKE_scene_collections_iterator_begin(&data->scene_collection_iter, scene);
@@ -2096,6 +2100,13 @@ void BKE_scene_objects_iterator_begin(BLI_Iterator *iter, void *data_in)
   BKE_scene_objects_iterator_next(iter);
 }
 
+void BKE_scene_objects_iterator_begin(BLI_Iterator *iter, void *data_in)
+{
+  Scene *scene = data_in;
+
+  scene_objects_iterator_begin(iter, scene, NULL);
+}
+
 /**
  * Ensures we only get each object once, even when included in several collections.
  */
@@ -2149,9 +2160,34 @@ void BKE_scene_objects_iterator_end(BLI_Iterator *iter)
   SceneObjectsIteratorData *data = iter->data;
   if (data) {
     BKE_scene_collections_iterator_end(&data->scene_collection_iter);
-    BLI_gset_free(data->visited, NULL);
+    if (data->visited != NULL) {
+      BLI_gset_free(data->visited, NULL);
+    }
     MEM_freeN(data);
   }
 }
 
+/** Generate a new GSet (or extend given `objects_gset` if not NULL) with all objects referenced by
+ * all collections of given `scene`.
+ *
+ * \note: This will include objects without a base currently (because they would belong to excluded
+ * collections only e.g.). */
+GSet *BKE_scene_objects_as_gset(Scene *scene, GSet *objects_gset)
+{
+  BLI_Iterator iter;
+  scene_objects_iterator_begin(&iter, scene, objects_gset);
+  while (iter.valid) {
+    BKE_scene_objects_iterator_next(&iter);
+  }
+
+  /* `return_gset` is either given `objects_gset` (if non-NULL), or the GSet allocated by the
+   * iterator. Either way, we want to get it back, and prevent `BKE_scene_objects_iterator_end`
+   * from freeing it. */
+  GSet *return_gset = ((SceneObjectsIteratorData *)iter.data)->visited;
+  ((SceneObjectsIteratorData *)iter.data)->visited = NULL;
+  BKE_scene_objects_iterator_end(&iter);
+
+  return return_gset;
+}
+
 /** \} */