Fix error when using nested Visject Surface context during State Machines editing in Anim Graph

Source/Editor/Surface/Archetypes/Animation.StateMachine.cs

@@ -250,6 +250,9 @@ namespace FlaxEditor.Surface.Archetypes
                 set => Values[1] = value;
             }
 
+            /// <inheritdoc />
+            public VisjectSurfaceContext ParentContext => Context;
+
             /// <inheritdoc />
             public void OnContextCreated(VisjectSurfaceContext context)
             {
@@ -1311,6 +1314,9 @@ namespace FlaxEditor.Surface.Archetypes
                 set => Values[1] = value;
             }
 
+            /// <inheritdoc />
+            public VisjectSurfaceContext ParentContext => Context;
+
             /// <inheritdoc />
             public void OnContextCreated(VisjectSurfaceContext context)
             {
@@ -1682,6 +1688,9 @@ namespace FlaxEditor.Surface.Archetypes
                 set => RuleGraph = value;
             }
 
+            /// <inheritdoc />
+            public VisjectSurfaceContext ParentContext => SourceState.Context;
+
             /// <inheritdoc />
             public void OnContextCreated(VisjectSurfaceContext context)
             {

Source/Editor/Surface/ISurfaceContext.cs

@@ -25,6 +25,11 @@ namespace FlaxEditor.Surface
         /// </summary>
         byte[] SurfaceData { get; set; }
 
+        /// <summary>
+        /// Gets the context which owns this surface context (null for root).
+        /// </summary>
+        VisjectSurfaceContext ParentContext { get; }
+
         /// <summary>
         /// Called when Visject Surface context gets created for this surface data source. Can be used to link for some events.
         /// </summary>

Source/Editor/Surface/Undo/ConnectBoxesAction.cs

@@ -33,6 +33,14 @@ namespace FlaxEditor.Surface.Undo
 
             CaptureConnections(iB, out _inputBefore);
             CaptureConnections(oB, out _outputBefore);
+
+#if BUILD_DEBUG
+            // Validate handles
+            if (_context.Get(_surface) != iB.ParentNode.Context)
+                throw new System.Exception("Invalid ContextHandle");
+            if (_input.Get(iB.ParentNode.Context) != iB || _output.Get(oB.ParentNode.Context) != oB)
+                throw new System.Exception("Invalid BoxHandle");
+#endif
         }
 
         public void End()

Source/Editor/Surface/Undo/ContextHandle.cs

@@ -1,7 +1,9 @@
 // Copyright (c) 2012-2023 Wojciech Figat. All rights reserved.
 
 using System;
+using System.Text;
 using FlaxEngine;
+using FlaxEngine.Json;
 
 namespace FlaxEditor.Surface.Undo
 {
@@ -9,7 +11,7 @@ namespace FlaxEditor.Surface.Undo
     /// The helper structure for Surface context handle.
     /// </summary>
     [HideInEditor]
-    public struct ContextHandle
+    public struct ContextHandle : IEquatable<ContextHandle>
     {
         private readonly string[] _path;
 
@@ -42,6 +44,31 @@ namespace FlaxEditor.Surface.Undo
             }
         }
 
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ContextHandle"/> struct.
+        /// </summary>
+        /// <param name="child">The child context provider.</param>
+        public ContextHandle(ISurfaceContext child)
+        {
+            if (child == null)
+                throw new ArgumentNullException();
+            VisjectSurfaceContext parent = child.ParentContext;
+            VisjectSurfaceContext context = parent;
+            int count = 1;
+            while (parent != null)
+            {
+                count++;
+                parent = parent.Parent;
+            }
+            _path = new string[count];
+            _path[0] = child.SurfaceName;
+            for (int i = 1; i < count; i++)
+            {
+                _path[i] = context.Context.SurfaceName;
+                context = context.Parent;
+            }
+        }
+
         /// <summary>
         /// Gets the context.
         /// </summary>
@@ -69,5 +96,46 @@ namespace FlaxEditor.Surface.Undo
 
             return context;
         }
+
+        /// <inheritdoc />
+        public bool Equals(ContextHandle other)
+        {
+            return JsonSerializer.ValueEquals(_path, other._path);
+        }
+
+        /// <inheritdoc />
+        public override bool Equals(object obj)
+        {
+            return obj is ContextHandle other && Equals(other);
+        }
+
+        /// <inheritdoc />
+        public override int GetHashCode()
+        {
+            int hash = 17;
+            if (_path != null)
+            {
+                unchecked
+                {
+                    for (var i = 0; i < _path.Length; i++)
+                    {
+                        var item = _path[i];
+                        hash = hash * 23 + (item != null ? item.GetHashCode() : 0);
+                    }
+                }
+            }
+            return hash;
+        }
+
+        /// <inheritdoc />
+        public override string ToString()
+        {
+            var sb = new StringBuilder();
+            if (_path == null)
+                return string.Empty;
+            for (int i = _path.Length - 1; i >= 0; i--)
+                sb.Append(_path[i]).Append('/');
+            return sb.ToString();
+        }
     }
 }

Source/Editor/Surface/VisjectSurface.Context.cs

@@ -2,6 +2,8 @@
 
 using System;
 using System.Collections.Generic;
+using FlaxEditor.Surface.Undo;
+using FlaxEngine;
 
 namespace FlaxEditor.Surface
 {
@@ -9,7 +11,7 @@ namespace FlaxEditor.Surface
     {
         private VisjectSurfaceContext _root;
         private VisjectSurfaceContext _context;
-        private readonly Dictionary<ISurfaceContext, VisjectSurfaceContext> _contextCache = new Dictionary<ISurfaceContext, VisjectSurfaceContext>();
+        private readonly Dictionary<ContextHandle, VisjectSurfaceContext> _contextCache = new Dictionary<ContextHandle, VisjectSurfaceContext>();
 
         /// <summary>
         /// The surface context stack.
@@ -54,11 +56,12 @@ namespace FlaxEditor.Surface
                 return;
 
             // Get or create context
-            if (!_contextCache.TryGetValue(context, out VisjectSurfaceContext surfaceContext))
+            var contextHandle = new ContextHandle(context);
+            if (!_contextCache.TryGetValue(contextHandle, out VisjectSurfaceContext surfaceContext))
             {
                 surfaceContext = CreateContext(_context, context);
                 _context?.Children.Add(surfaceContext);
-                _contextCache.Add(context, surfaceContext);
+                _contextCache.Add(contextHandle, surfaceContext);
 
                 context.OnContextCreated(surfaceContext);
 
@@ -118,7 +121,8 @@ namespace FlaxEditor.Surface
             }
 
             // Check if has context in cache
-            if (_contextCache.TryGetValue(context, out VisjectSurfaceContext surfaceContext))
+            var contextHandle = new ContextHandle(context);
+            if (_contextCache.TryGetValue(contextHandle, out VisjectSurfaceContext surfaceContext))
             {
                 // Remove from navigation path
                 while (ContextStack.Contains(surfaceContext))
@@ -126,7 +130,7 @@ namespace FlaxEditor.Surface
 
                 // Dispose
                 surfaceContext.Clear();
-                _contextCache.Remove(context);
+                _contextCache.Remove(contextHandle);
             }
         }
 
@@ -147,7 +151,8 @@ namespace FlaxEditor.Surface
                 return;
 
             // Check if already in a path
-            if (_contextCache.TryGetValue(context, out VisjectSurfaceContext surfaceContext) && ContextStack.Contains(surfaceContext))
+            var contextHandle = new ContextHandle(context);
+            if (_contextCache.TryGetValue(contextHandle, out VisjectSurfaceContext surfaceContext) && ContextStack.Contains(surfaceContext))
             {
                 // Change stack
                 do

Source/Editor/Surface/VisjectSurfaceWindow.cs

@@ -899,6 +899,9 @@ namespace FlaxEditor.Surface
         /// <inheritdoc />
         public abstract byte[] SurfaceData { get; set; }
 
+        /// <inheritdoc />
+        public VisjectSurfaceContext ParentContext => null;
+
         /// <inheritdoc />
         public void OnContextCreated(VisjectSurfaceContext context)
         {

Source/Editor/Viewport/Previews/MaterialPreview.cs

@@ -366,6 +366,9 @@ namespace FlaxEditor.Viewport.Previews
             }
         }
 
+        /// <inheritdoc />
+        public VisjectSurfaceContext ParentContext => null;
+
         /// <inheritdoc />
         void ISurfaceContext.OnContextCreated(VisjectSurfaceContext context)
         {

Source/Editor/Windows/AssetReferencesGraphWindow.cs

@@ -408,6 +408,9 @@ namespace FlaxEditor.Windows
         /// <inheritdoc />
         public byte[] SurfaceData { get; set; }
 
+        /// <inheritdoc />
+        public VisjectSurfaceContext ParentContext => null;
+
         /// <inheritdoc />
         public void OnContextCreated(VisjectSurfaceContext context)
         {

Source/Editor/Windows/Assets/VisjectFunctionSurfaceWindow.cs

@@ -179,6 +179,9 @@ namespace FlaxEditor.Windows.Assets
         /// <inheritdoc />
         public abstract byte[] SurfaceData { get; set; }
 
+        /// <inheritdoc />
+        public VisjectSurfaceContext ParentContext => null;
+
         /// <inheritdoc />
         public void OnContextCreated(VisjectSurfaceContext context)
         {

Source/Editor/Windows/Assets/VisualScriptWindow.cs

@@ -1147,6 +1147,9 @@ namespace FlaxEditor.Windows.Assets
             }
         }
 
+        /// <inheritdoc />
+        public VisjectSurfaceContext ParentContext => null;
+
         /// <inheritdoc />
         public void OnContextCreated(VisjectSurfaceContext context)
         {

Source/Editor/Windows/Search/ContentSearchWindow.cs

@@ -105,6 +105,9 @@ namespace FlaxEngine.Windows.Search
             /// <inheritdoc />
             public byte[] SurfaceData { get; set; }
 
+            /// <inheritdoc />
+            public VisjectSurfaceContext ParentContext => null;
+
             /// <inheritdoc />
             public void OnContextCreated(VisjectSurfaceContext context)
             {