1 files changed, 66 insertions, 15 deletions

diff --git a/source/blender/editors/armature/editarmature_undo.c b/source/blender/editors/armature/editarmature_undo.c
index 217de06d99b..f6f97af32b9 100644
--- a/source/blender/editors/armature/editarmature_undo.c
+++ b/source/blender/editors/armature/editarmature_undo.c
@@ -27,26 +27,34 @@
  *  \ingroup edarmature
  */
 
+#include "MEM_guardedalloc.h"
+
+
+#include "CLG_log.h"
+
 #include "DNA_armature_types.h"
 #include "DNA_object_types.h"
 
-#include "MEM_guardedalloc.h"
-
 #include "BLI_math.h"
 #include "BLI_array_utils.h"
 
 #include "BKE_context.h"
+#include "BKE_layer.h"
 #include "BKE_undo_system.h"
 
 #include "DEG_depsgraph.h"
 
 #include "ED_armature.h"
 #include "ED_object.h"
+#include "ED_undo.h"
 #include "ED_util.h"
 
 #include "WM_types.h"
 #include "WM_api.h"
 
+/** We only need this locally. */
+static CLG_LogRef LOG = {"ed.undo.armature"};
+
 /* -------------------------------------------------------------------- */
 /** \name Undo Conversion
  * \{ */
@@ -121,13 +129,20 @@ static Object *editarm_object_from_context(bContext *C)
 
 /* -------------------------------------------------------------------- */
 /** \name Implements ED Undo System
+ *
+ * \note This is similar for all edit-mode types.
  * \{ */
 
-typedef struct ArmatureUndoStep {
-	UndoStep step;
-	/* note: will split out into list for multi-object-editmode. */
+typedef struct ArmatureUndoStep_Elem {
+	struct ArmatureUndoStep_Elem *next, *prev;
 	UndoRefID_Object obedit_ref;
 	UndoArmature data;
+} ArmatureUndoStep_Elem;
+
+typedef struct ArmatureUndoStep {
+	UndoStep step;
+	ArmatureUndoStep_Elem *elems;
+	uint                   elems_len;
 } ArmatureUndoStep;
 
 static bool armature_undosys_poll(bContext *C)
@@ -138,10 +153,24 @@ static bool armature_undosys_poll(bContext *C)
 static bool armature_undosys_step_encode(struct bContext *C, UndoStep *us_p)
 {
 	ArmatureUndoStep *us = (ArmatureUndoStep *)us_p;
-	us->obedit_ref.ptr = editarm_object_from_context(C);
-	bArmature *arm = us->obedit_ref.ptr->data;
-	undoarm_from_editarm(&us->data, arm);
-	us->step.data_size = us->data.undo_size;
+
+	ViewLayer *view_layer = CTX_data_view_layer(C);
+	uint objects_len = 0;
+	Object **objects = BKE_view_layer_array_from_objects_in_edit_mode_unique_data(view_layer, &objects_len);
+
+	us->elems = MEM_callocN(sizeof(*us->elems) * objects_len, __func__);
+	us->elems_len = objects_len;
+
+	for (uint i = 0; i < objects_len; i++) {
+		Object *ob = objects[i];
+		ArmatureUndoStep_Elem *elem = &us->elems[i];
+
+		elem->obedit_ref.ptr = ob;
+		bArmature *arm = elem->obedit_ref.ptr->data;
+		undoarm_from_editarm(&elem->data, arm);
+		us->step.data_size += elem->data.undo_size;
+	}
+	MEM_freeN(objects);
 	return true;
 }
 
@@ -152,24 +181,46 @@ static void armature_undosys_step_decode(struct bContext *C, UndoStep *us_p, int
 	BLI_assert(armature_undosys_poll(C));
 
 	ArmatureUndoStep *us = (ArmatureUndoStep *)us_p;
-	Object *obedit = us->obedit_ref.ptr;
-	bArmature *arm = obedit->data;
-	undoarm_to_editarm(&us->data, arm);
-	DEG_id_tag_update(&obedit->id, OB_RECALC_DATA);
+
+	for (uint i = 0; i < us->elems_len; i++) {
+		ArmatureUndoStep_Elem *elem = &us->elems[i];
+		Object *obedit = elem->obedit_ref.ptr;
+		bArmature *arm = obedit->data;
+		if (arm->edbo == NULL) {
+			/* Should never fail, may not crash but can give odd behavior. */
+			CLOG_ERROR(&LOG, "name='%s', failed to enter edit-mode for object '%s', undo state invalid", us_p->name, obedit->id.name);
+			continue;
+		}
+		undoarm_to_editarm(&elem->data, arm);
+		DEG_id_tag_update(&obedit->id, OB_RECALC_DATA);
+	}
+
+	/* The first element is always active */
+	ED_undo_object_set_active_or_warn(CTX_data_view_layer(C), us->elems[0].obedit_ref.ptr, us_p->name, &LOG);
+
 	WM_event_add_notifier(C, NC_GEOM | ND_DATA, NULL);
 }
 
 static void armature_undosys_step_free(UndoStep *us_p)
 {
 	ArmatureUndoStep *us = (ArmatureUndoStep *)us_p;
-	undoarm_free_data(&us->data);
+
+	for (uint i = 0; i < us->elems_len; i++) {
+		ArmatureUndoStep_Elem *elem = &us->elems[i];
+		undoarm_free_data(&elem->data);
+	}
+	MEM_freeN(us->elems);
 }
 
 static void armature_undosys_foreach_ID_ref(
         UndoStep *us_p, UndoTypeForEachIDRefFn foreach_ID_ref_fn, void *user_data)
 {
 	ArmatureUndoStep *us = (ArmatureUndoStep *)us_p;
-	foreach_ID_ref_fn(user_data, ((UndoRefID *)&us->obedit_ref));
+
+	for (uint i = 0; i < us->elems_len; i++) {
+		ArmatureUndoStep_Elem *elem = &us->elems[i];
+		foreach_ID_ref_fn(user_data, ((UndoRefID *)&elem->obedit_ref));
+	}
 }
 
 /* Export for ED_undo_sys. */