From 8965d365e62aec3c458fca9585e2a7d6d5c0fa50 Mon Sep 17 00:00:00 2001
From: Thor Brigsted <thorbrigsted@gmail.com>
Date: Tue, 21 Apr 2020 15:31:07 +0200
Subject: [PATCH] Add [DisallowMultipleNodes] for limiting how many nodes of a
 certain type can exist on any graph

---
 Scripts/Editor/NodeEditorAction.cs |  9 +++++++++
 Scripts/Editor/NodeGraphEditor.cs  | 12 +++++++++++-
 Scripts/Node.cs                    | 10 ++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/Scripts/Editor/NodeEditorAction.cs b/Scripts/Editor/NodeEditorAction.cs
index 8e840af..f9a81fb 100644
--- a/Scripts/Editor/NodeEditorAction.cs
+++ b/Scripts/Editor/NodeEditorAction.cs
@@ -439,6 +439,15 @@ namespace XNodeEditor {
             for (int i = 0; i < nodes.Length; i++) {
                 XNode.Node srcNode = nodes[i];
                 if (srcNode == null) continue;
+
+                // Check if user is allowed to add more of given node type
+                XNode.Node.DisallowMultipleNodesAttribute disallowAttrib;
+                Type nodeType = srcNode.GetType();
+                if (NodeEditorUtilities.GetAttrib(nodeType, out disallowAttrib)) {
+                    int typeCount = graph.nodes.Count(x => x.GetType() == nodeType);
+                    if (typeCount >= disallowAttrib.max) continue;
+                }
+
                 XNode.Node newNode = graphEditor.CopyNode(srcNode);
                 substitutes.Add(srcNode, newNode);
                 newNode.position = srcNode.position + offset;
diff --git a/Scripts/Editor/NodeGraphEditor.cs b/Scripts/Editor/NodeGraphEditor.cs
index f8e1eda..3806d39 100644
--- a/Scripts/Editor/NodeGraphEditor.cs
+++ b/Scripts/Editor/NodeGraphEditor.cs
@@ -62,7 +62,17 @@ namespace XNodeEditor {
                 string path = GetNodeMenuName(type);
                 if (string.IsNullOrEmpty(path)) continue;
 
-                menu.AddItem(new GUIContent(path), false, () => {
+                // Check if user is allowed to add more of given node type
+                XNode.Node.DisallowMultipleNodesAttribute disallowAttrib;
+                bool disallowed = false;
+                if (NodeEditorUtilities.GetAttrib(type, out disallowAttrib)) {
+                    int typeCount = target.nodes.Count(x => x.GetType() == type);
+                    if (typeCount >= disallowAttrib.max) disallowed = true;
+                }
+
+                // Add node entry to context menu
+                if (disallowed) menu.AddItem(new GUIContent(path), false, null);
+                else menu.AddItem(new GUIContent(path), false, () => {
                     XNode.Node node = CreateNode(type, pos);
                     NodeEditorWindow.current.AutoConnect(node);
                 });
diff --git a/Scripts/Node.cs b/Scripts/Node.cs
index 9d2721b..8e1a670 100644
--- a/Scripts/Node.cs
+++ b/Scripts/Node.cs
@@ -334,6 +334,16 @@ namespace XNode {
             }
         }
 
+        [AttributeUsage(AttributeTargets.Class, AllowMultiple = false)]
+        public class DisallowMultipleNodesAttribute : Attribute {
+            public int max;
+            /// <summary> Prevents Node of the same type (or subtype) to be added more than once (configurable) to a NodeGraph </summary>
+            /// <param name="max"> How many nodes to allow. Defaults to 1. </param>
+            public DisallowMultipleNodesAttribute(int max = 1) {
+                this.max = max;
+            }
+        }
+
         [AttributeUsage(AttributeTargets.Class, AllowMultiple = false)]
         public class NodeTintAttribute : Attribute {
             public Color color;