Simplify code for traversing for pair node for node zones

src/PixiEditor/Models/Handlers/INodeHandler.cs

@@ -9,6 +9,7 @@ using Drawie.Backend.Core;
 using PixiEditor.Models.Rendering;
 using PixiEditor.Models.Structures;
 using Drawie.Numerics;
+using PixiEditor.ViewModels.Nodes;
 
 namespace PixiEditor.Models.Handlers;
 
@@ -33,6 +34,7 @@ public interface INodeHandler : INotifyPropertyChanged, IDisposable
     public void TraverseForwards(Func<INodeHandler, INodeHandler, bool> func);
     public void TraverseForwards(Func<INodeHandler, INodeHandler, INodePropertyHandler, bool> func);
     public void TraverseForwards(Func<INodeHandler, INodeHandler, INodePropertyHandler, INodePropertyHandler, bool> func);
+    public List<NodeFrameViewModelBase> Frames { get; }
     public IReadOnlyDictionary<string, INodePropertyHandler> InputPropertyMap { get; }
     public IReadOnlyDictionary<string, INodePropertyHandler> OutputPropertyMap { get; }
 }

src/PixiEditor/ViewModels/Document/NodeGraphViewModel.cs

@@ -109,33 +109,12 @@ internal class NodeGraphViewModel : ViewModelBase, INodeGraphHandler, IDisposabl
 
         Connections.Add(connection);
         
-        AddToFramesPartOf(connection.InputNode);
-        AddToFramesPartOf(connection.OutputNode);
+        UpdatesFramesPartOf(connection.InputNode);
+        UpdatesFramesPartOf(connection.OutputNode);
 
         StructureTree.Update(this);
     }
 
-    private void AddToFramesPartOf(INodeHandler node)
-    {
-        node.TraverseBackwards(x =>
-        {
-            if (x is IPairNodeEndViewModel)
-                return false;
-
-            if (x is not IPairNodeStartViewModel)
-                return true;
-
-            var zone = Frames
-                .OfType<NodeZoneViewModel>()
-                .First(z => z.Start == x);
-
-            if (!zone.Nodes.Contains(node))
-                zone.Nodes.Add(node);
-
-            return true;
-        });
-    }
-
     public void RemoveConnection(Guid nodeId, string property)
     {
         var connection = Connections.FirstOrDefault(x =>
@@ -147,8 +126,8 @@ internal class NodeGraphViewModel : ViewModelBase, INodeGraphHandler, IDisposabl
             Connections.Remove(connection);
         }
 
-        RemoveFromFramesNotPartOf(connection.InputNode);
-        RemoveFromFramesNotPartOf(connection.OutputNode);
+        UpdatesFramesPartOf(connection.InputNode);
+        UpdatesFramesPartOf(connection.OutputNode);
         
         var node = AllNodes.FirstOrDefault(x => x.Id == nodeId);
         if (node != null)
@@ -163,30 +142,41 @@ internal class NodeGraphViewModel : ViewModelBase, INodeGraphHandler, IDisposabl
         StructureTree.Update(this);
     }
 
-    private void RemoveFromFramesNotPartOf(INodeHandler node)
+    private void UpdatesFramesPartOf(INodeHandler node)
     {
-        var framesPartOf =
-            Frames.OfType<NodeZoneViewModel>().Where(x => x.Nodes.Contains(node));
-
-        foreach (var frame in framesPartOf)
+        var lastKnownFramesPartOf = node.Frames.OfType<NodeZoneViewModel>().ToList();
+        var currentlyPartOf = new List<NodeZoneViewModel>();
+        
+        node.TraverseBackwards(x =>
         {
-            var stillConnected = false;
+            if (x is IPairNodeEndViewModel)
+                return false;
 
-            node.TraverseBackwards(x =>
-            {
-                if (x is not IPairNodeStartViewModel || frame.Start != x)
-                {
-                    return true;
-                }
+            if (x is not IPairNodeStartViewModel)
+                return true;
 
-                stillConnected = true;
-                return false;
-            });
+            var zone = Frames
+                .OfType<NodeZoneViewModel>()
+                .First(z => z.Start == x);
 
-            if (!stillConnected)
-            {
-                frame.Nodes.Remove(node);
-            }
+            currentlyPartOf.Add(zone);
+
+            return true;
+        });
+
+        foreach (var frame in currentlyPartOf)
+        {
+            if (!frame.Nodes.Contains(node))
+                frame.Nodes.Add(node);
+            
+            if (!node.Frames.Contains(frame))
+                node.Frames.Add(frame);
+        }
+
+        foreach (var removedFrom in lastKnownFramesPartOf.Except(currentlyPartOf))
+        {
+            removedFrom.Nodes.Remove(node);
+            node.Frames.Remove(removedFrom);
         }
     }
 

src/PixiEditor/ViewModels/Nodes/NodeViewModel.cs

@@ -436,6 +436,7 @@ internal abstract class NodeViewModel : ObservableObject, INodeHandler
         }
     }
 
+    public List<NodeFrameViewModelBase> Frames { get; } = [];
 
     public virtual void Dispose()
     {