diff --git a/Content.Server/Containers/ContainerExt.cs b/Content.Server/Containers/ContainerExt.cs
index b19b8d0a8d..9245b58375 100644
--- a/Content.Server/Containers/ContainerExt.cs
+++ b/Content.Server/Containers/ContainerExt.cs
@@ -4,21 +4,39 @@ namespace Content.Server.Containers
{
public static class ContainerExt
{
- public static int CountPrototypeOccurencesRecursive(this ContainerManagerComponent mgr, string prototypeId)
+ ///
+ /// Searches the entity, and the entities containers recursively for a prototypeId
+ ///
+ /// The entity to search
+ /// The prototypeId to find
+ /// Optional entity manager
+ /// True if entity is, or contains a prototype Id
+ public static bool ContainsPrototypeRecursive(this EntityUid entityUid, string prototypeId, IEntityManager? entityManager = null)
{
- var entMan = IoCManager.Resolve();
- int total = 0;
- foreach (var container in mgr.GetAllContainers())
+ IoCManager.Resolve(ref entityManager);
+ var metaQuery = entityManager.GetEntityQuery();
+ var managerQuery = entityManager.GetEntityQuery();
+ var stack = new Stack();
+ if (metaQuery.GetComponent(entityUid).EntityPrototype?.ID == prototypeId)
+ return true;
+ if (!managerQuery.TryGetComponent(entityUid, out var currentManager))
+ return false;
+ do
{
- foreach (var entity in container.ContainedEntities)
+ foreach (var container in currentManager.Containers.Values)
{
- if (entMan.GetComponent(entity).EntityPrototype?.ID == prototypeId) total++;
- if(!entMan.TryGetComponent(entity, out var component)) continue;
- total += component.CountPrototypeOccurencesRecursive(prototypeId);
+ foreach (var entity in container.ContainedEntities)
+ {
+ if (metaQuery.GetComponent(entity).EntityPrototype?.ID == prototypeId)
+ return true;
+ if (!managerQuery.TryGetComponent(entity, out var containerManager))
+ continue;
+ stack.Push(containerManager);
+ }
}
- }
+ } while (stack.TryPop(out currentManager));
- return total;
+ return false;
}
}
}
diff --git a/Content.Server/Objectives/Conditions/StealCondition.cs b/Content.Server/Objectives/Conditions/StealCondition.cs
index e2348743b4..98758cfccb 100644
--- a/Content.Server/Objectives/Conditions/StealCondition.cs
+++ b/Content.Server/Objectives/Conditions/StealCondition.cs
@@ -1,7 +1,6 @@
using Content.Server.Containers;
using Content.Server.Objectives.Interfaces;
using JetBrains.Annotations;
-using Robust.Shared.Containers;
using Robust.Shared.Prototypes;
using Robust.Shared.Serialization;
using Robust.Shared.Utility;
@@ -47,15 +46,9 @@ namespace Content.Server.Objectives.Conditions
{
get
{
- if (_mind?.OwnedEntity is not {Valid: true} owned) return 0f;
- if (!IoCManager.Resolve().TryGetComponent(owned, out var containerManagerComponent)) return 0f;
-
- // slightly ugly but fixing it would just be duplicating it with a different return value
- float count = containerManagerComponent.CountPrototypeOccurencesRecursive(_prototypeId);
- if (count >= 1)
- return 1;
- else
- return 0;
+ if (_mind?.OwnedEntity is not {Valid: true} owned)
+ return 0f;
+ return owned.ContainsPrototypeRecursive(_prototypeId)? 1f : 0f;
}
}