Add change state handling to the editor

osu.Game/Screens/Edit/Editor.cs

@@ -62,6 +62,7 @@ namespace osu.Game.Screens.Edit
 
         private IBeatmap playableBeatmap;
         private EditorBeatmap editorBeatmap;
+        private EditorChangeHandler changeHandler;
 
         private DependencyContainer dependencies;
 
@@ -100,9 +101,11 @@ namespace osu.Game.Screens.Edit
             }
 
             AddInternal(editorBeatmap = new EditorBeatmap(playableBeatmap));
-
             dependencies.CacheAs(editorBeatmap);
 
+            changeHandler = new EditorChangeHandler(editorBeatmap);
+            dependencies.CacheAs<IEditorChangeHandler>(changeHandler);
+
             EditorMenuBar menuBar;
 
             var fileMenuItems = new List<MenuItem>
@@ -147,6 +150,14 @@ namespace osu.Game.Screens.Edit
                                 new MenuItem("File")
                                 {
                                     Items = fileMenuItems
+                                },
+                                new MenuItem("Edit")
+                                {
+                                    Items = new[]
+                                    {
+                                        new EditorMenuItem("Undo", MenuItemType.Standard, undo),
+                                        new EditorMenuItem("Redo", MenuItemType.Standard, redo)
+                                    }
                                 }
                             }
                         }
@@ -233,6 +244,19 @@ namespace osu.Game.Screens.Edit
                         return true;
                     }
 
+                    break;
+
+                case Key.Z:
+                    if (e.ControlPressed)
+                    {
+                        if (e.ShiftPressed)
+                            redo();
+                        else
+                            undo();
+
+                        return true;
+                    }
+
                     break;
             }
 
@@ -297,6 +321,10 @@ namespace osu.Game.Screens.Edit
             return base.OnExiting(next);
         }
 
+        private void undo() => changeHandler.RestoreState(-1);
+
+        private void redo() => changeHandler.RestoreState(1);
+
         private void resetTrack(bool seekToStart = false)
         {
             Beatmap.Value.Track?.Stop();

osu.Game/Screens/Edit/EditorChangeHandler.cs

@@ -0,0 +1,108 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Text;
+using osu.Game.Beatmaps.Formats;
+using osu.Game.Rulesets.Objects;
+
+namespace osu.Game.Screens.Edit
+{
+    /// <summary>
+    /// Tracks changes to the <see cref="Editor"/>.
+    /// </summary>
+    public class EditorChangeHandler : IEditorChangeHandler
+    {
+        private readonly LegacyEditorBeatmapDiffer differ;
+        private readonly List<Stream> savedStates = new List<Stream>();
+        private int currentState = -1;
+
+        private readonly EditorBeatmap editorBeatmap;
+        private int bulkChangesStarted;
+        private bool isRestoring;
+
+        /// <summary>
+        /// Creates a new <see cref="EditorChangeHandler"/>.
+        /// </summary>
+        /// <param name="editorBeatmap">The <see cref="EditorBeatmap"/> to track the <see cref="HitObject"/>s of.</param>
+        public EditorChangeHandler(EditorBeatmap editorBeatmap)
+        {
+            this.editorBeatmap = editorBeatmap;
+
+            editorBeatmap.HitObjectAdded += hitObjectAdded;
+            editorBeatmap.HitObjectRemoved += hitObjectRemoved;
+            editorBeatmap.HitObjectUpdated += hitObjectUpdated;
+
+            differ = new LegacyEditorBeatmapDiffer(editorBeatmap);
+
+            // Initial state.
+            SaveState();
+        }
+
+        private void hitObjectAdded(HitObject obj) => SaveState();
+
+        private void hitObjectRemoved(HitObject obj) => SaveState();
+
+        private void hitObjectUpdated(HitObject obj) => SaveState();
+
+        public void BeginChange() => bulkChangesStarted++;
+
+        public void EndChange()
+        {
+            if (bulkChangesStarted == 0)
+                throw new InvalidOperationException($"Cannot call {nameof(EndChange)} without a previous call to {nameof(BeginChange)}.");
+
+            if (--bulkChangesStarted == 0)
+                SaveState();
+        }
+
+        /// <summary>
+        /// Saves the current <see cref="Editor"/> state.
+        /// </summary>
+        public void SaveState()
+        {
+            if (bulkChangesStarted > 0)
+                return;
+
+            if (isRestoring)
+                return;
+
+            var stream = new MemoryStream();
+
+            using (var sw = new StreamWriter(stream, Encoding.UTF8, 1024, true))
+                new LegacyBeatmapEncoder(editorBeatmap).Encode(sw);
+
+            if (currentState < savedStates.Count - 1)
+                savedStates.RemoveRange(currentState + 1, savedStates.Count - currentState - 1);
+
+            savedStates.Add(stream);
+            currentState = savedStates.Count - 1;
+        }
+
+        /// <summary>
+        /// Restores an older or newer state.
+        /// </summary>
+        /// <param name="direction">The direction to restore in. If less than 0, an older state will be used. If greater than 0, a newer state will be used.</param>
+        public void RestoreState(int direction)
+        {
+            if (bulkChangesStarted > 0)
+                return;
+
+            if (savedStates.Count == 0)
+                return;
+
+            int newState = Math.Clamp(currentState + direction, 0, savedStates.Count - 1);
+            if (currentState == newState)
+                return;
+
+            isRestoring = true;
+
+            differ.Patch(savedStates[currentState], savedStates[newState]);
+            currentState = newState;
+
+            isRestoring = false;
+        }
+    }
+}

osu.Game/Screens/Edit/IEditorChangeHandler.cs

@@ -0,0 +1,33 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using osu.Game.Rulesets.Objects;
+
+namespace osu.Game.Screens.Edit
+{
+    /// <summary>
+    /// Interface for a component that manages changes in the <see cref="Editor"/>.
+    /// </summary>
+    public interface IEditorChangeHandler
+    {
+        /// <summary>
+        /// Begins a bulk state change event. <see cref="EndChange"/> should be invoked soon after.
+        /// </summary>
+        /// <remarks>
+        /// This should be invoked when multiple changes to the <see cref="Editor"/> should be bundled together into one state change event.
+        /// When nested invocations are involved, a state change will not occur until an equal number of invocations of <see cref="EndChange"/> are received.
+        /// </remarks>
+        /// <example>
+        /// When a group of <see cref="HitObject"/>s are deleted, a single undo and redo state change should update the state of all <see cref="HitObject"/>.
+        /// </example>
+        void BeginChange();
+
+        /// <summary>
+        /// Ends a bulk state change event.
+        /// </summary>
+        /// <remarks>
+        /// This should be invoked as soon as possible after <see cref="BeginChange"/> to cause a state change.
+        /// </remarks>
+        void EndChange();
+    }
+}