BehaviorTree / Editor / GraphElements / Graph.cs
Graph.cs
Raw
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEditor.Experimental.GraphView;
using UnityEngine;
using UnityEngine.UIElements;

namespace AI.BT
{
    [System.Serializable]
    public class NodeCache
    {
        public List<BTNode> Nodes = new List<BTNode>();
    }

    public class Graph : UnityEditor.Experimental.GraphView.GraphView
    {
        public new class UxmlFactory : UxmlFactory<Graph, GraphView.UxmlTraits> { }

        [SerializeField]
        private NodeSearchWindow _nodeSearchWindow;

        [SerializeField]
        private EditorWindow _editorWindow;

        public Action<NodeElement> OnNodeSelected;

        public Action<NodeElement> OnNodeUnselected;

        public Action<NodeElement> OnNodeDoubleClicked;

        public Action<System.Type, Vector2> OnEntrySelected { get; internal set; }

        public Action<BTNode> OnNodeRemoved { get; internal set; }

        public Action<BTNode, BTNode, int> OnEdgeRemoved { get; internal set; }

        public Action<BTNode> OnNodeMoved { get; internal set; }

        public Action<BTNode, BTNode, int> OnEdgeCreated { get; internal set; }

        public Action<BTNode> OnNodePasted { get; internal set; }

        [SerializeField]
        private MiniMap _miniMap;

        public Graph() 
        {
            Insert(0, new GridBackground());

            this.AddManipulator(new ContentZoomer() { minScale = .01f, maxScale = 3.0f });
            this.AddManipulator(new ContentDragger());
            this.AddManipulator(new SelectionDragger());
            this.AddManipulator(new RectangleSelector());

            SetupZoom(.01f, 3.0f);

            nodeCreationRequest += CreateSearchWindow;
            

            serializeGraphElements += CopyElements;
            unserializeAndPaste += PasteElements;

            CreateMiniMap();
        }

        public void Initialize(EditorWindow window)
        {
            _editorWindow = window;
        }
        
        public void PopulateView(List<BTNode> nodes)
        {
            graphViewChanged -= OnGraphViewChanged;
            DeleteElements(graphElements);
            graphViewChanged += OnGraphViewChanged;

            //create node elements
            nodes.ForEach(n =>
            {
                CreateNodeElement(n);
            });

            //create edges
            nodes.ForEach(n =>
            {
                if(n.GetChildren() is List<BTNode> children)
                {
                    int childIndex = 0;
                    children.ForEach(c =>
                    {
                        
                        if (c == null) return;

                        NodeElement parent = FindNodeElement(n);
                        NodeElement child = FindNodeElement(c);

                        if(parent == null || child == null) return;

                        Edge edge = parent.m_Outputs[child.NodeRef.PortIndex].ConnectTo<FlowEdge>(child.m_Inputs[childIndex]);
                        AddElement(edge);
                    });
                }
            });
        }

        internal NodeElement FindNodeElement(BTNode n)
        {
            return GetNodeByGuid(n.m_ID) as NodeElement;
        }

        internal GraphElement CreateNodeElement(BTNode node)
        {
            var element =  new NodeElement(node);
            element.OnNodeSelected += OnNodeSelected;
            element.OnNodeDoubleClicked += OnNodeDoubleClicked;
            element.OnNodeUnselected += OnNodeUnselected;
            AddElement(element);
            return element;
        }

        private void PasteElements(string operationName, string data)
        {
            if (data == "") return;

            NodeCache nodeCache = new NodeCache();
            EditorJsonUtility.FromJsonOverwrite(data, nodeCache);

            foreach (BTNode node in nodeCache.Nodes)
            {
                OnNodePasted?.Invoke(ScriptableObject.Instantiate(node));
            }
        }

        private string CopyElements(IEnumerable<GraphElement> elements)
        {
            NodeCache nodeCache = new NodeCache();

            foreach (GraphElement element in elements)
            {
                if(element is NodeElement nodeelement)
                {
                    nodeCache.Nodes.Add(nodeelement.NodeRef);
                }
            }

            return EditorJsonUtility.ToJson(nodeCache);
        }

        private GraphViewChange OnGraphViewChanged(GraphViewChange graphViewChange)
        {
            if(graphViewChange.elementsToRemove != null)
            {
                graphViewChange.elementsToRemove.ForEach(e =>
                {
                    if(e is NodeElement node)
                    {
                        OnNodeRemoved?.Invoke(node.NodeRef);
                    }

                    if(e is Edge edge)
                    {
                        NodeElement parent = edge.output.node as NodeElement;
                        NodeElement child = edge.input.node as NodeElement;

                        OnEdgeRemoved?.Invoke(parent.NodeRef, child.NodeRef, ((PortElement)edge.output).m_Index);
                    }
                });
            }

            if(graphViewChange.edgesToCreate != null)
            {
                graphViewChange.edgesToCreate.ForEach(edge =>
                {
                    NodeElement parent = edge.output.node as NodeElement;
                    NodeElement child = edge.input.node as NodeElement;

                    OnEdgeCreated?.Invoke(parent.NodeRef, child.NodeRef, ((PortElement)edge.output).m_Index);
                });
            }

            if(graphViewChange.movedElements != null)
            {
                graphViewChange.movedElements.ForEach(e =>
                {
                    if(e is NodeElement node)
                    {
                       OnNodeMoved?.Invoke(node.NodeRef);
                    }
                });
            }

            return graphViewChange;
        }

        public override List<Port> GetCompatiblePorts(Port startPort, NodeAdapter nodeAdapter)
        {
           return ports.Where(endport =>
           {
               BTNode node = (startPort.node as NodeElement).NodeRef;
               bool is_notsame_direction = startPort.direction != endport.direction;
               bool is_not_same_node = startPort.node != endport.node;
               bool is_assignable = endport.portType.IsAssignableFrom(node.GetType());
               bool is_same = endport.portType == node.GetType();
               return is_notsame_direction && is_not_same_node && (is_assignable || is_same);

           }).ToList();
        }

        public override void BuildContextualMenu(ContextualMenuPopulateEvent evt)
        {

            base.BuildContextualMenu(evt);

            if(evt.target is Graph)
            {
                foreach (var type in TypeCache.GetTypesDerivedFrom<BTNode>())
                {
                    CreateMenuGroups(evt.menu, type, evt.localMousePosition);
                }
            }
        }

        private void CreateMenuGroups(DropdownMenu menu, System.Type type, Vector2 mousePosition)
        {
            if(type.IsAbstract)
            {
                foreach(var derived in TypeCache.GetTypesDerivedFrom(type))
                {
                    CreateMenuGroups(menu, derived, mousePosition);
                }
                menu.AppendSeparator();
            }
            else
            {
                var displayname = BTNode.GetDisplayName(type);
                menu.AppendAction(displayname, a => OnEntrySelected?.Invoke(a.userData as System.Type, GetLocalMousePosition(this, mousePosition)),
                    a => DropdownMenuAction.Status.Normal, type);
            }
        }

        private void CreateSearchWindow(NodeCreationContext context)
        {
            if(_nodeSearchWindow == null)
            {
                _nodeSearchWindow = ScriptableObject.CreateInstance<NodeSearchWindow>();
                _nodeSearchWindow.OnEntrySelected += OnSearchWindowEntrySelected;
            }

            SearchWindow.Open(new SearchWindowContext(context.screenMousePosition), _nodeSearchWindow);
        }

        private static Vector2 GetLocalMousePosition(Graph graph, Vector2 screenMousePosition, bool searchWindow = false)
        {
            var world_position = screenMousePosition;

            if(searchWindow && graph._editorWindow)
            {
                world_position -= graph._editorWindow.position.position;
            }

            var local_position = graph.contentViewContainer.WorldToLocal(world_position);

            return local_position;
        }

        private void OnSearchWindowEntrySelected(Type type, Vector2 mousePosition)
        {
            OnEntrySelected?.Invoke(type, GetLocalMousePosition(this, mousePosition, true));
        }

        #region UI

        private void CreateMiniMap()
        {
            _miniMap = new MiniMap()
            {
                anchored = true,
                visible = false
            };

            _miniMap.SetPosition(new Rect(15, 30, 200, 150));

            Add(_miniMap );
        }

        internal void ToggleMinMap()
        {
            _miniMap.visible = !_miniMap.visible;
        }
        #endregion
    }
}