From 32e1d1a3b5bac3b81ad4f580d1001459ed7e2556 Mon Sep 17 00:00:00 2001
From: DrSmugleaf <10968691+DrSmugleaf@users.noreply.github.com>
Date: Sun, 16 Jun 2024 04:20:54 -0700
Subject: [PATCH] Use EntityQuery for mob state system resolves (#29021)

---
 .../Systems/MobStateSystem.StateMachine.cs     |  7 ++++---
 Content.Shared/Mobs/Systems/MobStateSystem.cs  | 18 ++++++++----------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/Content.Shared/Mobs/Systems/MobStateSystem.StateMachine.cs b/Content.Shared/Mobs/Systems/MobStateSystem.StateMachine.cs
index 2fa522dea5..5928b3871f 100644
--- a/Content.Shared/Mobs/Systems/MobStateSystem.StateMachine.cs
+++ b/Content.Shared/Mobs/Systems/MobStateSystem.StateMachine.cs
@@ -16,7 +16,8 @@ public partial class MobStateSystem
     /// <returns>If the entity can be set to that MobState</returns>
     public bool HasState(EntityUid entity, MobState mobState, MobStateComponent? component = null)
     {
-        return Resolve(entity, ref component, false) && component.AllowedStates.Contains(mobState);
+        return _mobStateQuery.Resolve(entity, ref component, false) &&
+               component.AllowedStates.Contains(mobState);
     }
 
     /// <summary>
@@ -27,7 +28,7 @@ public partial class MobStateSystem
     /// <param name="origin">Entity that caused the state update (if applicable)</param>
     public void UpdateMobState(EntityUid entity, MobStateComponent? component = null, EntityUid? origin = null)
     {
-        if (!Resolve(entity, ref component))
+        if (!_mobStateQuery.Resolve(entity, ref component))
             return;
 
         var ev = new UpdateMobStateEvent {Target = entity, Component = component, Origin = origin};
@@ -46,7 +47,7 @@ public partial class MobStateSystem
     public void ChangeMobState(EntityUid entity, MobState mobState, MobStateComponent? component = null,
         EntityUid? origin = null)
     {
-        if (!Resolve(entity, ref component))
+        if (!_mobStateQuery.Resolve(entity, ref component))
             return;
 
         ChangeState(entity, component, mobState, origin: origin);
diff --git a/Content.Shared/Mobs/Systems/MobStateSystem.cs b/Content.Shared/Mobs/Systems/MobStateSystem.cs
index a3886dd42e..323efa2242 100644
--- a/Content.Shared/Mobs/Systems/MobStateSystem.cs
+++ b/Content.Shared/Mobs/Systems/MobStateSystem.cs
@@ -2,7 +2,6 @@ using Content.Shared.ActionBlocker;
 using Content.Shared.Administration.Logs;
 using Content.Shared.Mobs.Components;
 using Content.Shared.Standing;
-using Robust.Shared.GameStates;
 using Robust.Shared.Physics.Systems;
 using Robust.Shared.Timing;
 
@@ -20,9 +19,12 @@ public partial class MobStateSystem : EntitySystem
     [Dependency] private readonly IGameTiming _timing = default!;
     private ISawmill _sawmill = default!;
 
+    private EntityQuery<MobStateComponent> _mobStateQuery;
+
     public override void Initialize()
     {
         _sawmill = _logManager.GetSawmill("MobState");
+        _mobStateQuery = GetEntityQuery<MobStateComponent>();
         base.Initialize();
         SubscribeEvents();
     }
@@ -37,7 +39,7 @@ public partial class MobStateSystem : EntitySystem
     /// <returns>If the entity is alive</returns>
     public bool IsAlive(EntityUid target, MobStateComponent? component = null)
     {
-        if (!Resolve(target, ref component, false))
+        if (!_mobStateQuery.Resolve(target, ref component, false))
             return false;
         return component.CurrentState == MobState.Alive;
     }
@@ -50,7 +52,7 @@ public partial class MobStateSystem : EntitySystem
     /// <returns>If the entity is Critical</returns>
     public bool IsCritical(EntityUid target, MobStateComponent? component = null)
     {
-        if (!Resolve(target, ref component, false))
+        if (!_mobStateQuery.Resolve(target, ref component, false))
             return false;
         return component.CurrentState == MobState.Critical;
     }
@@ -63,7 +65,7 @@ public partial class MobStateSystem : EntitySystem
     /// <returns>If the entity is Dead</returns>
     public bool IsDead(EntityUid target, MobStateComponent? component = null)
     {
-        if (!Resolve(target, ref component, false))
+        if (!_mobStateQuery.Resolve(target, ref component, false))
             return false;
         return component.CurrentState == MobState.Dead;
     }
@@ -76,7 +78,7 @@ public partial class MobStateSystem : EntitySystem
     /// <returns>If the entity is Critical or Dead</returns>
     public bool IsIncapacitated(EntityUid target, MobStateComponent? component = null)
     {
-        if (!Resolve(target, ref component, false))
+        if (!_mobStateQuery.Resolve(target, ref component, false))
             return false;
         return component.CurrentState is MobState.Critical or MobState.Dead;
     }
@@ -89,14 +91,10 @@ public partial class MobStateSystem : EntitySystem
     /// <returns>If the entity is in an Invalid State</returns>
     public bool IsInvalidState(EntityUid target, MobStateComponent? component = null)
     {
-        if (!Resolve(target, ref component, false))
+        if (!_mobStateQuery.Resolve(target, ref component, false))
             return false;
         return component.CurrentState is MobState.Invalid;
     }
 
     #endregion
-
-    #region Private Implementation
-
-    #endregion
 }
-- 
2.51.2