Iterator to loop over objects based on a flag

3 files changed, 71 insertions, 4 deletions

diff --git a/source/blender/blenkernel/BKE_collection.h b/source/blender/blenkernel/BKE_collection.h
index 7f4360d2e9c..a3bbcc8687a 100644
--- a/source/blender/blenkernel/BKE_collection.h
+++ b/source/blender/blenkernel/BKE_collection.h
@@ -346,6 +346,20 @@ 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);
 
+/** Iterate over objects in the scene based on a flag.
+ *
+ * \note The object->flag is tested against flag.
+ * */
+typedef struct SceneObjectsIteratorExData {
+  struct Scene *scene;
+  int flag;
+  void *iter_data;
+} SceneObjectsIteratorExData;
+
+void BKE_scene_objects_iterator_begin_ex(struct BLI_Iterator *iter, void *data_in);
+void BKE_scene_objects_iterator_next_ex(struct BLI_Iterator *iter);
+void BKE_scene_objects_iterator_end_ex(struct BLI_Iterator *iter);
+
 /**
  * Generate a new #GSet (or extend given `objects_gset` if not NULL) with all objects referenced by
  * all collections of given `scene`.
diff --git a/source/blender/blenkernel/BKE_layer.h b/source/blender/blenkernel/BKE_layer.h
index 77a3223c064..c877035830c 100644
--- a/source/blender/blenkernel/BKE_layer.h
+++ b/source/blender/blenkernel/BKE_layer.h
@@ -436,15 +436,26 @@ void BKE_view_layer_visible_bases_iterator_end(BLI_Iterator *iter);
     IteratorBeginCb func_begin; \
     IteratorCb func_next, func_end; \
     void *data_in; \
-    struct ObjectsVisibleIteratorData data_ = {NULL}; \
-    data_.view_layer = _view_layer; \
-    data_.v3d = _v3d; \
+\
+    struct ObjectsVisibleIteratorData data_select_ = {NULL}; \
+    data_select_.view_layer = _view_layer; \
+    data_select_.v3d = _v3d; \
+\
+    struct SceneObjectsIteratorExData data_flag_ = {NULL}; \
+    data_flag_.scene = scene; \
+    data_flag_.flag = flag; \
 \
     if (flag == SELECT) { \
       func_begin = &BKE_view_layer_selected_objects_iterator_begin; \
       func_next = &BKE_view_layer_selected_objects_iterator_next; \
       func_end = &BKE_view_layer_selected_objects_iterator_end; \
-      data_in = &data_; \
+      data_in = &data_select_; \
+    } \
+    else if (flag != 0) { \
+      func_begin = BKE_scene_objects_iterator_begin_ex; \
+      func_next = BKE_scene_objects_iterator_next_ex; \
+      func_end = BKE_scene_objects_iterator_end_ex; \
+      data_in = &data_flag_; \
     } \
     else { \
       func_begin = BKE_scene_objects_iterator_begin; \
diff --git a/source/blender/blenkernel/intern/collection.c b/source/blender/blenkernel/intern/collection.c
index bdaea487cfb..c215321bc30 100644
--- a/source/blender/blenkernel/intern/collection.c
+++ b/source/blender/blenkernel/intern/collection.c
@@ -1929,6 +1929,48 @@ void BKE_scene_objects_iterator_begin(BLI_Iterator *iter, void *data_in)
   scene_objects_iterator_begin(iter, scene, NULL);
 }
 
+void BKE_scene_objects_iterator_begin_ex(BLI_Iterator *iter, void *data_in)
+{
+  SceneObjectsIteratorExData *data = data_in;
+
+  BKE_scene_objects_iterator_begin(iter, data->scene);
+
+  /* Pack the data. */
+  data->iter_data = iter->data;
+  iter->data = data_in;
+}
+
+void BKE_scene_objects_iterator_next_ex(struct BLI_Iterator *iter)
+{
+  /* Unpack the data. */
+  SceneObjectsIteratorExData *data = iter->data;
+  iter->data = data->iter_data;
+
+  BKE_scene_objects_iterator_next(iter);
+
+  Object *ob = iter->current;
+  if (ob && (ob->flag & data->flag) == 0) {
+    iter->skip = true;
+  }
+
+  /* Pack the data. */
+  data->iter_data = iter->data;
+  iter->data = data;
+}
+
+void BKE_scene_objects_iterator_end_ex(struct BLI_Iterator *iter)
+{
+  /* Unpack the data. */
+  SceneObjectsIteratorExData *data = iter->data;
+  iter->data = data->iter_data;
+
+  BKE_scene_objects_iterator_end(iter);
+
+  /* Pack the data. */
+  data->iter_data = iter->data;
+  iter->data = data;
+}
+
 /**
  * Ensures we only get each object once, even when included in several collections.
  */