From b506fdcf654864457adc0af0e1ac5ed3a8c99616 Mon Sep 17 00:00:00 2001
From: metalgearsloth <31366439+metalgearsloth@users.noreply.github.com>
Date: Sun, 27 Mar 2022 17:49:26 +1100
Subject: [PATCH] Reduce atmos allocs a bunch (#7228)

---
 .../AtmosObstructionEnumerator.cs             | 36 +++++++++++++++++++
 .../EntitySystems/AtmosphereSystem.Grid.cs    | 22 +++++++++---
 2 files changed, 54 insertions(+), 4 deletions(-)
 create mode 100644 Content.Server/Atmos/EntitySystems/AtmosObstructionEnumerator.cs

diff --git a/Content.Server/Atmos/EntitySystems/AtmosObstructionEnumerator.cs b/Content.Server/Atmos/EntitySystems/AtmosObstructionEnumerator.cs
new file mode 100644
index 0000000000..2a9d0729d6
--- /dev/null
+++ b/Content.Server/Atmos/EntitySystems/AtmosObstructionEnumerator.cs
@@ -0,0 +1,36 @@
+using System.Diagnostics.CodeAnalysis;
+using Content.Server.Atmos.Components;
+using Robust.Shared.Map;
+
+namespace Content.Server.Atmos.EntitySystems;
+
+public struct AtmosObstructionEnumerator
+{
+    private AnchoredEntitiesEnumerator _enumerator;
+    private EntityQuery<AirtightComponent> _query;
+
+    public AtmosObstructionEnumerator(AnchoredEntitiesEnumerator enumerator, EntityQuery<AirtightComponent> query)
+    {
+        _enumerator = enumerator;
+        _query = query;
+    }
+
+    public bool MoveNext([NotNullWhen(true)] out AirtightComponent? airtight)
+    {
+        if (!_enumerator.MoveNext(out var uid))
+        {
+            airtight = null;
+            return false;
+        }
+
+        // No rider, it makes it uglier.
+        // ReSharper disable once ConvertIfStatementToReturnStatement
+        if (!_query.TryGetComponent(uid.Value, out airtight))
+        {
+            // ReSharper disable once TailRecursiveCall
+            return MoveNext(out airtight);
+        }
+
+        return true;
+    }
+}
diff --git a/Content.Server/Atmos/EntitySystems/AtmosphereSystem.Grid.cs b/Content.Server/Atmos/EntitySystems/AtmosphereSystem.Grid.cs
index 4c6746dfc7..739f575f84 100644
--- a/Content.Server/Atmos/EntitySystems/AtmosphereSystem.Grid.cs
+++ b/Content.Server/Atmos/EntitySystems/AtmosphereSystem.Grid.cs
@@ -195,13 +195,25 @@ namespace Content.Server.Atmos.EntitySystems
         /// <returns></returns>
         public IEnumerable<AirtightComponent> GetObstructingComponents(IMapGrid mapGrid, Vector2i tile)
         {
-            foreach (var uid in mapGrid.GetAnchoredEntities(tile))
+            var airQuery = GetEntityQuery<AirtightComponent>();
+            var enumerator = mapGrid.GetAnchoredEntitiesEnumerator(tile);
+
+            while (enumerator.MoveNext(out var uid))
             {
-                if (TryComp<AirtightComponent>(uid, out var ac))
-                    yield return ac;
+                if (!airQuery.TryGetComponent(uid.Value, out var airtight)) continue;
+                yield return airtight;
             }
         }
 
+        public AtmosObstructionEnumerator GetObstructingComponentsEnumerator(IMapGrid mapGrid, Vector2i tile)
+        {
+            var ancEnumerator = mapGrid.GetAnchoredEntitiesEnumerator(tile);
+            var airQuery = GetEntityQuery<AirtightComponent>();
+
+            var enumerator = new AtmosObstructionEnumerator(ancEnumerator, airQuery);
+            return enumerator;
+        }
+
         private AtmosDirection GetBlockedDirections(IMapGrid mapGrid, Vector2i indices)
         {
             var value = AtmosDirection.Invalid;
@@ -717,7 +729,9 @@ namespace Content.Server.Atmos.EntitySystems
         {
             var directions = AtmosDirection.Invalid;
 
-            foreach (var obstructingComponent in GetObstructingComponents(mapGrid, tile))
+            var enumerator = GetObstructingComponentsEnumerator(mapGrid, tile);
+
+            while (enumerator.MoveNext(out var obstructingComponent))
             {
                 if (!obstructingComponent.AirBlocked)
                     continue;