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; } }