From 40bd40c52ed794b20cc7772e0d8913d65744afd2 Mon Sep 17 00:00:00 2001
From: Jaex <Jaex@users.noreply.github.com>
Date: Wed, 27 Dec 2023 10:27:54 +0300
Subject: [PATCH] Code refactoring

---
 .../Helpers/ImageEditorHistory.cs             | 33 +++++++++----------
 .../Helpers/ImageEditorMemento.cs             |  2 +-
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/ShareX.ScreenCaptureLib/Helpers/ImageEditorHistory.cs b/ShareX.ScreenCaptureLib/Helpers/ImageEditorHistory.cs
index 35c86a0a1..9e0883959 100644
--- a/ShareX.ScreenCaptureLib/Helpers/ImageEditorHistory.cs
+++ b/ShareX.ScreenCaptureLib/Helpers/ImageEditorHistory.cs
@@ -35,9 +35,9 @@ internal class ImageEditorHistory : IDisposable
         public bool CanUndo => undoMementoStack.Count > 0;
         public bool CanRedo => redoMementoStack.Count > 0;
 
+        private readonly ShapeManager shapeManager;
         private Stack<ImageEditorMemento> undoMementoStack = new Stack<ImageEditorMemento>();
         private Stack<ImageEditorMemento> redoMementoStack = new Stack<ImageEditorMemento>();
-        private readonly ShapeManager shapeManager;
 
         public ImageEditorHistory(ShapeManager shapeManager)
         {
@@ -50,7 +50,7 @@ private void AddMemento(ImageEditorMemento memento)
 
             foreach (ImageEditorMemento redoMemento in redoMementoStack)
             {
-                redoMemento.Dispose();
+                redoMemento?.Dispose();
             }
 
             redoMementoStack.Clear();
@@ -85,24 +85,23 @@ public void Undo()
         {
             if (CanUndo)
             {
-                ImageEditorMemento redoMemento;
-                ImageEditorMemento memento = undoMementoStack.Pop();
+                ImageEditorMemento undoMemento = undoMementoStack.Pop();
 
-                if (memento.Shapes != null)
+                if (undoMemento.Shapes != null)
                 {
-                    if (memento.Canvas == null)
+                    if (undoMemento.Canvas == null)
                     {
-                        redoMemento = GetMementoFromShapes();
+                        ImageEditorMemento redoMemento = GetMementoFromShapes();
                         redoMementoStack.Push(redoMemento);
 
-                        shapeManager.RestoreState(memento);
+                        shapeManager.RestoreState(undoMemento);
                     }
                     else
                     {
-                        redoMemento = GetMementoFromCanvas();
+                        ImageEditorMemento redoMemento = GetMementoFromCanvas();
                         redoMementoStack.Push(redoMemento);
 
-                        shapeManager.RestoreState(memento);
+                        shapeManager.RestoreState(undoMemento);
                     }
                 }
             }
@@ -118,15 +117,15 @@ public void Redo()
                 {
                     if (redoMemento.Canvas == null)
                     {
-                        ImageEditorMemento memento = GetMementoFromShapes();
-                        undoMementoStack.Push(memento);
+                        ImageEditorMemento undoMemento = GetMementoFromShapes();
+                        undoMementoStack.Push(undoMemento);
 
                         shapeManager.RestoreState(redoMemento);
                     }
                     else
                     {
-                        ImageEditorMemento memento = GetMementoFromCanvas();
-                        undoMementoStack.Push(memento);
+                        ImageEditorMemento undoMemento = GetMementoFromCanvas();
+                        undoMementoStack.Push(undoMemento);
 
                         shapeManager.RestoreState(redoMemento);
                     }
@@ -136,16 +135,16 @@ public void Redo()
 
         public void Dispose()
         {
-            foreach (ImageEditorMemento memento in undoMementoStack)
+            foreach (ImageEditorMemento undoMemento in undoMementoStack)
             {
-                memento.Dispose();
+                undoMemento?.Dispose();
             }
 
             undoMementoStack.Clear();
 
             foreach (ImageEditorMemento redoMemento in redoMementoStack)
             {
-                redoMemento.Dispose();
+                redoMemento?.Dispose();
             }
 
             redoMementoStack.Clear();
diff --git a/ShareX.ScreenCaptureLib/Helpers/ImageEditorMemento.cs b/ShareX.ScreenCaptureLib/Helpers/ImageEditorMemento.cs
index d270d7960..3f740f1bc 100644
--- a/ShareX.ScreenCaptureLib/Helpers/ImageEditorMemento.cs
+++ b/ShareX.ScreenCaptureLib/Helpers/ImageEditorMemento.cs
@@ -49,7 +49,7 @@ public void Dispose()
         {
             foreach (BaseShape shape in Shapes)
             {
-                shape.Dispose();
+                shape?.Dispose();
             }
 
             Shapes.Clear();