From 18c34595e2a1cc8ec4a33d3a333a91a37beb6fb5 Mon Sep 17 00:00:00 2001
From: metalgearsloth <comedian_vs_clown@hotmail.com>
Date: Tue, 20 Aug 2024 22:19:28 +1000
Subject: [PATCH] compquery

---
 .../GameObjects/EntityManager.Components.cs   | 149 ++++++++++++++++++
 .../GameObjects/IEntityManager.Components.cs  |  10 ++
 .../Systems/PrototypeReloadSystem.cs          |   5 +-
 3 files changed, 160 insertions(+), 4 deletions(-)

diff --git a/Robust.Shared/GameObjects/EntityManager.Components.cs b/Robust.Shared/GameObjects/EntityManager.Components.cs
index adb2a3ee2ec..29e370523f0 100644
--- a/Robust.Shared/GameObjects/EntityManager.Components.cs
+++ b/Robust.Shared/GameObjects/EntityManager.Components.cs
@@ -1135,6 +1135,34 @@ public NetComponentEnumerable GetNetComponents(EntityUid uid, MetaDataComponent?
             return comps;
         }
 
+        /// <inheritdoc />
+        public ComponentQueryEnumerator ComponentQueryEnumerator(ComponentRegistry registry)
+        {
+            if (registry.Count == 0)
+            {
+                return new ComponentQueryEnumerator(new Dictionary<EntityUid, IComponent>());
+            }
+
+            var comp1 = registry.First().Value;
+            var trait1 = _entTraitArray[_componentFactory.GetArrayIndex(comp1.Component.GetType())];
+
+            return new ComponentQueryEnumerator(trait1);
+        }
+
+        /// <inheritdoc />
+        public CompRegistryEntityEnumerator CompRegistryQueryEnumerator(ComponentRegistry registry)
+        {
+            if (registry.Count == 0)
+            {
+                return new CompRegistryEntityEnumerator(this, new Dictionary<EntityUid, IComponent>(), registry);
+            }
+
+            var comp1 = registry.First().Value;
+            var trait1 = _entTraitArray[_componentFactory.GetArrayIndex(comp1.Component.GetType())];
+
+            return new CompRegistryEntityEnumerator(this, trait1, registry);
+        }
+
         public AllEntityQueryEnumerator<TComp1> AllEntityQueryEnumerator<TComp1>()
         where TComp1 : IComponent
         {
@@ -1723,6 +1751,127 @@ internal bool ResolveInternal(EntityUid uid, [NotNullWhen(true)] ref TComp1? com
         #endregion
     }
 
+    #region ComponentRegistry Query
+
+    /// <summary>
+    /// Returns entities that match the ComponentRegistry.
+    /// </summary>
+    public struct CompRegistryEntityEnumerator : IDisposable
+    {
+        private IEntityManager _entManager;
+
+        private Dictionary<EntityUid, IComponent>.Enumerator _traitDict;
+        private ComponentRegistry _registry;
+
+        public CompRegistryEntityEnumerator(
+            IEntityManager entManager,
+            Dictionary<EntityUid, IComponent> traitDict, ComponentRegistry registry)
+        {
+            _entManager = entManager;
+            _traitDict = traitDict.GetEnumerator();
+            _registry = registry;
+        }
+
+        public bool MoveNext(out EntityUid uid)
+        {
+            while (true)
+            {
+                if (!_traitDict.MoveNext())
+                {
+                    uid = default;
+                    return false;
+                }
+
+                var current = _traitDict.Current;
+
+                if (current.Value.Deleted)
+                {
+                    continue;
+                }
+
+                var idx = -1;
+                var found = true;
+
+                foreach (var comp in _registry)
+                {
+                    idx++;
+
+                    // First one is us
+                    if (idx == 0)
+                        continue;
+
+                    if (!_entManager.TryGetComponent(current.Key, comp.Value.Component.GetType(), out var nextComp) ||
+                        nextComp.Deleted)
+                    {
+                        found = false;
+                        break;
+                    }
+                }
+
+                if (!found)
+                    continue;
+
+                uid = current.Key;
+                return true;
+            }
+        }
+
+        public void Dispose()
+        {
+            _traitDict.Dispose();
+        }
+    }
+
+    /// <summary>
+    /// Non-generic version of <see cref="AllEntityQueryEnumerator{TComp1}"/>
+    /// </summary>
+    public struct ComponentQueryEnumerator : IDisposable
+    {
+        private Dictionary<EntityUid, IComponent>.Enumerator _traitDict;
+
+        public ComponentQueryEnumerator(
+            Dictionary<EntityUid, IComponent> traitDict)
+        {
+            _traitDict = traitDict.GetEnumerator();
+        }
+
+        public bool MoveNext(out EntityUid uid, [NotNullWhen(true)] out IComponent? comp1)
+        {
+            while (true)
+            {
+                if (!_traitDict.MoveNext())
+                {
+                    uid = default;
+                    comp1 = default;
+                    return false;
+                }
+
+                var current = _traitDict.Current;
+
+                if (current.Value.Deleted)
+                {
+                    continue;
+                }
+
+                uid = current.Key;
+                comp1 = current.Value;
+                return true;
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool MoveNext([NotNullWhen(true)] out IComponent? comp1)
+        {
+            return MoveNext(out _, out comp1);
+        }
+
+        public void Dispose()
+        {
+            _traitDict.Dispose();
+        }
+    }
+    #endregion
+
     #region Query
 
     /// <summary>
diff --git a/Robust.Shared/GameObjects/IEntityManager.Components.cs b/Robust.Shared/GameObjects/IEntityManager.Components.cs
index 8e23583c23d..c371a3a055d 100644
--- a/Robust.Shared/GameObjects/IEntityManager.Components.cs
+++ b/Robust.Shared/GameObjects/IEntityManager.Components.cs
@@ -427,6 +427,16 @@ public partial interface IEntityManager
         /// </summary>
         List<(EntityUid Uid, T Component)> AllComponentsList<T>() where T : IComponent;
 
+        /// <summary>
+        /// <see cref="ComponentQueryEnumerator"/>
+        /// </summary>
+        public ComponentQueryEnumerator ComponentQueryEnumerator(ComponentRegistry registry);
+
+        /// <summary>
+        /// <see cref="CompRegistryQueryEnumerator"/>
+        /// </summary>
+        public CompRegistryEntityEnumerator CompRegistryQueryEnumerator(ComponentRegistry registry);
+
         AllEntityQueryEnumerator<TComp1> AllEntityQueryEnumerator<TComp1>()
             where TComp1 : IComponent;
 
diff --git a/Robust.Shared/GameObjects/Systems/PrototypeReloadSystem.cs b/Robust.Shared/GameObjects/Systems/PrototypeReloadSystem.cs
index f0a73252b94..50b2b43b582 100644
--- a/Robust.Shared/GameObjects/Systems/PrototypeReloadSystem.cs
+++ b/Robust.Shared/GameObjects/Systems/PrototypeReloadSystem.cs
@@ -1,7 +1,4 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using Robust.Shared.IoC;
+using Robust.Shared.IoC;
 using Robust.Shared.Prototypes;
 
 namespace Robust.Shared.GameObjects;