From 0b36b71b860bcb69465ae470e0e926e85e25e091 Mon Sep 17 00:00:00 2001
From: RustyStriker <avivr903@tutanota.com>
Date: Wed, 7 Sep 2022 23:47:55 +0300
Subject: [PATCH] fix image deletion makes other undos impossible

---
 src/ui.rs   |  2 +-
 src/undo.rs | 24 ++++++++++++++++++++----
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/ui.rs b/src/ui.rs
index eb3db37..4271d62 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -418,7 +418,7 @@ pub fn inspector_sys(
                         std::mem::swap(&mut nid, &mut  *id);
                         let transform = transforms.get(e).map(|t| *t).unwrap_or_default();
                         let handle = image_handles.get(e).unwrap().clone();
-                        undo.push(UndoItem::DeleteImage { id: nid, transform: transform, handle });
+                        undo.push(UndoItem::DeleteImage { id: nid, transform: transform, handle, old_entity: e });
                     }
                     coms.entity(e).despawn_recursive();
                     **selected = None;
diff --git a/src/undo.rs b/src/undo.rs
index f299a09..85a7445 100644
--- a/src/undo.rs
+++ b/src/undo.rs
@@ -16,7 +16,7 @@ pub enum UndoItem {
     ChangeZ { entity: Entity, from: f32, to: f32 },
     ReScale { entity: Entity, from: Vec2, to: Vec2 },
     CreateImage { entity: Entity },
-    DeleteImage { id: ImageData, transform: Transform, handle: Handle<Image> },
+    DeleteImage { id: ImageData, transform: Transform, handle: Handle<Image>, old_entity: Entity },
 }
 pub struct UndoStack {
     items: Vec<UndoItem>,
@@ -170,11 +170,11 @@ pub fn undo_redo_sys(
                         let mut nid = ImageData::default();
                         std::mem::swap(&mut *id, &mut nid);
 
-                        coms.entity(*entity).despawn();
-                        *item = UndoItem::DeleteImage { id: nid, transform, handle: handle.clone() };
+                        coms.entity(*entity).despawn_recursive();
+                        *item = UndoItem::DeleteImage { id: nid, transform, handle: handle.clone(), old_entity: *entity };
                     }
                 },
-                UndoItem::DeleteImage { transform, handle, id } => {
+                UndoItem::DeleteImage { transform, handle, id, old_entity } => {
                     let mut nid = ImageData::default();
                     std::mem::swap(id, &mut nid);
 
@@ -185,7 +185,23 @@ pub fn undo_redo_sys(
                     })
                     .insert(nid).id();
 
+                    let oe = *old_entity;
                     *item = UndoItem::CreateImage { entity: e };
+                    
+                    // Update each item with the same entity id
+                    let f = |ne: &mut Entity| { if *ne == oe { *ne = e; }};
+                    for i in stack.items.iter_mut() {
+                        match i {
+                            UndoItem::CreateShape { entity } => f(entity),
+                            UndoItem::MoveItem { shape, .. } => f(shape),
+                            UndoItem::Rotate { entity, .. } => f(entity),
+                            UndoItem::NameChange { entity, .. } => f(entity),
+                            UndoItem::NoteChange { entity, .. } => f(entity),
+                            UndoItem::ChangeZ { entity, .. } => f(entity),
+                            UndoItem::ReScale { entity, .. } => f(entity),
+                            _ => {},
+                        }
+                    }
                 },
                 UndoItem::NoteChange { entity, from } => {
                     if let Ok(mut sd) = shapes.get_mut(*entity) {