From 97cd919a7b14d7555c3144db65f80406dc5a1d96 Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Mon, 5 Jul 2021 15:59:26 +1000
Subject: [PATCH] Improve undo redo object linking robustness. See #1475.

---
 src/blenderbim/blenderbim/bim/__init__.py |  2 ++
 src/blenderbim/blenderbim/bim/handler.py  | 16 ++++++++++--
 src/blenderbim/blenderbim/bim/ifc.py      | 31 ++++++++++++++++++-----
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/blenderbim/blenderbim/bim/__init__.py b/src/blenderbim/blenderbim/bim/__init__.py
index 952e46b1a7..d3433934b1 100644
--- a/src/blenderbim/blenderbim/bim/__init__.py
+++ b/src/blenderbim/blenderbim/bim/__init__.py
@@ -110,7 +110,9 @@ if bpy is not None:
         for cls in classes:
             bpy.utils.register_class(cls)
         bpy.app.handlers.depsgraph_update_post.append(on_register)
+        bpy.app.handlers.undo_pre.append(handler.undo_pre)
         bpy.app.handlers.undo_post.append(handler.undo_post)
+        bpy.app.handlers.redo_pre.append(handler.redo_pre)
         bpy.app.handlers.redo_post.append(handler.redo_post)
         bpy.app.handlers.load_post.append(handler.setDefaultProperties)
         bpy.app.handlers.load_post.append(handler.loadIfcStore)
diff --git a/src/blenderbim/blenderbim/bim/handler.py b/src/blenderbim/blenderbim/bim/handler.py
index 32723669e0..f4d9945518 100644
--- a/src/blenderbim/blenderbim/bim/handler.py
+++ b/src/blenderbim/blenderbim/bim/handler.py
@@ -114,12 +114,23 @@ def loadIfcStore(scene):
     purge_module_data()
 
 
+@persistent
+def undo_pre(scene):
+    IfcStore.update_undo_redo_stack_objects()
+
+
 @persistent
 def undo_post(scene):
     if IfcStore.last_transaction != bpy.context.scene.BIMProperties.last_transaction:
         IfcStore.last_transaction = bpy.context.scene.BIMProperties.last_transaction
         IfcStore.undo()
-    IfcStore.reload_linked_elements(should_reload_selected=True)
+    IfcStore.update_undo_redo_stack_objects()
+    IfcStore.reload_linked_elements(objects=[bpy.data.objects.get(o) for o in IfcStore.undo_redo_stack_objects])
+
+
+@persistent
+def redo_pre(scene):
+    IfcStore.update_undo_redo_stack_objects()
 
 
 @persistent
@@ -127,7 +138,8 @@ def redo_post(scene):
     if IfcStore.last_transaction != bpy.context.scene.BIMProperties.last_transaction:
         IfcStore.last_transaction = bpy.context.scene.BIMProperties.last_transaction
         IfcStore.redo()
-    IfcStore.reload_linked_elements(should_reload_selected=True)
+    IfcStore.update_undo_redo_stack_objects()
+    IfcStore.reload_linked_elements(objects=[bpy.data.objects.get(o) for o in IfcStore.undo_redo_stack_objects])
 
 
 @persistent
diff --git a/src/blenderbim/blenderbim/bim/ifc.py b/src/blenderbim/blenderbim/bim/ifc.py
index acc3b71817..2da38e04d5 100644
--- a/src/blenderbim/blenderbim/bim/ifc.py
+++ b/src/blenderbim/blenderbim/bim/ifc.py
@@ -16,6 +16,7 @@ class IfcStore:
     library_path = ""
     library_file = None
     element_listeners = set()
+    undo_redo_stack_objects = set()
     current_transaction = ""
     last_transaction = ""
     history = []
@@ -70,19 +71,36 @@ class IfcStore:
     def add_element_listener(callback):
         IfcStore.element_listeners.add(callback)
 
+    """Keeps track of selected object names, typically during undo and redo
+
+    When any Blender object is stored outside a Blender PointerProperty, such as
+    in a regular Python list, there is the likely probability that the object
+    will be invalidated when undo or redo occurs. Object invalidation seems to
+    only occur for selected objects either pre/post undo/redo event, including
+    selected objects for consecutive undo/redos.
+
+    So if I first select o1, then o2, then o3, then press undo, o3 will be
+    invalidated. If instead I press undo twice, o3 and o2 will be invalidated.
+    """
     @staticmethod
-    def reload_linked_elements(should_reload_selected=False):
+    def update_undo_redo_stack_objects():
+        if bpy.context.active_object:
+            objects = set([o.name for o in bpy.context.selected_objects + [bpy.context.active_object]])
+        else:
+            objects = set([o.name for o in bpy.context.selected_objects])
+        IfcStore.undo_redo_stack_objects |= objects
+
+    @staticmethod
+    def reload_linked_elements(objects=None):
         file = IfcStore.get_file()
         if not file:
             return
-        if should_reload_selected:
-            objects = bpy.context.selected_objects
-            if bpy.context.active_object:
-                objects += [bpy.context.active_object]
-        else:
+        if objects is None:
             objects = bpy.data.objects
 
         for obj in objects:
+            if not obj:
+                continue
             if not obj.BIMObjectProperties.ifc_definition_id:
                 continue
             element = file.by_id(obj.BIMObjectProperties.ifc_definition_id)
@@ -183,6 +201,7 @@ class IfcStore:
 
     @staticmethod
     def begin_transaction(operator):
+        IfcStore.undo_redo_stack_objects = set()
         IfcStore.current_transaction = str(uuid.uuid4())
         operator.transaction_key = IfcStore.current_transaction