Add compreg methods to entitymanager (#5655)

RELEASE-NOTES.md

@@ -39,7 +39,7 @@ END TEMPLATE-->
 
 ### New features
 
-*None yet*
+* Add EntityManager overloads for ComponentRegistration that's faster than the generic methods.
 
 ### Bugfixes
 

Robust.Benchmarks/EntityManager/HasComponentBenchmark.cs

@@ -0,0 +1,96 @@
+using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Engines;
+using JetBrains.Annotations;
+using Robust.Shared.Analyzers;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Map;
+using Robust.UnitTesting.Server;
+
+namespace Robust.Benchmarks.EntityManager;
+
+[Virtual]
+public partial class HasComponentBenchmark
+{
+    private static readonly Consumer Consumer = new();
+
+    private ISimulation _simulation = default!;
+    private IEntityManager _entityManager = default!;
+
+    private ComponentRegistration _compReg = default!;
+
+    private A _dummyA = new();
+
+    [UsedImplicitly]
+    [Params(1, 10, 100, 1000)]
+    public int N;
+
+    [GlobalSetup]
+    public void GlobalSetup()
+    {
+        _simulation = RobustServerSimulation
+            .NewSimulation()
+            .RegisterComponents(f => f.RegisterClass<A>())
+            .InitializeInstance();
+
+        _entityManager = _simulation.Resolve<IEntityManager>();
+        var map = _simulation.CreateMap().Uid;
+        var coords = new EntityCoordinates(map, default);
+        _compReg = _entityManager.ComponentFactory.GetRegistration(typeof(A));
+
+        for (var i = 0; i < N; i++)
+        {
+            var uid = _entityManager.SpawnEntity(null, coords);
+            _entityManager.AddComponent<A>(uid);
+        }
+    }
+
+    [Benchmark]
+    public void HasComponentGeneric()
+    {
+        for (var i = 2; i <= N+1; i++)
+        {
+            var uid = new EntityUid(i);
+            var result = _entityManager.HasComponent<A>(uid);
+            Consumer.Consume(result);
+        }
+    }
+
+    [Benchmark]
+    public void HasComponentCompReg()
+    {
+        for (var i = 2; i <= N+1; i++)
+        {
+            var uid = new EntityUid(i);
+            var result = _entityManager.HasComponent(uid, _compReg);
+            Consumer.Consume(result);
+        }
+    }
+
+    [Benchmark]
+    public void HasComponentType()
+    {
+        for (var i = 2; i <= N+1; i++)
+        {
+            var uid = new EntityUid(i);
+            var result = _entityManager.HasComponent(uid, typeof(A));
+            Consumer.Consume(result);
+        }
+    }
+
+    [Benchmark]
+    public void HasComponentGetType()
+    {
+        for (var i = 2; i <= N+1; i++)
+        {
+            var uid = new EntityUid(i);
+            var type = _dummyA.GetType();
+            var result = _entityManager.HasComponent(uid, type);
+            Consumer.Consume(result);
+        }
+    }
+
+    [ComponentProtoName("A")]
+    public sealed partial class A : Component
+    {
+    }
+}

Robust.Shared/GameObjects/EntityManager.Components.cs

@@ -197,17 +197,23 @@ namespace Robust.Shared.GameObjects
             {
                 var reg = _componentFactory.GetRegistration(name);
 
-                if (HasComponent(target, reg.Type))
+                if (removeExisting)
                 {
-                    if (!removeExisting)
-                        continue;
-
-                    RemoveComponent(target, reg.Type, metadata);
+                    var comp = _componentFactory.GetComponent(reg);
+                    _serManager.CopyTo(entry.Component, ref comp, notNullableOverride: true);
+                    AddComponentInternal(target, comp, reg, overwrite: true, metadata: metadata);
                 }
+                else
+                {
+                    if (HasComponent(target, reg))
+                    {
+                        continue;
+                    }
 
-                var comp = _componentFactory.GetComponent(reg);
-                _serManager.CopyTo(entry.Component, ref comp, notNullableOverride: true);
-                AddComponent(target, comp, metadata: metadata);
+                    var comp = _componentFactory.GetComponent(reg);
+                    _serManager.CopyTo(entry.Component, ref comp, notNullableOverride: true);
+                    AddComponentInternal(target, comp, reg, overwrite: false, metadata: metadata);
+                }
             }
         }
 
@@ -315,6 +321,22 @@ namespace Robust.Shared.GameObjects
             AddComponentInternal(uid, component, overwrite, false, metadata);
         }
 
+        private void AddComponentInternal<T>(
+            EntityUid uid,
+            T component,
+            ComponentRegistration compReg,
+            bool overwrite = false,
+            MetaDataComponent? metadata = null) where T : IComponent
+        {
+            if (!MetaQuery.Resolve(uid, ref metadata, false))
+                throw new ArgumentException($"Entity {uid} is not valid.", nameof(uid));
+
+            DebugTools.Assert(component.Owner == default);
+            component.Owner = uid;
+
+            AddComponentInternal(uid, component, compReg, overwrite, skipInit: false, metadata);
+        }
+
         private void AddComponentInternal<T>(EntityUid uid, T component, bool overwrite, bool skipInit, MetaDataComponent? metadata) where T : IComponent
         {
             if (!MetaQuery.ResolveInternal(uid, ref metadata, false))
@@ -731,6 +753,14 @@ namespace Robust.Shared.GameObjects
             return uid.HasValue && HasComponent<T>(uid.Value);
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        [Pure]
+        public bool HasComponent(EntityUid uid, ComponentRegistration reg)
+        {
+            var dict = _entTraitArray[reg.Idx.Value];
+            return dict.TryGetValue(uid, out var comp) && !comp.Deleted;
+        }
+
         /// <inheritdoc />
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         [Pure]
@@ -943,6 +973,23 @@ namespace Robust.Shared.GameObjects
             return false;
         }
 
+        /// <inheritdoc />
+        public bool TryGetComponent(EntityUid uid, ComponentRegistration reg, [NotNullWhen(true)] out IComponent? component)
+        {
+            var dict = _entTraitArray[reg.Idx.Value];
+            if (dict.TryGetValue(uid, out var comp))
+            {
+                if (!comp.Deleted)
+                {
+                    component = comp;
+                    return true;
+                }
+            }
+
+            component = null;
+            return false;
+        }
+
         /// <inheritdoc />
         public bool TryGetComponent(EntityUid uid, Type type, [NotNullWhen(true)] out IComponent? component)
         {

Robust.Shared/GameObjects/IEntityManager.Components.cs

@@ -176,6 +176,14 @@ namespace Robust.Shared.GameObjects
         /// <returns>True if the entity has the component type, otherwise false.</returns>
         bool HasComponent<T>([NotNullWhen(true)] EntityUid? uid) where T : IComponent;
 
+        /// <summary>
+        ///     Checks if the entity has a component type.
+        /// </summary>
+        /// <param name="uid">Entity UID to check.</param>
+        /// <param name="reg">The component registration to check for.</param>
+        /// <returns>True if the entity has the component type, otherwise false.</returns>
+        bool HasComponent(EntityUid uid, ComponentRegistration reg);
+
         /// <summary>
         ///     Checks if the entity has a component type.
         /// </summary>
@@ -294,6 +302,15 @@ namespace Robust.Shared.GameObjects
         /// <returns>If the component existed in the entity.</returns>
         bool TryGetComponent<T>([NotNullWhen(true)] EntityUid? uid, [NotNullWhen(true)] out T? component) where T : IComponent?;
 
+        /// <summary>
+        ///     Returns the component of a specific type.
+        /// </summary>
+        /// <param name="uid">Entity UID to check.</param>
+        /// <param name="reg">The component registration to check for.</param>
+        /// <param name="component">Component of the specified type (if exists).</param>
+        /// <returns>If the component existed in the entity.</returns>
+        bool TryGetComponent(EntityUid uid, ComponentRegistration reg, [NotNullWhen(true)] out IComponent? component);
+
         /// <summary>
         ///     Returns the component of a specific type.
         /// </summary>