From ce1465ae98a3797073d2494fed31f18dfdb35a71 Mon Sep 17 00:00:00 2001
From: Wojciech Figat <wojtek@figat.pl>
Date: Fri, 5 Aug 2022 11:27:37 +0200
Subject: [PATCH] Fix crash when reloading game scripts in Editor with Anim
 Graph that contains game structs as params

---
 Source/Editor/Surface/AnimGraphSurface.cs     | 37 +++++++++++--------
 .../Surface/VisjectSurface.Serialization.cs   |  4 +-
 Source/Engine/Level/Actors/AnimatedModel.cpp  |  8 ++++
 Source/Engine/Level/Actors/AnimatedModel.h    |  1 +
 4 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/Source/Editor/Surface/AnimGraphSurface.cs b/Source/Editor/Surface/AnimGraphSurface.cs
index be32871c1..7953851ee 100644
--- a/Source/Editor/Surface/AnimGraphSurface.cs
+++ b/Source/Editor/Surface/AnimGraphSurface.cs
@@ -269,8 +269,6 @@ namespace FlaxEditor.Surface
         /// </summary>
         protected VisjectCM _cmStateMachineTransitionMenu;
 
-        private bool _isRegisteredForScriptsReload;
-
         /// <inheritdoc />
         public AnimGraphSurface(IVisjectSurfaceOwner owner, Action onSave, FlaxEditor.Undo undo)
         : base(owner, onSave, undo, CreateStyle())
@@ -280,14 +278,9 @@ namespace FlaxEditor.Surface
             if (customNodes != null && customNodes.Count > 0)
             {
                 AddCustomNodes(customNodes);
-
-                // Check if any of the nodes comes from the game scripts - those can be reloaded at runtime so prevent crashes
-                if (Editor.Instance.CodeEditing.AnimGraphNodes.HasTypeFromGameScripts)
-                {
-                    _isRegisteredForScriptsReload = true;
-                    ScriptsBuilder.ScriptsReloadBegin += OnScriptsReloadBegin;
-                }
             }
+
+            ScriptsBuilder.ScriptsReloadBegin += OnScriptsReloadBegin;
         }
 
         private static SurfaceStyle CreateStyle()
@@ -301,8 +294,26 @@ namespace FlaxEditor.Surface
 
         private void OnScriptsReloadBegin()
         {
-            Owner.OnSurfaceClose();
+            // Check if any of the nodes comes from the game scripts - those can be reloaded at runtime so prevent crashes
+            bool hasTypeFromGameScripts = Editor.Instance.CodeEditing.AnimGraphNodes.HasTypeFromGameScripts;
 
+            // Check any surface parameter comes from Game scripts module to handle scripts reloads in Editor
+            if (!hasTypeFromGameScripts)
+            {
+                foreach (var param in Parameters)
+                {
+                    if (FlaxEngine.Scripting.IsTypeFromGameScripts(param.Type.Type))
+                    {
+                        hasTypeFromGameScripts = true;
+                        break;
+                    }
+                }
+            }
+
+            if (!hasTypeFromGameScripts)
+                return;
+
+            Owner.OnSurfaceClose();
             // TODO: make reload soft: dispose default primary context menu, update existing custom nodes to new ones or remove if invalid
         }
 
@@ -463,11 +474,7 @@ namespace FlaxEditor.Surface
                 _cmStateMachineTransitionMenu.Dispose();
                 _cmStateMachineTransitionMenu = null;
             }
-            if (_isRegisteredForScriptsReload)
-            {
-                _isRegisteredForScriptsReload = false;
-                ScriptsBuilder.ScriptsReloadBegin -= OnScriptsReloadBegin;
-            }
+            ScriptsBuilder.ScriptsReloadBegin -= OnScriptsReloadBegin;
             NodesCache.Wait();
 
             base.OnDestroy();
diff --git a/Source/Editor/Surface/VisjectSurface.Serialization.cs b/Source/Editor/Surface/VisjectSurface.Serialization.cs
index 07a420c1b..bb7fbed23 100644
--- a/Source/Editor/Surface/VisjectSurface.Serialization.cs
+++ b/Source/Editor/Surface/VisjectSurface.Serialization.cs
@@ -28,7 +28,7 @@ namespace FlaxEditor.Surface
         /// The method calls the <see cref="ISurfaceContext.SurfaceData"/> getter to load the surface data bytes.
         /// </remarks>
         /// <returns>True if failed, otherwise false.</returns>
-        public bool Load()
+        public virtual bool Load()
         {
             Enabled = false;
 
@@ -62,7 +62,7 @@ namespace FlaxEditor.Surface
         /// <remarks>
         /// The method calls the <see cref="ISurfaceContext.SurfaceData"/> setter to assign the result bytes. Sets null value if failed.
         /// </remarks>
-        public void Save()
+        public virtual void Save()
         {
             var wasEdited = IsEdited;
 
diff --git a/Source/Engine/Level/Actors/AnimatedModel.cpp b/Source/Engine/Level/Actors/AnimatedModel.cpp
index a5957e702..086b85e26 100644
--- a/Source/Engine/Level/Actors/AnimatedModel.cpp
+++ b/Source/Engine/Level/Actors/AnimatedModel.cpp
@@ -885,6 +885,14 @@ bool AnimatedModel::IntersectsEntry(const Ray& ray, Real& distance, Vector3& nor
     return result;
 }
 
+void AnimatedModel::OnDeleteObject()
+{
+    // Ensure this object is no longer referenced for anim update
+    Animations::RemoveFromUpdate(this);
+    
+    ModelInstanceActor::OnDeleteObject();
+}
+
 void AnimatedModel::OnTransformChanged()
 {
     // Base
diff --git a/Source/Engine/Level/Actors/AnimatedModel.h b/Source/Engine/Level/Actors/AnimatedModel.h
index 2dc9e86ac..16f2740dd 100644
--- a/Source/Engine/Level/Actors/AnimatedModel.h
+++ b/Source/Engine/Level/Actors/AnimatedModel.h
@@ -367,6 +367,7 @@ public:
     void Deserialize(DeserializeStream& stream, ISerializeModifier* modifier) override;
     bool IntersectsEntry(int32 entryIndex, const Ray& ray, Real& distance, Vector3& normal) override;
     bool IntersectsEntry(const Ray& ray, Real& distance, Vector3& normal, int32& entryIndex) override;
+    void OnDeleteObject() override;
 
 protected:
     // [ModelInstanceActor]