From c561d684eba1e0cea290b5d07a2b2603132bce72 Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Wed, 6 Mar 2024 19:01:36 +0100
Subject: [PATCH] Refactor undo logic for actors reparenting in Editor

#1741
---
 Source/Editor/SceneGraph/ActorNode.cs         |  36 +--
 Source/Editor/SceneGraph/GUI/ActorTreeNode.cs | 221 +++++-------------
 .../Editor/Undo/Actions/ParentActorsAction.cs | 199 ++++++++++++++++
 .../Editor/Undo/Actions/PasteActorsAction.cs  |   1 -
 4 files changed, 257 insertions(+), 200 deletions(-)
 create mode 100644 Source/Editor/Undo/Actions/ParentActorsAction.cs

diff --git a/Source/Editor/SceneGraph/ActorNode.cs b/Source/Editor/SceneGraph/ActorNode.cs
index 5c6adb055..91bf26103 100644
--- a/Source/Editor/SceneGraph/ActorNode.cs
+++ b/Source/Editor/SceneGraph/ActorNode.cs
@@ -316,41 +316,7 @@ namespace FlaxEditor.SceneGraph
         {
             base.OnParentChanged();
 
-            // Update UI (special case if actor is spawned and added to existing scene tree)
-            var parentTreeNode = (parentNode as ActorNode)?.TreeNode;
-            if (parentTreeNode != null && !parentTreeNode.IsLayoutLocked)
-            {
-                parentTreeNode.IsLayoutLocked = true;
-                _treeNode.Parent = parentTreeNode;
-                _treeNode.IndexInParent = _actor.OrderInParent;
-                parentTreeNode.IsLayoutLocked = false;
-
-                // Skip UI update if node won't be in a view
-                if (parentTreeNode.IsCollapsed)
-                {
-                    TreeNode.UnlockChildrenRecursive();
-                }
-                else
-                {
-                    // Try to perform layout at the level where it makes it the most performant (the least computations)
-                    var tree = parentTreeNode.ParentTree;
-                    if (tree != null)
-                    {
-                        if (tree.Parent is FlaxEngine.GUI.Panel treeParent)
-                            treeParent.PerformLayout();
-                        else
-                            tree.PerformLayout();
-                    }
-                    else
-                    {
-                        parentTreeNode.PerformLayout();
-                    }
-                }
-            }
-            else
-            {
-                _treeNode.Parent = parentTreeNode;
-            }
+            _treeNode.OnParentChanged(_actor, parentNode as ActorNode);
         }
 
         /// <inheritdoc />
diff --git a/Source/Editor/SceneGraph/GUI/ActorTreeNode.cs b/Source/Editor/SceneGraph/GUI/ActorTreeNode.cs
index 270647433..6056cf68b 100644
--- a/Source/Editor/SceneGraph/GUI/ActorTreeNode.cs
+++ b/Source/Editor/SceneGraph/GUI/ActorTreeNode.cs
@@ -3,6 +3,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using FlaxEditor.Actions;
 using FlaxEditor.Content;
 using FlaxEditor.GUI;
 using FlaxEditor.GUI.Drag;
@@ -14,7 +15,6 @@ using FlaxEditor.Windows.Assets;
 using FlaxEngine;
 using FlaxEngine.GUI;
 using FlaxEngine.Utilities;
-using Object = FlaxEngine.Object;
 
 namespace FlaxEditor.SceneGraph.GUI
 {
@@ -82,8 +82,51 @@ namespace FlaxEditor.SceneGraph.GUI
             UpdateText();
         }
 
+        internal void OnParentChanged(Actor actor, ActorNode parentNode)
+        {
+            // Update cached value
+            _orderInParent = actor.OrderInParent;
+
+            // Update UI (special case if actor is spawned and added to existing scene tree)
+            var parentTreeNode = parentNode?.TreeNode;
+            if (parentTreeNode != null && !parentTreeNode.IsLayoutLocked)
+            {
+                parentTreeNode.IsLayoutLocked = true;
+                Parent = parentTreeNode;
+                IndexInParent = _orderInParent;
+                parentTreeNode.IsLayoutLocked = false;
+
+                // Skip UI update if node won't be in a view
+                if (parentTreeNode.IsCollapsed)
+                {
+                    UnlockChildrenRecursive();
+                }
+                else
+                {
+                    // Try to perform layout at the level where it makes it the most performant (the least computations)
+                    var tree = parentTreeNode.ParentTree;
+                    if (tree != null)
+                    {
+                        if (tree.Parent is Panel treeParent)
+                            treeParent.PerformLayout();
+                        else
+                            tree.PerformLayout();
+                    }
+                    else
+                    {
+                        parentTreeNode.PerformLayout();
+                    }
+                }
+            }
+            else
+            {
+                Parent = parentTreeNode;
+            }
+        }
+
         internal void OnOrderInParentChanged()
         {
+            // Use cached value to check if we need to update UI layout (and update siblings order at once)
             if (Parent is ActorTreeNode parent)
             {
                 var anyChanged = false;
@@ -419,134 +462,6 @@ namespace FlaxEditor.SceneGraph.GUI
             _dragHandlers.OnDragLeave();
         }
 
-        [Serializable]
-        private class ReparentAction : IUndoAction
-        {
-            [Serialize]
-            private Guid[] _ids;
-
-            [Serialize]
-            private int _actorsCount;
-
-            [Serialize]
-            private Guid[] _prefabIds;
-
-            [Serialize]
-            private Guid[] _prefabObjectIds;
-
-            public ReparentAction(Actor actor)
-            : this(new List<Actor> { actor })
-            {
-            }
-
-            public ReparentAction(List<Actor> actors)
-            {
-                var allActors = new List<Actor>(Mathf.NextPowerOfTwo(actors.Count));
-
-                for (int i = 0; i < actors.Count; i++)
-                {
-                    GetAllActors(allActors, actors[i]);
-                }
-
-                var allScripts = new List<Script>(allActors.Capacity);
-                GetAllScripts(allActors, allScripts);
-
-                int allCount = allActors.Count + allScripts.Count;
-                _actorsCount = allActors.Count;
-                _ids = new Guid[allCount];
-                _prefabIds = new Guid[allCount];
-                _prefabObjectIds = new Guid[allCount];
-
-                for (int i = 0; i < allActors.Count; i++)
-                {
-                    _ids[i] = allActors[i].ID;
-                    _prefabIds[i] = allActors[i].PrefabID;
-                    _prefabObjectIds[i] = allActors[i].PrefabObjectID;
-                }
-
-                for (int i = 0; i < allScripts.Count; i++)
-                {
-                    int j = _actorsCount + i;
-                    _ids[j] = allScripts[i].ID;
-                    _prefabIds[j] = allScripts[i].PrefabID;
-                    _prefabObjectIds[j] = allScripts[i].PrefabObjectID;
-                }
-            }
-
-            public ReparentAction(Script script)
-            {
-                _actorsCount = 0;
-                _ids = new Guid[] { script.ID };
-                _prefabIds = new Guid[] { script.PrefabID };
-                _prefabObjectIds = new Guid[] { script.PrefabObjectID };
-            }
-
-            private void GetAllActors(List<Actor> allActors, Actor actor)
-            {
-                allActors.Add(actor);
-
-                for (int i = 0; i < actor.ChildrenCount; i++)
-                {
-                    var child = actor.GetChild(i);
-                    if (!allActors.Contains(child))
-                    {
-                        GetAllActors(allActors, child);
-                    }
-                }
-            }
-
-            private void GetAllScripts(List<Actor> allActors, List<Script> allScripts)
-            {
-                for (int i = 0; i < allActors.Count; i++)
-                {
-                    var actor = allActors[i];
-                    for (int j = 0; j < actor.ScriptsCount; j++)
-                    {
-                        allScripts.Add(actor.GetScript(j));
-                    }
-                }
-            }
-
-            /// <inheritdoc />
-            public string ActionString => string.Empty;
-
-            /// <inheritdoc />
-            public void Do()
-            {
-                // Note: prefab links are broken by the C++ backend on actor reparenting
-            }
-
-            /// <inheritdoc />
-            public void Undo()
-            {
-                // Restore links
-                for (int i = 0; i < _actorsCount; i++)
-                {
-                    var actor = Object.Find<Actor>(ref _ids[i]);
-                    if (actor != null && _prefabIds[i] != Guid.Empty)
-                    {
-                        Actor.Internal_LinkPrefab(Object.GetUnmanagedPtr(actor), ref _prefabIds[i], ref _prefabObjectIds[i]);
-                    }
-                }
-                for (int i = _actorsCount; i < _ids.Length; i++)
-                {
-                    var script = Object.Find<Script>(ref _ids[i]);
-                    if (script != null && _prefabIds[i] != Guid.Empty)
-                    {
-                        Script.Internal_LinkPrefab(Object.GetUnmanagedPtr(script), ref _prefabIds[i], ref _prefabObjectIds[i]);
-                    }
-                }
-            }
-
-            /// <inheritdoc />
-            public void Dispose()
-            {
-                _ids = null;
-                _prefabIds = null;
-                _prefabObjectIds = null;
-            }
-        }
-
         /// <inheritdoc />
         protected override DragDropEffect OnDragDropHeader(DragData data)
         {
@@ -593,46 +508,24 @@ namespace FlaxEditor.SceneGraph.GUI
             // Drag actors
             if (_dragActors != null && _dragActors.HasValidDrag)
             {
-                bool worldPositionLock = Root.GetKey(KeyboardKeys.Control) == false;
-                var singleObject = _dragActors.Objects.Count == 1;
-                if (singleObject)
-                {
-                    var targetActor = _dragActors.Objects[0].Actor;
-                    var customAction = targetActor.HasPrefabLink ? new ReparentAction(targetActor) : null;
-                    using (new UndoBlock(ActorNode.Root.Undo, targetActor, "Change actor parent", customAction))
-                    {
-                        targetActor.SetParent(newParent, worldPositionLock, true);
-                        targetActor.OrderInParent = newOrder;
-                    }
-                }
-                else
-                {
-                    var targetActors = _dragActors.Objects.ConvertAll(x => x.Actor);
-                    var customAction = targetActors.Any(x => x.HasPrefabLink) ? new ReparentAction(targetActors) : null;
-                    using (new UndoMultiBlock(ActorNode.Root.Undo, targetActors, "Change actors parent", customAction))
-                    {
-                        for (int i = 0; i < targetActors.Count; i++)
-                        {
-                            var targetActor = targetActors[i];
-                            targetActor.SetParent(newParent, worldPositionLock, true);
-                            targetActor.OrderInParent = newOrder;
-                        }
-                    }
-                }
-
+                bool worldPositionsStays = Root.GetKey(KeyboardKeys.Control) == false;
+                var objects = new SceneObject[_dragActors.Objects.Count];
+                for (int i = 0; i < objects.Length; i++)
+                    objects[i] = _dragActors.Objects[i].Actor;
+                var action = new ParentActorsAction(objects, newParent, newOrder, worldPositionsStays);
+                ActorNode.Root.Undo?.AddAction(action);
+                action.Do();
                 result = DragDropEffect.Move;
             }
             // Drag scripts
             else if (_dragScripts != null && _dragScripts.HasValidDrag)
             {
-                foreach (var script in _dragScripts.Objects)
-                {
-                    var customAction = script.HasPrefabLink ? new ReparentAction(script) : null;
-                    using (new UndoBlock(ActorNode.Root.Undo, script, "Change script parent", customAction))
-                    {
-                        script.SetParent(newParent, true);
-                    }
-                }
+                var objects = new SceneObject[_dragScripts.Objects.Count];
+                for (int i = 0; i < objects.Length; i++)
+                    objects[i] = _dragScripts.Objects[i];
+                var action = new ParentActorsAction(objects, newParent, newOrder);
+                ActorNode.Root.Undo?.AddAction(action);
+                action.Do();
                 Select();
                 result = DragDropEffect.Move;
             }
diff --git a/Source/Editor/Undo/Actions/ParentActorsAction.cs b/Source/Editor/Undo/Actions/ParentActorsAction.cs
new file mode 100644
index 000000000..56957ff84
--- /dev/null
+++ b/Source/Editor/Undo/Actions/ParentActorsAction.cs
@@ -0,0 +1,199 @@
+// Copyright (c) 2012-2024 Wojciech Figat. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using FlaxEngine;
+using Object = FlaxEngine.Object;
+
+namespace FlaxEditor.Actions
+{
+    /// <summary>
+    /// Implementation of <see cref="IUndoAction"/> used to change parent for <see cref="Actor"/> or <see cref="Script"/>.
+    /// </summary>
+    /// <seealso cref="FlaxEditor.IUndoAction" />
+    [Serializable]
+    class ParentActorsAction : IUndoAction
+    {
+        private struct Item
+        {
+            public Guid ID;
+            public Guid Parent;
+            public int OrderInParent;
+            public Transform LocalTransform;
+        }
+
+        [Serialize]
+        private bool _worldPositionsStays;
+
+        [Serialize]
+        private Guid _newParent;
+
+        [Serialize]
+        private int _newOrder;
+
+        [Serialize]
+        private Item[] _items;
+
+        [Serialize]
+        private Guid[] _idsForPrefab;
+
+        [Serialize]
+        private Guid[] _prefabIds;
+
+        [Serialize]
+        private Guid[] _prefabObjectIds;
+
+        public ParentActorsAction(SceneObject[] objects, Actor newParent, int newOrder, bool worldPositionsStays = true)
+        {
+            // Sort source objects to provide deterministic behavior
+            Array.Sort(objects, SortObjects);
+
+            // Cache initial state for undo
+            _worldPositionsStays = worldPositionsStays;
+            _newParent = newParent.ID;
+            _newOrder = newOrder;
+            _items = new Item[objects.Length];
+            for (int i = 0; i < objects.Length; i++)
+            {
+                var obj = objects[i];
+                _items[i] = new Item
+                {
+                    ID = obj.ID,
+                    Parent = obj.Parent?.ID ?? Guid.Empty,
+                    OrderInParent = obj.OrderInParent,
+                    LocalTransform = obj is Actor actor ? actor.LocalTransform : Transform.Identity,
+                };
+            }
+
+            // Collect all objects that have prefab links so they can be restored on undo
+            var prefabs = new List<SceneObject>();
+            for (int i = 0; i < objects.Length; i++)
+                GetAllPrefabs(prefabs, objects[i]);
+            if (prefabs.Count != 0)
+            {
+                // Cache ids of all objects
+                _idsForPrefab = new Guid[prefabs.Count];
+                _prefabIds = new Guid[prefabs.Count];
+                _prefabObjectIds = new Guid[prefabs.Count];
+                for (int i = 0; i < prefabs.Count; i++)
+                {
+                    var obj = prefabs[i];
+                    _idsForPrefab[i] = obj.ID;
+                    _prefabIds[i] = obj.PrefabID;
+                    _prefabObjectIds[i] = obj.PrefabObjectID;
+                }
+            }
+        }
+
+        private static int SortObjects(SceneObject a, SceneObject b)
+        {
+            // By parent
+            var aParent = Object.GetUnmanagedPtr(a.Parent);
+            var bParent = Object.GetUnmanagedPtr(b.Parent);
+            if (aParent == bParent)
+            {
+                // By index in parent
+                var aOrder = a.OrderInParent;
+                var bOrder = b.OrderInParent;
+                return aOrder.CompareTo(bOrder);
+            }
+            return aParent.CompareTo(bParent);
+        }
+
+        private static void GetAllPrefabs(List<SceneObject> result, SceneObject obj)
+        {
+            if (result.Contains(obj))
+                return;
+            if (obj.HasPrefabLink)
+                result.Add(obj);
+            if (obj is Actor actor)
+            {
+                for (int i = 0; i < actor.ScriptsCount; i++)
+                    GetAllPrefabs(result, actor.GetScript(i));
+                for (int i = 0; i < actor.ChildrenCount; i++)
+                    GetAllPrefabs(result, actor.GetChild(i));
+            }
+        }
+
+        public string ActionString => "Change parent";
+
+        public void Do()
+        {
+            // Perform action
+            var newParent = Object.Find<Actor>(ref _newParent);
+            if (newParent == null)
+            {
+                Editor.LogError("Missing actor to change objects parent.");
+                return;
+            }
+            var order = _newOrder;
+            var scenes = new HashSet<Scene> { newParent.Scene };
+            for (int i = 0; i < _items.Length; i++)
+            {
+                var item = _items[i];
+                var obj = Object.Find<SceneObject>(ref item.ID);
+                if (obj != null)
+                {
+                    scenes.Add(obj.Parent.Scene);
+                    if (obj is Actor actor)
+                        actor.SetParent(newParent, _worldPositionsStays, true);
+                    else
+                        obj.Parent = newParent;
+                    if (order != -1)
+                        obj.OrderInParent = order++;
+                }
+            }
+
+            // Prefab links are broken by the C++ backend on actor reparenting
+
+            // Mark scenes as edited
+            foreach (var scene in scenes)
+                Editor.Instance.Scene.MarkSceneEdited(scene);
+        }
+
+        public void Undo()
+        {
+            // Restore state
+            for (int i = 0; i < _items.Length; i++)
+            {
+                var item = _items[i];
+                var obj = Object.Find<SceneObject>(ref item.ID);
+                if (obj != null)
+                {
+                    var parent = Object.Find<Actor>(ref item.Parent);
+                    if (parent != null)
+                        obj.Parent = parent;
+                    if (obj is Actor actor)
+                        actor.LocalTransform = item.LocalTransform;
+                }
+            }
+            for (int j = 0; j < _items.Length; j++) // TODO: find a better way ensure the order is properly restored when moving back multiple objects
+            for (int i = 0; i < _items.Length; i++)
+            {
+                var item = _items[i];
+                var obj = Object.Find<SceneObject>(ref item.ID);
+                if (obj != null)
+                    obj.OrderInParent = item.OrderInParent;
+            }
+
+            // Restore prefab links (if any was in use)
+            if (_idsForPrefab != null)
+            {
+                for (int i = 0; i < _idsForPrefab.Length; i++)
+                {
+                    var obj = Object.Find<SceneObject>(ref _idsForPrefab[i]);
+                    if (obj != null && _prefabIds[i] != Guid.Empty)
+                        SceneObject.Internal_LinkPrefab(Object.GetUnmanagedPtr(obj), ref _prefabIds[i], ref _prefabObjectIds[i]);
+                }
+            }
+        }
+
+        public void Dispose()
+        {
+            _items = null;
+            _idsForPrefab = null;
+            _prefabIds = null;
+            _prefabObjectIds = null;
+        }
+    }
+}
diff --git a/Source/Editor/Undo/Actions/PasteActorsAction.cs b/Source/Editor/Undo/Actions/PasteActorsAction.cs
index 87a21ec2d..e43d3bb4d 100644
--- a/Source/Editor/Undo/Actions/PasteActorsAction.cs
+++ b/Source/Editor/Undo/Actions/PasteActorsAction.cs
@@ -2,7 +2,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using FlaxEditor.SceneGraph;
 using FlaxEngine;