Spooky stuff

Editor/Scripts/Editor/Attributes/NodeAttributes.cs

@@ -27,25 +27,28 @@ namespace AssetGraph.Core.Attributes
         }
     }
     
-    [AttributeUsage(AttributeTargets.Field)]
+    // Base type for Input and Output attributes
-    public class Input : Attribute
+    public abstract class PortAttribute : Attribute
     {
         public string Label { get; }
 
-        public Input(string _label)
+        protected PortAttribute(string label)
         {
-            Label = _label;
+            Label = label;
         }
     }
 
+    // Input attribute inherits from base type
     [AttributeUsage(AttributeTargets.Field)]
-    public class Output : Attribute
+    public class Input : PortAttribute
     {
-        public string Label { get; }
+        public Input(string label) : base(label) { }
-
-        public Output(string _label)
-        {
-            Label = _label;
     }
+
+    // Output attribute inherits from base type
+    [AttributeUsage(AttributeTargets.Field)]
+    public class Output : PortAttribute
+    {
+        public Output(string label) : base(label) { }
     }
 }

Editor/Scripts/Editor/Windows/AssetGraphDataEditor.cs

@@ -8,29 +8,24 @@ namespace AssetGraph.Core
     {
         public override void OnInspectorGUI()
         {
-            if (GUILayout.Button("Open"))
+            var assetGraphData = (AssetGraphData)target;
+
+            DrawButton("Open", () => AssetGraphEditorWindow.Open(assetGraphData));
+            DrawButton("Calculate Dependency Graph",
+                () =>
                 {
-                AssetGraphEditorWindow.Open((AssetGraphData)target);
+                    assetGraphData.runOrder =
+                        assetGraphData.GetExecutionOrder(assetGraphData.Nodes, assetGraphData.Connections, true);
+                });
+            DrawButton("Run Graph", assetGraphData.RunGraph);
         }
 
-            if(GUILayout.Button("Calculate Dependancy Graph"))
+        private void DrawButton(string label, System.Action action)
         {
-                var bleh = (AssetGraphData)target;
+            if (GUILayout.Button(label))
-                bleh.runOrder = bleh.GetExecutionOrder(bleh.Nodes, bleh.Connections, true);
-            }
-
-            if(GUILayout.Button("Run Graph"))
             {
-                var bleh = (AssetGraphData)target;
+                action?.Invoke();
-                bleh.RunGraph();
             }
-            
-        }
-        
-        [MenuItem("Test/TEser", false, 1)]
-        public static void OpenBatchingWindow()
-        {
-            
         }
     }
 }

Editor/Scripts/Editor/Windows/AssetGraphEdgeConnector.cs

@@ -1,115 +1,141 @@
 using System.Collections.Generic;
 using ImageProcessingGraph.Editor;
 using UnityEditor.Experimental.GraphView;
-using UnityEditor.MemoryProfiler;
 using UnityEngine;
 using UnityEngine.UIElements;
 
 namespace AssetGraph.Core
 {
-
     public class AssetGraphEdgeConnector : IEdgeConnectorListener
     {
-        private GraphViewChange m_GraphViewChange;
+        private readonly GraphViewChange graphViewChange;
-        private List<Edge> m_EdgesToCreate;
+        private readonly List<Edge> edgesToCreate;
-        private List<GraphElement> m_EdgesToDelete;
+        private readonly List<GraphElement> edgesToDelete;
-
+        private readonly AssetGraphViewWindow window;
-        private AssetGraphViewWindow window;
 
         public AssetGraphEdgeConnector(AssetGraphViewWindow window)
         {
-            this.m_EdgesToCreate = new List<Edge>();
-            this.m_EdgesToDelete = new List<GraphElement>();
-            this.m_GraphViewChange.edgesToCreate = this.m_EdgesToCreate;
             this.window = window;
+            edgesToCreate = new List<Edge>();
+            edgesToDelete = new List<GraphElement>();
+            graphViewChange = new GraphViewChange
+            {
+                edgesToCreate = edgesToCreate
+            };
         }
 
         public void OnDropOutsidePort(Edge edge, Vector2 position)
         {
+            HandleSearchProvider(edge);
+
+            RemoveExistingConnections(edge);
+
+            if (edge.input?.node != null)
+            {
+                ((AssetGraphNodeEditor)edge.input.node)
+                    .ToggleExposedVariable((AssetGraphPort)edge.input, true);
+            }
+        }
+
+        public void OnDrop(GraphView graphView, Edge edge)
+        {
+            RemoveExistingConnections(edge);
+
+            if (edge.input?.node != null)
+            {
+                ((AssetGraphNodeEditor)edge.input.node)
+                    .ToggleExposedVariable((AssetGraphPort)edge.input, true);
+            }
+
+            UpdateConnections(graphView, edge);
+        }
+
+        private void HandleSearchProvider(Edge edge)
+        {
+            // Set the target for the search provider
             window.searchProvider.target = (VisualElement)window.focusController.focusedElement;
 
-            var mousePos =
+            var mousePos = window.ChangeCoordinatesTo(window, window.cachedMousePos - window.Window.position.position);
-                window.ChangeCoordinatesTo(window, window.cachedMousePos - window.Window.position.position);
             var graphMousePosition = window.contentViewContainer.WorldToLocal(mousePos);
 
-            var searchProviderInstance = ScriptableObject.CreateInstance<AssetGraphSearchProvider.CustomSearchProviderForEdge>();
+            var searchProviderInstance =
-            searchProviderInstance.Init(window.searchProvider, edge, (AssetGraphPort)edge.input, (AssetGraphPort)edge.output);
+                ScriptableObject.CreateInstance<AssetGraphSearchProvider.CustomSearchProviderForEdge>();
+            searchProviderInstance.Init(window.searchProvider, edge, (AssetGraphPort)edge.input,
+                (AssetGraphPort)edge.output);
 
             SearchWindow.Open(
                 new SearchWindowContext(GUIUtility.GUIToScreenPoint(Event.current.mousePosition)),
                 searchProviderInstance
             );
-    
-            // Remove connections as you did
-            List<GraphConnection> connections = new List<GraphConnection>();
-            foreach (var conn in window.asset.Connections)
-            {
-                if (conn.internalEdge == edge)
-                    connections.Add(conn);
-            }
-            foreach (var VARIABLE in connections)
-            {
-                window.asset.Connections.Remove(VARIABLE);
         }
 
-            if (edge.input != null)
+        private void RemoveExistingConnections(Edge edge)
-                if (edge.input.node != null)
-                    ((AssetGraphNodeEditor)edge.input.node).ToggleExposedVariable((AssetGraphPort)edge.input, true);
-        }
-
-
-
-        public void OnDrop(UnityEditor.Experimental.GraphView.GraphView graphView, Edge edge)
         {
-            List<GraphConnection> connections = new List<GraphConnection>();
+            // Collect and remove connections associated with the edge
-            
+            var connectionsToRemove = new List<GraphConnection>();
-            foreach (var conn in window.asset.Connections)
+            foreach (var connection in window.asset.Connections)
             {
-                if (conn.internalEdge == edge)
+                if (connection.internalEdge == edge)
                 {
-                    connections.Add(conn);
+                    connectionsToRemove.Add(connection);
                 }
             }
 
-            foreach (var VARIABLE in connections)
+            foreach (var connection in connectionsToRemove)
             {
-                window.asset.Connections.Remove(VARIABLE);
+                window.asset.Connections.Remove(connection);
+            }
         }
 
-            ((AssetGraphNodeEditor)edge.input.node).ToggleExposedVariable((AssetGraphPort)edge.input, true);
+        private void UpdateConnections(GraphView graphView, Edge edge)
+        {
+            // Clear the list of created and deleted edges
+            edgesToCreate.Clear();
+            edgesToDelete.Clear();
 
+            // Add the current edge to edgesToCreate
+            edgesToCreate.Add(edge);
 
-            this.m_EdgesToCreate.Clear();
+            // Handle single-capacity input ports
-            this.m_EdgesToCreate.Add(edge);
-            this.m_EdgesToDelete.Clear();
             if (edge.input.capacity == Port.Capacity.Single)
             {
-                foreach (Edge connection in edge.input.connections)
+                foreach (var connection in edge.input.connections)
                 {
                     if (connection != edge)
-                        this.m_EdgesToDelete.Add((GraphElement)connection);
+                    {
+                        edgesToDelete.Add(connection);
+                    }
                 }
             }
 
+            // Handle single-capacity output ports
             if (edge.output.capacity == Port.Capacity.Single)
             {
-                foreach (Edge connection in edge.output.connections)
+                foreach (var connection in edge.output.connections)
                 {
                     if (connection != edge)
-                        this.m_EdgesToDelete.Add((GraphElement)connection);
+                    {
+                        edgesToDelete.Add(connection);
+                    }
                 }
             }
 
-            if (this.m_EdgesToDelete.Count > 0)
+            // Delete unwanted connections
-                graphView.DeleteElements((IEnumerable<GraphElement>)this.m_EdgesToDelete);
+            if (edgesToDelete.Count > 0)
-            List<Edge> edgesToCreate = this.m_EdgesToCreate;
-            if (graphView.graphViewChanged != null)
-                edgesToCreate = graphView.graphViewChanged(this.m_GraphViewChange).edgesToCreate;
-            foreach (Edge edge1 in edgesToCreate)
             {
-                graphView.AddElement((GraphElement)edge1);
+                graphView.DeleteElements(edgesToDelete);
-                edge.input.Connect(edge1);
+            }
-                edge.output.Connect(edge1);
+
+            // Create new connections
+            var createdEdges = graphView.graphViewChanged != null
+                ? graphView.graphViewChanged(graphViewChange).edgesToCreate
+                : edgesToCreate;
+
+            foreach (var createdEdge in createdEdges)
+            {
+                graphView.AddElement(createdEdge);
+                edge.input.Connect(createdEdge);
+                edge.output.Connect(createdEdge);
             }
         }
     }

Editor/Scripts/Editor/Windows/AssetGraphEditorWindow.cs

@@ -10,43 +10,47 @@ namespace AssetGraph.Core
         [SerializeField] private AssetGraphData currentGraph;
         [SerializeField] private SerializedObject serializedObject;
         [SerializeField] private AssetGraphViewWindow currentView;
+
         public AssetGraphData CurrentGraph => currentGraph;
 
         public static void Open(AssetGraphData asset)
         {
-            var existingWindows = Resources.FindObjectsOfTypeAll<AssetGraphEditorWindow>();
+            foreach (var window in Resources.FindObjectsOfTypeAll<AssetGraphEditorWindow>())
-            foreach (var w in existingWindows)
             {
-                if (w.CurrentGraph == asset)
+                if (window.CurrentGraph == asset)
                 {
-                    w.Focus(); // 👁 focus the OG window
+                    window.Focus();
                     return;
                 }
             }
             
-            var window = CreateWindow<AssetGraphEditorWindow>(typeof(SceneView));
+            var newWindow = CreateWindow<AssetGraphEditorWindow>(typeof(SceneView));
-            window.titleContent = new GUIContent($"{asset.name}",
+            newWindow.Initialize(asset);
-                EditorGUIUtility.ObjectContent(null, typeof(AssetGraphData)).image);
-            window.Load(asset);
-            window.Focus();
         }
 
-
+        private void Initialize(AssetGraphData asset)
-
+        {
-        void OnEnable()
+            titleContent = new GUIContent(
+                asset.name,
+                EditorGUIUtility.ObjectContent(null, typeof(AssetGraphData)).image
+            );
+            Load(asset);
+            Focus();
+        }
+
+        private void OnEnable()
+        {
+            if (currentGraph != null)
             {
-            if(currentGraph != null)
                 DrawGraph();
             }
+        }
 
         private void OnGUI()
         {
             if (currentGraph != null)
             {
-                if(EditorUtility.IsDirty(currentGraph))
+                hasUnsavedChanges = EditorUtility.IsDirty(currentGraph);
-                    this.hasUnsavedChanges = true;
-                else
-                    this.hasUnsavedChanges = false;
             }
         }
 
@@ -56,19 +60,23 @@ namespace AssetGraph.Core
             DrawGraph();
         }
 
-        public void DrawGraph()
+        private void DrawGraph()
         {
             serializedObject = new SerializedObject(currentGraph);
-            currentView = new AssetGraphViewWindow(serializedObject, this);
+            currentView = new AssetGraphViewWindow(serializedObject, this)
-            currentView.graphViewChanged += OnChange;
+            {
+                graphViewChanged = OnGraphChanged
+            };
+
+            rootVisualElement.Clear();
             rootVisualElement.style.flexGrow = 1;
             rootVisualElement.Add(currentView);
         }
 
-        private GraphViewChange OnChange(GraphViewChange graphviewchange)
+        private GraphViewChange OnGraphChanged(GraphViewChange graphViewChange)
         {
             EditorUtility.SetDirty(currentGraph);
-            return graphviewchange;
+            return graphViewChange;
         }
     }
 }

Editor/Scripts/Editor/Windows/AssetGraphNodeEditor.cs

@@ -10,88 +10,115 @@ using UnityEditor.Experimental.GraphView;
 using UnityEditor.UIElements;
 using UnityEngine;
 using UnityEngine.UIElements;
+using Input = AssetGraph.Core.Attributes.Input;
 
 namespace AssetGraph.Core
 {
     public class AssetGraphNodeEditor : Node
     {
-        private AssetGraphNode graphNode;
+        private readonly StyleSheet defaultStyleSheet;
-        public AssetGraphNode GraphNode => graphNode;
+        private readonly StyleSheet errorStyleSheet;
+        private readonly NodeInfoAttribute info;
 
-        public List<Port> AllPorts = new List<Port>();
+        public AssetGraphNode GraphNode { get; }
-        public List<Port> InputPorts { get; }
+        public List<Port> AllPorts { get; } = new List<Port>();
-        public List<Port> OutputPorts { get; }
+        public List<Port> InputPorts { get; } = new List<Port>();
+        public List<Port> OutputPorts { get; } = new List<Port>();
 
-        public AssetGraphViewWindow window;
+        private readonly AssetGraphViewWindow window;
-        
-        private StyleSheet defaaStyleSheet;
-        private StyleSheet errorStyleSheet;
-        private NodeInfoAttribute info;
 
         public AssetGraphNodeEditor(AssetGraphNode node, AssetGraphViewWindow window)
         {
-            this.AddToClassList("image-node-visual");
             this.window = window;
-            graphNode = node;
+            GraphNode = node;
 
+            // Initialize node UI
+            info = InitializeNode(node);
 
-            Type typeInfo = node.GetType();
+            // Collect input and output ports
-            info = typeInfo.GetCustomAttribute<NodeInfoAttribute>();
+            CreatePorts(node.GetType());
-            title = info.Title;
-            this.name = typeInfo.Name;
 
+            // Load stylesheets
+            defaultStyleSheet = LoadStyleSheet("Node.uss", "Packages/com.chromium.imageprocessingrah/Node.uss");
+            errorStyleSheet = LoadStyleSheet("NodeError.uss", "Packages/com.chromium.imageprocessingrah/NodeError.uss");
+            styleSheets.Add(defaultStyleSheet);
 
-            string[] depths = info.MenuItem.Split('/');
+            // Set up error handling and cleanup
-            foreach (var depth in depths)
+            SetUpErrorHandling();
-            {
-                this.AddToClassList(depth.ToLower().Replace(' ', '-'));
         }
 
-            this.InputPorts = new List<Port>();
+        /// <summary>
-            this.OutputPorts = new List<Port>();
+        /// Initializes the node, including title, name, and class lists.
+        /// </summary>
+        private NodeInfoAttribute InitializeNode(AssetGraphNode node)
+        {
+            var typeInfo = node.GetType();
+            var nodeInfo = typeInfo.GetCustomAttribute<NodeInfoAttribute>();
 
-            List<AssetGraph.Core.Attributes.Input> inputs = new List<AssetGraph.Core.Attributes.Input>();
+            title = nodeInfo.Title;
-            List<FieldInfo> inputFieldInfo = new List<FieldInfo>();
+            name = typeInfo.Name;
-            List<FieldInfo> outputFieldInfo = new List<FieldInfo>();
 
-            FieldInfo[] fields = typeInfo.GetFields(BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public);
+            foreach (var depth in nodeInfo.MenuItem.Split('/'))
+            {
+                AddToClassList(depth.ToLower().Replace(' ', '-'));
+            }
+
+            return nodeInfo;
+        }
+
+        /// <summary>
+        /// Creates input and output ports for the node from its fields.
+        /// </summary>
+        private void CreatePorts(Type typeInfo)
+        {
+            var fields = typeInfo.GetFields(BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public);
+
+            // Collect input and output fields
+            var inputFields = fields.Where(f => f.GetCustomAttribute<Input>() != null).ToList();
+            var outputFields = fields.Where(f => f.GetCustomAttribute<Output>() != null).ToList();
+
+            // Create ports
+            CreateInputPorts(inputFields);
+            CreateOutputPorts(outputFields);
+        }
+
+        private void CreateInputPorts(List<FieldInfo> fields)
+        {
             foreach (var field in fields)
             {
-                if (field.GetCustomAttribute(typeof(AssetGraph.Core.Attributes.Input)) != null)
+                var port = CreatePort(field, true, field.GetCustomAttribute<Input>().Label);
-                {
+                InputPorts.Add(port);
-                    AssetGraph.Core.Attributes.Input input = field.GetCustomAttribute<AssetGraph.Core.Attributes.Input>();
+                AllPorts.Add(port);
-                    inputs.Add(input);
+
-                    inputFieldInfo.Add(field);
+                inputContainer.Add(port);
+                ExposeVariableToPort(port, field);
+                ((AssetGraphPort)port).fieldInfo = field;
+            }
         }
 
-                if (field.GetCustomAttribute(typeof(Output)) != null)
+        private void CreateOutputPorts(List<FieldInfo> fields)
         {
-                    Output output = field.GetCustomAttribute<Output>();
+            foreach (var field in fields)
-                    outputFieldInfo.Add(field);
+            {
+                var port = CreatePort(field, false, field.GetCustomAttribute<Output>().Label);
+                OutputPorts.Add(port);
+                AllPorts.Add(port);
+
+                outputContainer.Add(port);
+                ((AssetGraphPort)port).fieldInfo = field;
+            }
         }
 
+        private Port CreatePort(FieldInfo field, bool isInput, string label)
+        {
+            var port = AssetGraphPort.Create(window.edgeConnectorListener, isInput, field.FieldType);
+            if (!string.IsNullOrEmpty(label)) port.portName = label;
+            return port;
         }
 
-            CreateInputPorts(inputFieldInfo);
+        private void SetUpErrorHandling()
-            CreateOutputPorts(outputFieldInfo);
-            
-            defaaStyleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>("Assets/Unity Image Processing/Node.uss");
-            if (defaaStyleSheet == null)
         {
-                defaaStyleSheet = EditorGUIUtility.Load("Packages/com.chromium.imageprocessingrah/Node.uss") as StyleSheet;
+            GraphNode.onFailed += () => styleSheets.Add(errorStyleSheet);
-            }
-            styleSheets.Add(defaaStyleSheet);
-            
-            errorStyleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>("Assets/Unity Image Processing/NodeError.uss");
-            if (errorStyleSheet == null)
-            {
-                errorStyleSheet = EditorGUIUtility.Load("Packages/com.chromium.imageprocessingrah/NodeError.uss") as StyleSheet;
-            }
-
-            graphNode.onFailed += () =>
-            {
-                styleSheets.Add(errorStyleSheet);
-            };
 
             window.asset.OnRun += () =>
             {
@@ -100,221 +127,134 @@ namespace AssetGraph.Core
             };
         }
 
-        private void CreateInputPorts(List<FieldInfo> fields)
+        private StyleSheet LoadStyleSheet(string path, string fallbackPath)
         {
-            for (var index = 0; index < fields.Count; index++)
+            var styleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>($"Assets/Unity Image Processing/{path}");
-            {
+            return styleSheet ?? EditorGUIUtility.Load(fallbackPath) as StyleSheet;
-                var field = fields[index];
-                var port = AssetGraphPort.Create(window.edgeConnectorListener, true, field.FieldType);
-                
-                string label = field.GetCustomAttribute<Core.Attributes.Input>().Label;
-                if (label != "")
-                    port.portName = label; 
-                InputPorts.Add(port);
-                AllPorts.Add(port);
-                
-
-                inputContainer.Add(port);
-                ExposeVariableToPort(port, field);
-                port.fieldInfo = field;
-
-            }
         }
 
-        private void CreateOutputPorts(List<FieldInfo> fields)
+        /// <summary>
-        {
+        /// Exposes a variable on the port for editing when it's not connected.
-            for (var index = 0; index < fields.Count; index++)
+        /// </summary>
-            {
-                var field = fields[index];
-                var port = AssetGraphPort.Create(window.edgeConnectorListener, false, field.FieldType);
-                
-                string label = field.GetCustomAttribute<Output>().Label;
-                if (label != "")
-                    port.portName = label; 
-                OutputPorts.Add(port);
-                AllPorts.Add(port);
-                outputContainer.Add(port);
-                
-                port.fieldInfo = field;
-            }
-        }
-
-        // Exposes a variable on the port for editing when it's not connected
         public void ExposeVariableToPort(Port port, FieldInfo field)
         {
-            VisualElement NewElement = new VisualElement();
+            var variableContainer = ((AssetGraphPort)port).ExposedPropertyContainer;
-            var ExposedPropertyContainer = ((AssetGraphPort)port).ExposedPropertyContainer;
-            Type containerType = null;
 
-            if (ExposedPropertyContainer == null)
+            if (variableContainer == null)
             {
-                NewElement.name = "property-field-container";
+                var newElement = CreatePropertyFieldForType(field.FieldType, field.GetValue(GraphNode));
-                VisualElement the = CreatePropertyFieldForType(field.FieldType, field.GetValue(graphNode));
+                if (newElement != null)
+                {
+                    variableContainer = newElement;
+                    ((AssetGraphPort)port).ExposedPropertyContainer = newElement;
 
-                if(the != null)
+                    // Add the edited element to the port
-                    containerType = the.GetType();
+                    port.Add(newElement);
+                    RegisterFieldChangeCallback(field, newElement);
+                    variableContainer.style.display = DisplayStyle.Flex;
+                }
+            }
 
-                NewElement.Add(the);
+        }
-                ((AssetGraphPort)port).ExposedPropertyContainer = the;
+
-                ExposedPropertyContainer = ((AssetGraphPort)port).ExposedPropertyContainer;
+        private void RegisterFieldChangeCallback(FieldInfo field, VisualElement element)
+        {
+            // Register value change callbacks for various field types
+            switch (element)
+            {
+                case IntegerField intField:
+                    if (field.FieldType == typeof(GreyscaleValue))
+                    {
+                        // Handle GreyscaleValue specifically
+                        intField.RegisterValueChangedCallback(evt =>
+                        {
+                            var value = (GreyscaleValue)field.GetValue(GraphNode);
+
+                            if (intField is GreyscaleField greyscaleField && value != null)
+                            {
+                                if (evt.newValue > greyscaleField.minMax.Item2)
+                                    value.value = greyscaleField.minMax.Item2;
+                                else if (evt.newValue < greyscaleField.minMax.Item1)
+                                    value.value = greyscaleField.minMax.Item1;
+                                else
+                                    value.value = evt.newValue;
+
+                                field.SetValue(GraphNode, value);
+                                intField.SetValueWithoutNotify((int)value.value);
+
+                            }
+                        });
                     }
                     else
                     {
-                containerType = ExposedPropertyContainer.GetType();
+                        // Default case for IntegerField
+                        intField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
                     }
 
-            if (containerType == null)
+                    break;
-                return;
 
-            if (ExposedPropertyContainer.GetType() == typeof(IntegerField))
+                case FloatField floatField:
-            {
+                    floatField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
-                var intField = ExposedPropertyContainer as IntegerField;
+                    break;
-                intField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is FloatField floatField)
-            {
-                floatField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is Toggle boolField)
-            {
-                boolField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is TextField stringField)
-            {
-                stringField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is ColorField colorField)
-            {
-                colorField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is Vector3Field vector3Field)
-            {
-                vector3Field.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is Vector2Field vector2Field)
-            {
-                vector2Field.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is ObjectField objectField)
-            {
-                objectField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // Update the field with the new value
-                });
-            }
-            else if (ExposedPropertyContainer is EnumField enumField)
-            {
-                enumField.RegisterValueChangedCallback(evt =>
-                {
-                    field.SetValue(graphNode, evt.newValue); // 🎯 Update the field with the new enum value
-                });
-            }
-            else if (ExposedPropertyContainer.GetType() == typeof(GreyscaleField))
-            {
-                var greyscaleField = ExposedPropertyContainer as GreyscaleField;
-                greyscaleField.RegisterValueChangedCallback(evt =>
-                {
-                    var value = (GreyscaleValue)field.GetValue(graphNode);
 
-                    if (evt.newValue > greyscaleField.minMax.Item2)
+                case Toggle boolField:
-                        value.value = greyscaleField.minMax.Item2;
+                    boolField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
-                    else if (evt.newValue < greyscaleField.minMax.Item1) value.value = greyscaleField.minMax.Item1;
+                    break;
 
-                    value.value = evt.newValue;
+                case TextField stringField:
-                });
+                    stringField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
+                    break;
+
+                case ColorField colorField:
+                    colorField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
+                    break;
+
+                case Vector3Field vector3Field:
+                    vector3Field.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
+                    break;
+
+                case Vector2Field vector2Field:
+                    vector2Field.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
+                    break;
+
+                case EnumField enumField:
+                    enumField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
+                    break;
+
+                case ObjectField objectField:
+                    objectField.RegisterValueChangedCallback(evt => field.SetValue(GraphNode, evt.newValue));
+                    break;
             }
-            
-            port.Add(NewElement);
-        }
-
-        public void ToggleExposedVariable(Port port, bool value)
-        {
-            AssetGraphPort assetGraphPort = port as AssetGraphPort;
-            if(assetGraphPort.ExposedPropertyContainer != null)
-                assetGraphPort.ExposedPropertyContainer.style.display = value ? DisplayStyle.Flex : DisplayStyle.None;
         }
 
         private VisualElement CreatePropertyFieldForType(Type type, object value)
         {
-            if (type == typeof(int))
+            return type switch
             {
-                var intField = new IntegerField { value = (int)value };
+                { } when type == typeof(int) => new IntegerField { value = (int)value },
-                return intField;
+                { } when type == typeof(float) => new FloatField { value = (float)value },
+                { } when type == typeof(bool) => new Toggle { value = (bool)value },
+                { } when type == typeof(string) => new TextField { value = (string)value },
+                { } when type == typeof(Color) => new ColorField { value = (Color)value },
+                { } when type == typeof(Vector3) => new Vector3Field { value = (Vector3)value },
+                { } when type == typeof(Vector2) => new Vector2Field { value = (Vector2)value },
+                { } when type == typeof(GreyscaleValue) => new GreyscaleField { value = ((GreyscaleValue)value).value },
+                { } when type.IsEnum => new EnumField((Enum)value),
+                { } when typeof(UnityEngine.Object).IsAssignableFrom(type) => new ObjectField
+                {
+                    value = (UnityEngine.Object)value,
+                    objectType = type
+                },
+                _ => null
+            };
         }
-            else if (type == typeof(float))
-            {
-                var floatField = new FloatField { value = (float)value };
-                return floatField;
-            }
-            else if (type == typeof(bool))
-            {
-                var boolField = new Toggle { value = (bool)value };
-                return boolField;
-            }
-            else if (type == typeof(string))
-            {
-                var stringField = new TextField { value = (string)value };
-                return stringField;
-            }
-            else if (type == typeof(Color))
-            {
-                var colorField = new ColorField() { value = (Color)value };
-                return colorField;
-            }
-            else if (type == typeof(Vector3))
-            {
-                var vector3Field = new Vector3Field { value = (Vector3)value };
-                return vector3Field;
-            }
-            else if (type == typeof(Vector2))
-            {
-                var vector2Field = new Vector2Field { value = (Vector2)value };
-                return vector2Field;
-            }
-            else if (type == typeof(Texture2D))
-            {
-                var objectField = new ObjectField { value = (Texture2D)value, objectType = typeof(Texture2D) };
-                return objectField;
-            }
-            else if (type.IsEnum) // 💥✨ ENUMS, BABY! 💥✨
-            {
-                var enumField = new EnumField((Enum)value);
-                return enumField;
-            }
-            else if (type == typeof(GreyscaleValue))
-            {
 
-                var greyscaleValue = (GreyscaleValue)value;
+        public void ToggleExposedVariable(Port port, bool isVisible)
-                var intField = new GreyscaleField { value = (int)greyscaleValue.value };
-                return intField;
-            }
-            else if (typeof(UnityEngine.Object).IsAssignableFrom(type))
         {
-                var objectField = new ObjectField { value = (UnityEngine.Object)value, objectType = typeof(UnityEngine.Object) };
+            var assetGraphPort = port as AssetGraphPort;
-                return objectField;
+            if (assetGraphPort?.ExposedPropertyContainer != null)
-            }
+                assetGraphPort.ExposedPropertyContainer.style.display =
-            return null;
+                    isVisible ? DisplayStyle.Flex : DisplayStyle.None;
         }
 
         public override void BuildContextualMenu(ContextualMenuPopulateEvent evt)
@@ -326,16 +266,16 @@ namespace AssetGraph.Core
 
         private void OpenDocumentation(DropdownMenuAction obj)
         {
-            if (info.DocumentationationURL != null || !string.IsNullOrEmpty(info.DocumentationationURL))
+            if (!string.IsNullOrEmpty(info.DocumentationationURL))
             {
                 Application.OpenURL(info.DocumentationationURL);
             }
             else
             {
-                Debug.LogWarning($"No documentation URL provided for node {this.graphNode.asset.name}!");
+                Debug.LogWarning($"No documentation URL provided for node {GraphNode.asset.name}!");
             }
         }
 
-        public void SavePosition() => graphNode.SetPosition(GetPosition());
+        public void SavePosition() => GraphNode.SetPosition(GetPosition());
     }
 }

Editor/Scripts/Editor/Windows/AssetGraphSearchProvider.cs

@@ -25,8 +25,8 @@ namespace AssetGraph.Core
         {
             this.target = target;
             this.title = title;
-            this.ImportPortTypes = importPortTypes;
+            ImportPortTypes = importPortTypes;
-            this.ExportPortTypes = exportPortTypes;
+            ExportPortTypes = exportPortTypes;
             this.portID = portID;
         }
     }
@@ -39,6 +39,7 @@ namespace AssetGraph.Core
         public static List<SearchContextElement> elements;
         private Assembly[] assemblies;
 
+        // Create Search Tree for Nodes
         public List<SearchTreeEntry> CreateSearchTree(SearchWindowContext context)
         {
             List<SearchTreeEntry> tree = new List<SearchTreeEntry>();
@@ -49,76 +50,46 @@ namespace AssetGraph.Core
             foreach (var type in TypeCache.GetTypesWithAttribute<NodeInfoAttribute>())
             {
                 var attr = type.GetCustomAttribute<NodeInfoAttribute>();
-                NodeInfoAttribute info = attr as NodeInfoAttribute;
+                if (attr == null || string.IsNullOrEmpty(attr.MenuItem)) continue;
+
                 var node = Activator.CreateInstance(type);
-                if (string.IsNullOrEmpty(info.MenuItem)) continue;
 
-                var Fields = type.GetFields();
+                var fields = type.GetFields();
-                Dictionary<int, Type> ImportPortTypes = new Dictionary<int, Type>();
+                var importPortTypes = GetPortsByAttribute(fields, typeof(Core.Attributes.Input));
-                Dictionary<int, Type> OutputPortTypes = new Dictionary<int, Type>();
+                var exportPortTypes = GetPortsByAttribute(fields, typeof(Output));
 
-                foreach (var field in Fields)
+                elements.Add(new SearchContextElement(node, attr.MenuItem, importPortTypes, exportPortTypes, 0));
-                {
-                    if (field.GetCustomAttribute<Core.Attributes.Input>() != null)
-                        ImportPortTypes.Add(ImportPortTypes.Count, field.FieldType);
-                    if (field.GetCustomAttribute<Output>() != null)
-                        OutputPortTypes.Add(OutputPortTypes.Count, field.FieldType);
             }
 
-                elements.Add(new SearchContextElement(node, info.MenuItem, ImportPortTypes, OutputPortTypes, 0));
+            elements.Sort(SortElementsByTitle);
-            }
 
-            elements.Sort((entry1, entry2) =>
+            var groups = new HashSet<string>();
-            {
-                string[] splits1 = entry1.title.Split('/');
-                string[] splits2 = entry2.title.Split('/');
-
-                for (int i = 0; i < splits1.Length; i++)
-                {
-                    if (i >= splits2.Length) return 1;
-
-                    int value = splits1[i].CompareTo(splits2[i]);
-                    if (value != 0)
-                    {
-                        if (splits1.Length != splits2.Length && (i == splits1.Length - 1 || i == splits2.Length - 1))
-                            return splits1.Length < splits2.Length ? 1 : -1;
-
-                        return value;
-                    }
-                }
-
-                return 0;
-            });
-
-            List<string> groups = new List<string>();
 
             foreach (var element in elements)
             {
                 string[] entryTitle = element.title.Split('/');
-
+                string currentGroup = "";
-                string groupName = "";
 
                 for (int i = 0; i < entryTitle.Length - 1; i++)
                 {
-                    groupName += entryTitle[i];
+                    currentGroup += entryTitle[i];
-                    if (!groups.Contains(groupName))
+                    if (groups.Add(currentGroup))
+                        tree.Add(new SearchTreeGroupEntry(new GUIContent(entryTitle[i]), i + 1));
+
+                    currentGroup += '/';
+                }
+
+                tree.Add(new SearchTreeEntry(new GUIContent(entryTitle.Last()))
                 {
-                        tree.Add(new SearchTreeGroupEntry(new GUIContent(groupName), i + 1));
+                    level = entryTitle.Length,
-                        groups.Add(groupName);
+                    userData = element
-                    }
+                });
-
-                    groupName += '/';
-                }
-
-                SearchTreeEntry entry = new SearchTreeEntry(new GUIContent(entryTitle.Last()));
-                entry.level = entryTitle.Length;
-                entry.userData = element;
-                tree.Add(entry);
             }
 
             return tree;
         }
 
+        // Handle Selection from Search Tree
         public bool OnSelectEntry(SearchTreeEntry SearchTreeEntry, SearchWindowContext context)
         {
             var mousePos =
@@ -126,7 +97,6 @@ namespace AssetGraph.Core
             var graphMousePosition = graph.contentViewContainer.WorldToLocal(mousePos);
 
             SearchContextElement element = (SearchContextElement)SearchTreeEntry.userData;
-
             AssetGraphNode node = (AssetGraphNode)element.target;
             node.SetPosition(new Rect(graphMousePosition, new Vector2()));
             graph.Add(node);
@@ -135,13 +105,14 @@ namespace AssetGraph.Core
             return true;
         }
 
+        // Create Search Tree for Edge Connections
         public List<SearchTreeEntry> CreateSearchTreeForEdge(SearchWindowContext context, Edge edge)
         {
             var tree = new List<SearchTreeEntry>();
             tree.Add(new SearchTreeGroupEntry(new GUIContent("Compatible Nodes"), 0));
 
             elements = new List<SearchContextElement>();
-            var groups = new List<string>();
+            var groups = new HashSet<string>();
 
             var sourcePort = edge.output ?? edge.input;
             bool isSourceOutput = edge.output != null;
@@ -156,68 +127,43 @@ namespace AssetGraph.Core
             foreach (var type in TypeCache.GetTypesWithAttribute<NodeInfoAttribute>())
             {
                 var attr = type.GetCustomAttribute<NodeInfoAttribute>();
-                if (string.IsNullOrEmpty(attr.MenuItem)) continue;
+                if (attr == null || string.IsNullOrEmpty(attr.MenuItem)) continue;
 
                 var node = Activator.CreateInstance(type);
                 var fields = type.GetFields();
 
-                // Get all ports
+                var inputPorts = GetPortsByAttributeList(fields, typeof(Core.Attributes.Input));
-                var inputPorts = fields.Where(f => f.GetCustomAttribute<Core.Attributes.Input>() != null).ToList();
+                var outputPorts = GetPortsByAttributeList(fields, typeof(Output));
-                var outputPorts = fields.Where(f => f.GetCustomAttribute<Output>() != null).ToList();
 
-                // REVERSED LOGIC:
-                // If dragging from output, we want compatible INPUT ports
-                // If dragging from input, we want compatible OUTPUT ports
                 var compatiblePorts = isSourceOutput
-                    ? inputPorts.Where(f => f.FieldType.IsAssignableFrom(targetType)).ToList()
+                    ? inputPorts.Where(port => port.FieldType.IsAssignableFrom(targetType)).ToList()
-                    : outputPorts.Where(f => targetType.IsAssignableFrom(f.FieldType)).ToList();
+                    : outputPorts.Where(port => targetType.IsAssignableFrom(port.FieldType)).ToList();
 
-                if (compatiblePorts.Count == 0) continue;
+                if (!compatiblePorts.Any()) continue;
 
-                // Build group hierarchy
+                string[] menuPath = attr.MenuItem.Split('/');
-                var menuPath = attr.MenuItem.Split('/');
+                string currentGroupPath = "";
-                var currentGroupPath = "";
 
                 for (int i = 0; i < menuPath.Length - 1; i++)
                 {
                     currentGroupPath += menuPath[i];
-                    if (!groups.Contains(currentGroupPath))
+                    if (groups.Add(currentGroupPath))
-                    {
                         tree.Add(new SearchTreeGroupEntry(new GUIContent(menuPath[i]), i + 1));
-                        groups.Add(currentGroupPath);
-                    }
 
                     currentGroupPath += "/";
                 }
 
-                // // Add node entry
+                foreach (var compatiblePort in compatiblePorts)
-                // var nodeTitle = $"{type.Name}: {menuPath.Last()}";
-                // tree.Add(new SearchTreeEntry(new GUIContent(nodeTitle))
-                // {
-                //     level = menuPath.Length,
-                //     userData = new SearchContextElement(
-                //         node,
-                //         nodeTitle,
-                //         inputPorts.ToDictionary(f => inputPorts.IndexOf(f), f => f.FieldType),
-                //         outputPorts.ToDictionary(f => outputPorts.IndexOf(f), f => f.FieldType))
-                // });
-
-                // Add compatible port entries
-                foreach (var portField in compatiblePorts)
                 {
-                    var portList = portField.GetCustomAttribute<Core.Attributes.Input>() != null ? inputPorts : outputPorts;
+                    tree.Add(new SearchTreeEntry(new GUIContent($"{attr.Title}: {compatiblePort.Name}"))
-                    int portIndex = portList.IndexOf(portField);
-
-                    var portTitle = $"{attr.Title}: {portField.Name}";
-                    tree.Add(new SearchTreeEntry(new GUIContent(portTitle))
                     {
                         level = menuPath.Length,
                         userData = new SearchContextElement(
                             node,
-                            portTitle,
+                            attr.Title,
-                            inputPorts.ToDictionary(f => inputPorts.IndexOf(f), f => f.FieldType),
+                            inputPorts.ToDictionary(field => inputPorts.IndexOf(field), field => field.FieldType),
-                            outputPorts.ToDictionary(f => outputPorts.IndexOf(f), f => f.FieldType),
+                            outputPorts.ToDictionary(field => outputPorts.IndexOf(field), field => field.FieldType),
-                            portIndex)
+                            inputPorts.IndexOf(compatiblePort))
                     });
                 }
             }
@@ -225,6 +171,25 @@ namespace AssetGraph.Core
             return tree;
         }
 
+        private static Dictionary<int, Type> GetPortsByAttribute(FieldInfo[] fields, Type attributeType)
+        {
+            return fields
+                .Where(f => f.GetCustomAttribute(attributeType) != null)
+                .Select((f, index) => new { f, index })
+                .ToDictionary(a => a.index, a => a.f.FieldType);
+        }
+
+        private static List<FieldInfo> GetPortsByAttributeList(FieldInfo[] fields, Type attributeType)
+        {
+            return fields.Where(f => f.GetCustomAttribute(attributeType) != null).ToList();
+        }
+
+        private static int SortElementsByTitle(SearchContextElement a, SearchContextElement b)
+        {
+            return string.Compare(a.title, b.title, StringComparison.Ordinal);
+        }
+
+        // CustomSearchProviderForEdge
         public class CustomSearchProviderForEdge : ScriptableObject, ISearchWindowProvider
         {
             private AssetGraphSearchProvider original;
@@ -232,12 +197,10 @@ namespace AssetGraph.Core
             private AssetGraphPort inputPort;
             private AssetGraphPort outputPort;
 
-            public CustomSearchProviderForEdge(AssetGraphSearchProvider original, Edge edge, AssetGraphPort inputPort, AssetGraphPort outputPort)
+            public CustomSearchProviderForEdge(AssetGraphSearchProvider original, Edge edge, AssetGraphPort inputPort,
+                AssetGraphPort outputPort)
             {
-                this.original = original;
+                Init(original, edge, inputPort, outputPort);
-                this.edge = edge;
-                this.inputPort = inputPort;
-                this.outputPort = outputPort;
             }
 
             public void Init(AssetGraphSearchProvider original, Edge edge, AssetGraphPort inputPort,
@@ -256,35 +219,28 @@ namespace AssetGraph.Core
 
             public bool OnSelectEntry(SearchTreeEntry selectedEntry, SearchWindowContext context)
             {
-                var mousePos =
+                var mousePos = original.graph.ChangeCoordinatesTo(original.graph,
-                    original.graph.ChangeCoordinatesTo(original.graph, context.screenMousePosition - original.graph.Window.position.position);
+                    context.screenMousePosition - original.graph.Window.position.position);
                 var graphMousePosition = original.graph.contentViewContainer.WorldToLocal(mousePos);
 
-                SearchContextElement element = (SearchContextElement)selectedEntry.userData;
+                var element = (SearchContextElement)selectedEntry.userData;
+                var node = (AssetGraphNode)element.target;
 
-                AssetGraphNode node = (AssetGraphNode)element.target;
                 node.SetPosition(new Rect(graphMousePosition, new Vector2()));
                 original.graph.Add(node);
                 node.asset = original.graph.asset;
 
-
+                Edge newEdge = new Edge();
-                Edge edge = new Edge();
-                
                 if (inputPort != null)
-                {
+                    newEdge = inputPort.ConnectTo(original.graph.nodeDictionary[node.ID].OutputPorts[element.portID]);
-                    edge = inputPort.ConnectTo(original.graph.nodeDictionary[node.ID].OutputPorts[element.portID]);
-                }
                 else if (outputPort != null)
-                {
+                    newEdge = outputPort.ConnectTo(original.graph.nodeDictionary[node.ID].InputPorts[element.portID]);
-                    edge = outputPort.ConnectTo(original.graph.nodeDictionary[node.ID].InputPorts[element.portID]);
-                }
 
-                original.graph.CreateEdge(edge);
+                original.graph.CreateEdge(newEdge);
-                original.graph.AddElement((GraphElement)edge);
+                original.graph.AddElement(newEdge);
 
                 return true;
             }
         }
-
     }
 }