Replace count prototype (#10212)

* Replace count prototype

* Doc improvement

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)
+        /// <summary>
+        /// Searches the entity, and the entities containers recursively for a prototypeId
+        /// </summary>
+        /// <param name="entityUid">The entity to search</param>
+        /// <param name="prototypeId">The prototypeId to find</param>
+        /// <param name="entityManager">Optional entity manager</param>
+        /// <returns>True if entity is, or contains a prototype Id</returns>
+        public static bool ContainsPrototypeRecursive(this EntityUid entityUid, string prototypeId, IEntityManager? entityManager = null)
         {
-            var entMan = IoCManager.Resolve<IEntityManager>();
+            IoCManager.Resolve(ref entityManager);
-            int total = 0;
+            var metaQuery = entityManager.GetEntityQuery<MetaDataComponent>();
-            foreach (var container in mgr.GetAllContainers())
+            var managerQuery = entityManager.GetEntityQuery<ContainerManagerComponent>();
+            var stack = new Stack<ContainerManagerComponent>();
+            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<MetaDataComponent>(entity).EntityPrototype?.ID == prototypeId) total++;
+                    foreach (var entity in container.ContainedEntities)
-                    if(!entMan.TryGetComponent<ContainerManagerComponent?>(entity, out var component)) continue;
+                    {
-                    total += component.CountPrototypeOccurencesRecursive(prototypeId);
+                        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;
         }
     }
 }

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 (_mind?.OwnedEntity is not {Valid: true} owned)
-                if (!IoCManager.Resolve<IEntityManager>().TryGetComponent<ContainerManagerComponent?>(owned, out var containerManagerComponent)) return 0f;
+                    return 0f;
-
+                return owned.ContainsPrototypeRecursive(_prototypeId)? 1f : 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;
             }
         }