Nerfed/Nerfed.Editor/Systems/EditorHierarchyWindow.cs

using ImGuiNET;
using MoonTools.ECS;
using Nerfed.Editor.Components;
using Nerfed.Runtime.Components;

namespace Nerfed.Editor.Systems
{
    // Window that draws entities.
    internal class EditorHierarchyWindow : MoonTools.ECS.DebugSystem
    {
        private const ImGuiTreeNodeFlags baseFlags = ImGuiTreeNodeFlags.OpenOnArrow | ImGuiTreeNodeFlags.OpenOnDoubleClick | ImGuiTreeNodeFlags.SpanAvailWidth;

        private readonly Filter rootEntitiesWithTransformFilter;
        //private readonly Filter rootEntitiesFilterBroken;
        private readonly Filter rootEntitiesFilter;

        private readonly EditorHierachySelectionSystem hierachySelectionSystem;

        public EditorHierarchyWindow(World world) : base(world)
        {
            rootEntitiesWithTransformFilter = FilterBuilder.Include<LocalTransform>().Exclude<Child>().Build();

            // TODO: this doesn't work.
            //rootEntitiesFilterBroken = FilterBuilder.Exclude<Child>().Build();

            // Maybe the parent/child functions should add a root component when not being a child.
            rootEntitiesFilter = FilterBuilder.Include<Root>().Build();

            // Maybe instead of a root, if we need a component that is always on an entity and has some use we could create something like a VersionComponent which only hold an int.
            // The version would update each time something changes on the entity.
            // Or a EditorComponent, just a component that always gets added when in editor mode.

            hierachySelectionSystem = new EditorHierachySelectionSystem(world);
        }

        public override void Update(TimeSpan delta)
        {
            ImGui.Begin("Hierarchy");

            ImGuiTreeNodeFlags flags = baseFlags;
            flags |= ImGuiTreeNodeFlags.DefaultOpen;

            if (ImGui.TreeNodeEx("World", flags))
            {
                foreach (Entity entity in rootEntitiesWithTransformFilter.Entities)
                {
                    DrawEntityAndChildren(entity);
                }

                foreach (Entity entity in rootEntitiesFilter.Entities)
                {
                    DrawEntityAndChildren(entity);
                }

                ImGui.TreePop();
            }

            ImGui.End();

            hierachySelectionSystem.Update(delta);
        }

        private void DrawEntityAndChildren(in Entity entity)
        {
            ImGuiTreeNodeFlags flags = baseFlags;

            if (!World.HasInRelation<ChildParentRelation>(entity))
            {
                flags |= ImGuiTreeNodeFlags.Leaf;
            }

            if (World.Has<SelectedInHierachy>(entity))
            {
                flags |= ImGuiTreeNodeFlags.Selected;
            }

            if (ImGui.TreeNodeEx($"{entity.ID} | {GetTag(entity)}", flags))
            {
                if (ImGui.IsItemClicked() && !ImGui.IsItemToggledOpen())
                {
                    World.Set(entity, new ClickedInHierarchy());
                }

                ReverseSpanEnumerator<Entity> childEntities = World.InRelations<ChildParentRelation>(entity);
                foreach (Entity childEntity in childEntities)
                {
                    DrawEntityAndChildren(childEntity);
                }

                ImGui.TreePop();
            }
        }

        // System for handling the selected entities in the hierachy.
        private class EditorHierachySelectionSystem : MoonTools.ECS.System
        {
            private readonly Filter selectedEntities;
            private readonly Filter clickedEntities;

            public EditorHierachySelectionSystem(World world) : base(world)
            {
                selectedEntities = FilterBuilder.Include<SelectedInHierachy>().Build();
                clickedEntities = FilterBuilder.Include<ClickedInHierarchy>().Build();
            }

            public override void Update(TimeSpan delta)
            {
                ImGuiIOPtr io = ImGui.GetIO();

                if (!clickedEntities.Empty && !io.KeyCtrl)
                {
                    foreach (Entity entity in selectedEntities.Entities)
                    {
                        Remove<SelectedInHierachy>(entity);
                    }
                }

                foreach (Entity entity in clickedEntities.Entities)
                {
                    // Unselect.
                    if (Has<SelectedInHierachy>(entity))
                    {
                        Remove<SelectedInHierachy>(entity);
                    }
                    // Select.
                    else
                    {
                        Set(entity, new SelectedInHierachy());
                    }

                    Remove<ClickedInHierarchy>(entity);
                }
            }
        }
    }
}