Cache ItemSizePrototype in StorageSystem (#22611)

Content.Shared/Storage/EntitySystems/SharedStorageSystem.cs

@@ -1,3 +1,4 @@
+using System.Collections.Frozen;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Content.Shared.ActionBlocker;
@@ -51,8 +52,13 @@ public abstract class SharedStorageSystem : EntitySystem
     [ValidatePrototypeId<ItemSizePrototype>]
     public const string DefaultStorageMaxItemSize = "Normal";
 
+    private ItemSizePrototype _defaultStorageMaxItemSize = default!;
+
     public bool CheckingCanInsert;
 
+    private readonly List<ItemSizePrototype> _sortedSizes = new();
+    private FrozenDictionary<string, ItemSizePrototype> _nextSmallest = FrozenDictionary<string, ItemSizePrototype>.Empty;
+
     /// <inheritdoc />
     public override void Initialize()
     {
@@ -61,6 +67,7 @@ public abstract class SharedStorageSystem : EntitySystem
         _itemQuery = GetEntityQuery<ItemComponent>();
         _stackQuery = GetEntityQuery<StackComponent>();
         _xformQuery = GetEntityQuery<TransformComponent>();
+        _prototype.PrototypesReloaded += OnPrototypesReloaded;
 
         SubscribeLocalEvent<StorageComponent, ComponentInit>(OnComponentInit, before: new[] { typeof(SharedContainerSystem) });
         SubscribeLocalEvent<StorageComponent, GetVerbsEvent<UtilityVerb>>(AddTransferVerbs);
@@ -82,6 +89,39 @@ public abstract class SharedStorageSystem : EntitySystem
         SubscribeAllEvent<StorageSetItemLocationEvent>(OnSetItemLocation);
         SubscribeAllEvent<StorageInsertItemIntoLocationEvent>(OnInsertItemIntoLocation);
         SubscribeAllEvent<StorageRemoveItemEvent>(OnRemoveItem);
+        UpdatePrototypeCache();
+    }
+
+    public override void Shutdown()
+    {
+        _prototype.PrototypesReloaded -= OnPrototypesReloaded;
+    }
+
+    private void OnPrototypesReloaded(PrototypesReloadedEventArgs args)
+    {
+        if (args.ByType.ContainsKey(typeof(ItemSizePrototype))
+            || (args.Removed?.ContainsKey(typeof(ItemSizePrototype)) ?? false))
+        {
+            UpdatePrototypeCache();
+        }
+    }
+
+    private void UpdatePrototypeCache()
+    {
+        _defaultStorageMaxItemSize = _prototype.Index<ItemSizePrototype>(DefaultStorageMaxItemSize);
+        _sortedSizes.Clear();
+        _sortedSizes.AddRange(_prototype.EnumeratePrototypes<ItemSizePrototype>());
+        _sortedSizes.Sort();
+
+        var nextSmallest = new KeyValuePair<string, ItemSizePrototype>[_sortedSizes.Count];
+        for (var i = 0; i < _sortedSizes.Count; i++)
+        {
+            var k = _sortedSizes[i].ID;
+            var v = _sortedSizes[Math.Max(i - 1, 0)];
+            nextSmallest[i] = new(k, v);
+        }
+
+        _nextSmallest = nextSmallest.ToFrozenDictionary();
     }
 
     private void OnComponentInit(EntityUid uid, StorageComponent storageComp, ComponentInit args)
@@ -601,7 +641,7 @@ public abstract class SharedStorageSystem : EntitySystem
             return true;
         }
 
-        var maxSize = ItemSystem.GetSizePrototype(GetMaxItemSize((uid, storageComp)));
+        var maxSize = GetMaxItemSize((uid, storageComp));
         if (ItemSystem.GetSizePrototype(item.Size) > maxSize)
         {
             reason = "comp-storage-too-big";
@@ -609,7 +649,7 @@ public abstract class SharedStorageSystem : EntitySystem
         }
 
         if (TryComp<StorageComponent>(insertEnt, out var insertStorage)
-            && ItemSystem.GetSizePrototype(GetMaxItemSize((insertEnt, insertStorage))) >= maxSize)
+            && GetMaxItemSize((insertEnt, insertStorage)) >= maxSize)
         {
             reason = "comp-storage-too-big";
             return false;
@@ -1009,25 +1049,21 @@ public abstract class SharedStorageSystem : EntitySystem
         return sum;
     }
 
-    public ProtoId<ItemSizePrototype> GetMaxItemSize(Entity<StorageComponent?> uid)
+    public ItemSizePrototype GetMaxItemSize(Entity<StorageComponent?> uid)
     {
         if (!Resolve(uid, ref uid.Comp))
-            return DefaultStorageMaxItemSize;
+            return _defaultStorageMaxItemSize;
 
         // If we specify a max item size, use that
         if (uid.Comp.MaxItemSize != null)
-            return uid.Comp.MaxItemSize.Value;
+            return _prototype.Index(uid.Comp.MaxItemSize.Value);
 
         if (!_itemQuery.TryGetComponent(uid, out var item))
-            return DefaultStorageMaxItemSize;
+            return _defaultStorageMaxItemSize;
-        var size = ItemSystem.GetSizePrototype(item.Size);
 
         // if there is no max item size specified, the value used
-        // is one below the item size of the storage entity, clamped at ItemSize.Tiny
+        // is one below the item size of the storage entity.
-        var sizes = _prototype.EnumeratePrototypes<ItemSizePrototype>().ToList();
+        return _nextSmallest[item.Size];
-        sizes.Sort();
-        var currentSizeIndex = sizes.IndexOf(size);
-        return sizes[Math.Max(currentSizeIndex - 1, 0)].ID;
     }
 
     private void OnStackCountChanged(EntityUid uid, MetaDataComponent component, StackCountChangedEvent args)