Add new EnsureEntity variants (space-wizards#4586)

metalgearsloth · Nov 20, 2023 · 202182e · 202182e
1 parent 96cb52e
commit 202182e
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 6 deletions.
diff --git a/Robust.Shared.CompNetworkGenerator/ComponentNetworkGenerator.cs b/Robust.Shared.CompNetworkGenerator/ComponentNetworkGenerator.cs
@@ -167,7 +167,7 @@ private static string GenerateSource(in GeneratorExecutionContext context, IName
                         getStateInit.Append($@"
                 {name} = GetNetEntitySet(component.{name}),");
                         handleStateSetters.Append($@"
-            component.{name} = EnsureEntitySet<{componentName}>(state.{name}, uid);");
+            EnsureEntitySet<{componentName}>(state.{name}, uid, component.{name});");
 
                         break;
                     case GlobalEntityUidListName:
@@ -177,7 +177,7 @@ private static string GenerateSource(in GeneratorExecutionContext context, IName
                         getStateInit.Append($@"
                 {name} = GetNetEntityList(component.{name}),");
                         handleStateSetters.Append($@"
-            component.{name} = EnsureEntityList<{componentName}>(state.{name}, uid);");
+            EnsureEntityList<{componentName}>(state.{name}, uid, component.{name});");
 
                         break;
                     default:

diff --git a/Robust.Shared/GameObjects/EntityManager.Network.cs b/Robust.Shared/GameObjects/EntityManager.Network.cs
@@ -256,8 +256,7 @@ public virtual EntityCoordinates EnsureCoordinates<T>(NetCoordinates netCoordina
     /// <inheritdoc />
     public HashSet<EntityUid> GetEntitySet(HashSet<NetEntity> netEntities)
     {
-        var entities = new HashSet<EntityUid>();
-        entities.EnsureCapacity(netEntities.Count);
+        var entities = new HashSet<EntityUid>(netEntities.Count);
 
         foreach (var netEntity in netEntities)
         {
@@ -292,6 +291,16 @@ public HashSet<EntityUid> EnsureEntitySet<T>(HashSet<NetEntity> netEntities, Ent
         return entities;
     }
 
+    public void EnsureEntitySet<T>(HashSet<NetEntity> netEntities, EntityUid callerEntity, HashSet<EntityUid> entities)
+    {
+        entities.Clear();
+        entities.EnsureCapacity(netEntities.Count);
+        foreach (var netEntity in netEntities)
+        {
+            entities.Add(EnsureEntity<T>(netEntity, callerEntity));
+        }
+    }
+
     /// <inheritdoc />
     public List<EntityUid> EnsureEntityList<T>(List<NetEntity> netEntities, EntityUid callerEntity)
     {
@@ -305,6 +314,16 @@ public List<EntityUid> EnsureEntityList<T>(List<NetEntity> netEntities, EntityUi
         return entities;
     }
 
+    public void EnsureEntityList<T>(List<NetEntity> netEntities, EntityUid callerEntity, List<EntityUid> entities)
+    {
+        entities.Clear();
+        entities.EnsureCapacity(netEntities.Count);
+        foreach (var netEntity in netEntities)
+        {
+            entities.Add(EnsureEntity<T>(netEntity, callerEntity));
+        }
+    }
+
     /// <inheritdoc />
     public List<EntityUid> GetEntityList(ICollection<NetEntity> netEntities)
     {

diff --git a/Robust.Shared/GameObjects/EntitySystem.Proxy.cs b/Robust.Shared/GameObjects/EntitySystem.Proxy.cs
@@ -1024,12 +1024,24 @@ protected HashSet<EntityUid> EnsureEntitySet<T>(HashSet<NetEntity> netEntities,
         return EntityManager.EnsureEntitySet<T>(netEntities, callerEntity);
     }
 
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    protected void EnsureEntitySet<T>(HashSet<NetEntity> netEntities, EntityUid callerEntity, HashSet<EntityUid> entities)
+    {
+        EntityManager.EnsureEntitySet<T>(netEntities, callerEntity, entities);
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     protected List<EntityUid> EnsureEntityList<T>(List<NetEntity> netEntities, EntityUid callerEntity)
     {
         return EntityManager.EnsureEntityList<T>(netEntities, callerEntity);
     }
 
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    protected void EnsureEntityList<T>(List<NetEntity> netEntities, EntityUid callerEntity, List<EntityUid> entities)
+    {
+        EntityManager.EnsureEntityList<T>(netEntities, callerEntity, entities);
+    }
+
     /// <summary>
     ///     Returns the <see cref="EntityUid"/> of a <see cref="NetEntity"/>. Returns <see cref="EntityUid.Invalid"/> if it doesn't exist.
     /// </summary>

diff --git a/Robust.Shared/Physics/Systems/SharedJointSystem.Relay.cs b/Robust.Shared/Physics/Systems/SharedJointSystem.Relay.cs
@@ -43,8 +43,7 @@ private void OnRelayHandleState(EntityUid uid, JointRelayTargetComponent compone
         if (args.Current is not JointRelayComponentState state)
             return;
 
-        component.Relayed.Clear();
-        component.Relayed.UnionWith(EnsureEntitySet<JointRelayTargetComponent>(state.Entities, uid));
+        EnsureEntitySet<JointRelayTargetComponent>(state.Entities, uid, component.Relayed);
     }
 
     private void OnRelayShutdown(EntityUid uid, JointRelayTargetComponent component, ComponentShutdown args)