1 files changed, 20 insertions, 24 deletions

diff --git a/source/blender/editors/undo/ed_undo.c b/source/blender/editors/undo/ed_undo.c
index 7d8f72f3779..989863b34af 100644
--- a/source/blender/editors/undo/ed_undo.c
+++ b/source/blender/editors/undo/ed_undo.c
@@ -911,9 +911,7 @@ void ED_undo_object_editmode_restore_helper(struct bContext *C,
  * and local collections may be used.
  * \{ */
 
-static int undo_editmode_objects_from_view_layer_prepare(ViewLayer *view_layer,
-                                                         Object *obact,
-                                                         int *r_active_index)
+static int undo_editmode_objects_from_view_layer_prepare(ViewLayer *view_layer, Object *obact)
 {
   const short object_type = obact->type;
 
@@ -929,9 +927,6 @@ static int undo_editmode_objects_from_view_layer_prepare(ViewLayer *view_layer,
   LISTBASE_FOREACH (Base *, base, &view_layer->object_bases) {
     Object *ob = base->object;
     if ((ob->type == object_type) && (ob->mode & OB_MODE_EDIT)) {
-      if (ob == obact) {
-        *r_active_index = len;
-      }
       ID *id = ob->data;
       if ((id->tag & LIB_TAG_DOIT) == 0) {
         len += 1;
@@ -944,16 +939,18 @@ static int undo_editmode_objects_from_view_layer_prepare(ViewLayer *view_layer,
 
 Object **ED_undo_editmode_objects_from_view_layer(ViewLayer *view_layer, uint *r_len)
 {
-  Object *obact = OBACT(view_layer);
-  if ((obact == NULL) || (obact->mode & OB_MODE_EDIT) == 0) {
+  Base *baseact = BASACT(view_layer);
+  if ((baseact == NULL) || (baseact->object->mode & OB_MODE_EDIT) == 0) {
     return MEM_mallocN(0, __func__);
   }
-  int active_index = 0;
-  const int len = undo_editmode_objects_from_view_layer_prepare(view_layer, obact, &active_index);
-  const short object_type = obact->type;
+  const int len = undo_editmode_objects_from_view_layer_prepare(view_layer, baseact->object);
+  const short object_type = baseact->object->type;
   int i = 0;
   Object **objects = MEM_malloc_arrayN(len, sizeof(*objects), __func__);
-  LISTBASE_FOREACH (Base *, base, &view_layer->object_bases) {
+  /* Base iteration, starting with the active-base to ensure it's the first item in the array.
+   * Looping over the active-base twice is OK as the tag check prevents it being handled twice. */
+  for (Base *base = baseact, *base_next = FIRSTBASE(view_layer); base;
+       base = base_next, base_next = base_next ? base_next->next : NULL) {
     Object *ob = base->object;
     if ((ob->type == object_type) && (ob->mode & OB_MODE_EDIT)) {
       ID *id = ob->data;
@@ -964,25 +961,25 @@ Object **ED_undo_editmode_objects_from_view_layer(ViewLayer *view_layer, uint *r
     }
   }
   BLI_assert(i == len);
-  if (active_index > 0) {
-    SWAP(Object *, objects[0], objects[active_index]);
-  }
+  BLI_assert(objects[0] == baseact->object);
   *r_len = len;
   return objects;
 }
 
 Base **ED_undo_editmode_bases_from_view_layer(ViewLayer *view_layer, uint *r_len)
 {
-  Object *obact = OBACT(view_layer);
-  if ((obact == NULL) || (obact->mode & OB_MODE_EDIT) == 0) {
+  Base *baseact = BASACT(view_layer);
+  if ((baseact == NULL) || (baseact->object->mode & OB_MODE_EDIT) == 0) {
     return MEM_mallocN(0, __func__);
   }
-  int active_index = 0;
-  const int len = undo_editmode_objects_from_view_layer_prepare(view_layer, obact, &active_index);
-  const short object_type = obact->type;
+  const int len = undo_editmode_objects_from_view_layer_prepare(view_layer, baseact->object);
+  const short object_type = baseact->object->type;
   int i = 0;
   Base **base_array = MEM_malloc_arrayN(len, sizeof(*base_array), __func__);
-  LISTBASE_FOREACH (Base *, base, &view_layer->object_bases) {
+  /* Base iteration, starting with the active-base to ensure it's the first item in the array.
+   * Looping over the active-base twice is OK as the tag check prevents it being handled twice. */
+  for (Base *base = BASACT(view_layer), *base_next = FIRSTBASE(view_layer); base;
+       base = base_next, base_next = base_next ? base_next->next : NULL) {
     Object *ob = base->object;
     if ((ob->type == object_type) && (ob->mode & OB_MODE_EDIT)) {
       ID *id = ob->data;
@@ -992,10 +989,9 @@ Base **ED_undo_editmode_bases_from_view_layer(ViewLayer *view_layer, uint *r_len
       }
     }
   }
+
   BLI_assert(i == len);
-  if (active_index > 0) {
-    SWAP(Base *, base_array[0], base_array[active_index]);
-  }
+  BLI_assert(base_array[0] == baseact);
   *r_len = len;
   return base_array;
 }