weh

metalgearsloth · Aug 31, 2024 · 8e96788 · 8e96788
1 parent 474ac3a
commit 8e96788
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 118 deletions.
diff --git a/Robust.Shared/GameObjects/EntityManager.Components.cs b/Robust.Shared/GameObjects/EntityManager.Components.cs
@@ -3,6 +3,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
@@ -257,7 +258,7 @@ public void AddComponent(
             MetaDataComponent? metadata = null)
         {
             var copy = _componentFactory.GetComponent(entry);
-            AddComponent(uid, copy, overwrite, metadata);
+            AddComponentInternal(uid, copy, false, metadata);
         }
 
         /// <inheritdoc />
@@ -1123,27 +1124,11 @@ public ComponentQueryEnumerator ComponentQueryEnumerator(ComponentRegistry regis
         {
             if (registry.Count == 0)
             {
-                return new ComponentQueryEnumerator(new Dictionary<EntityUid, IComponent>());
+                return new ComponentQueryEnumerator(_world, QueryDescription.Null);
             }
 
-            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);
+            var query = new QueryDescription(registry.GetTypes());
+            return new ComponentQueryEnumerator(_world, query);
         }
 
         public AllEntityQueryEnumerator<TComp1> AllEntityQueryEnumerator<TComp1>()
@@ -1634,121 +1619,53 @@ internal bool ResolveInternal(EntityUid uid, [NotNullWhen(true)] ref TComp1? com
     #region ComponentRegistry Query
 
     /// <summary>
-    /// Returns entities that match the ComponentRegistry.
+    /// Non-generic version of <see cref="AllEntityQueryEnumerator{TComp1}"/>
     /// </summary>
-    public struct CompRegistryEntityEnumerator : IDisposable
+    public struct ComponentQueryEnumerator
     {
-        private IEntityManager _entManager;
-
-        private Dictionary<EntityUid, IComponent>.Enumerator _traitDict;
-        private ComponentRegistry _registry;
+        private QueryDescription _desc;
+        private ArchChunkEnumerator _chunkEnumerator;
+        private int _index;
 
-        public CompRegistryEntityEnumerator(
-            IEntityManager entManager,
-            Dictionary<EntityUid, IComponent> traitDict, ComponentRegistry registry)
+        public ComponentQueryEnumerator(
+            World world,
+            QueryDescription desc)
         {
-            _entManager = entManager;
-            _traitDict = traitDict.GetEnumerator();
-            _registry = registry;
+            _desc = desc;
+            var query = world.Query(desc);
+            _chunkEnumerator = query.ChunkIterator(world).GetEnumerator();
+            if (_chunkEnumerator.MoveNext())
+            {
+                _index = _chunkEnumerator.Current.Size;
+            }
         }
 
         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)
+                if (--_index < 0)
                 {
-                    idx++;
-
-                    // First one is us
-                    if (idx == 0)
-                        continue;
-
-                    if (!_entManager.TryGetComponent(current.Key, comp.Value.Component.GetType(), out var nextComp) ||
-                        nextComp.Deleted)
+                    if (!_chunkEnumerator.MoveNext())
                     {
-                        found = false;
-                        break;
+                        uid = EntityUid.Invalid;
+                        return false;
                     }
-                }
 
-                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;
+                    _index = _chunkEnumerator.Current.Size - 1;
                 }
 
-                var current = _traitDict.Current;
-
-                if (current.Value.Deleted)
+                // Deletion check moment
+                foreach (var comp in _desc.All)
                 {
-                    continue;
+                    if (((IComponent)_chunkEnumerator.Current.Get(_index, comp)!).Deleted)
+                        return MoveNext(out uid);
                 }
 
-                uid = current.Key;
-                comp1 = current.Value;
+                uid = _chunkEnumerator.Current.EntityReference(_index);
                 return true;
             }
         }
-
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public bool MoveNext([NotNullWhen(true)] out IComponent? comp1)
-        {
-            return MoveNext(out _, out comp1);
-        }
-
-        public void Dispose()
-        {
-            _traitDict.Dispose();
-        }
     }
     #endregion
 

diff --git a/Robust.Shared/GameObjects/IEntityManager.Components.cs b/Robust.Shared/GameObjects/IEntityManager.Components.cs
@@ -434,11 +434,6 @@ public partial interface IEntityManager
         /// </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/Prototypes/EntityPrototype.cs b/Robust.Shared/Prototypes/EntityPrototype.cs
@@ -454,6 +454,19 @@ public ComponentRegistry(Dictionary<string, EntityPrototype.ComponentRegistryEnt
         {
         }
 
+        public ComponentType[] GetTypes()
+        {
+            var arr = new ComponentType[Count];
+            var idx = 0;
+
+            foreach (var comp in Values)
+            {
+                arr[idx++] = comp.Component.GetType();
+            }
+
+            return arr;
+        }
+
         public bool TryGetComponent(string componentName, [NotNullWhen(true)] out IComponent? component)
         {
             var success = TryGetValue(componentName, out var comp);