From 0e840d816ace750525329094d9bb2b92d8eb230f Mon Sep 17 00:00:00 2001
From: PJB3005 <pieterjan.briers+git@gmail.com>
Date: Sat, 18 Jan 2025 00:50:22 +0100
Subject: [PATCH] Make GasMixture enumerable

I noticed that enumerating gases is frequently done in an annoying way with Enum.GetValues. So I made it better. Now GasMixture is IEnumerable<(Gas gas, float moles)> and it just works.
---
 Content.Shared/Atmos/GasMixture.cs           | 45 ++++++++++++++++++--
 Content.Tests/Shared/Atmos/GasMixtureTest.cs | 30 +++++++++++++
 2 files changed, 72 insertions(+), 3 deletions(-)
 create mode 100644 Content.Tests/Shared/Atmos/GasMixtureTest.cs

diff --git a/Content.Shared/Atmos/GasMixture.cs b/Content.Shared/Atmos/GasMixture.cs
index 0f1efba976..deca8faaed 100644
--- a/Content.Shared/Atmos/GasMixture.cs
+++ b/Content.Shared/Atmos/GasMixture.cs
@@ -1,4 +1,5 @@
-using System.Diagnostics.CodeAnalysis;
+using System.Collections;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Runtime.CompilerServices;
 using Content.Shared.Atmos.EntitySystems;
@@ -13,12 +14,12 @@ namespace Content.Shared.Atmos
     /// </summary>
     [Serializable]
     [DataDefinition]
-    public sealed partial class GasMixture : IEquatable<GasMixture>, ISerializationHooks
+    public sealed partial class GasMixture : IEquatable<GasMixture>, ISerializationHooks, IEnumerable<(Gas gas, float moles)>
     {
         public static GasMixture SpaceGas => new() {Volume = Atmospherics.CellVolume, Temperature = Atmospherics.TCMB, Immutable = true};
 
         // No access, to ensure immutable mixtures are never accidentally mutated.
-        [Access(typeof(SharedAtmosphereSystem), typeof(SharedAtmosDebugOverlaySystem), Other = AccessPermissions.None)]
+        [Access(typeof(SharedAtmosphereSystem), typeof(SharedAtmosDebugOverlaySystem), typeof(GasEnumerator), Other = AccessPermissions.None)]
         [DataField]
         public float[] Moles = new float[Atmospherics.AdjustedNumberOfGases];
 
@@ -249,6 +250,16 @@ namespace Content.Shared.Atmos
             return new GasMixtureStringRepresentation(TotalMoles, Temperature, Pressure, molesPerGas);
         }
 
+        GasEnumerator GetEnumerator()
+        {
+            return new GasEnumerator(this);
+        }
+
+        IEnumerator<(Gas gas, float moles)> IEnumerable<(Gas gas, float moles)>.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+
         public override bool Equals(object? obj)
         {
             if (obj is GasMixture mix)
@@ -289,6 +300,11 @@ namespace Content.Shared.Atmos
             return hashCode.ToHashCode();
         }
 
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+
         public GasMixture Clone()
         {
             if (Immutable)
@@ -302,5 +318,28 @@ namespace Content.Shared.Atmos
             };
             return newMixture;
         }
+
+        public struct GasEnumerator(GasMixture mixture) : IEnumerator<(Gas gas, float moles)>
+        {
+            private int _idx = -1;
+
+            public void Dispose()
+            {
+                // Nada.
+            }
+
+            public bool MoveNext()
+            {
+                return ++_idx < Atmospherics.TotalNumberOfGases;
+            }
+
+            public void Reset()
+            {
+                _idx = -1;
+            }
+
+            public (Gas gas, float moles) Current => ((Gas)_idx, mixture.Moles[_idx]);
+            object? IEnumerator.Current => Current;
+        }
     }
 }
diff --git a/Content.Tests/Shared/Atmos/GasMixtureTest.cs b/Content.Tests/Shared/Atmos/GasMixtureTest.cs
new file mode 100644
index 0000000000..5068d48581
--- /dev/null
+++ b/Content.Tests/Shared/Atmos/GasMixtureTest.cs
@@ -0,0 +1,30 @@
+using Content.Shared.Atmos;
+using NUnit.Framework;
+
+namespace Content.Tests.Shared.Atmos;
+
+[TestFixture, TestOf(typeof(GasMixture))]
+[Parallelizable(ParallelScope.All)]
+public sealed class GasMixtureTest
+{
+    [Test]
+    public void TestEnumerate()
+    {
+        var mixture = new GasMixture();
+        mixture.SetMoles(Gas.Oxygen, 20);
+        mixture.SetMoles(Gas.Nitrogen, 10);
+        mixture.SetMoles(Gas.Plasma, 80);
+
+        var expectedList = new (Gas, float)[Atmospherics.TotalNumberOfGases];
+        for (var i = 0; i < Atmospherics.TotalNumberOfGases; i++)
+        {
+            expectedList[i].Item1 = (Gas)i;
+        }
+
+        expectedList[(int)Gas.Oxygen].Item2 = 20f;
+        expectedList[(int)Gas.Nitrogen].Item2 = 10f;
+        expectedList[(int)Gas.Plasma].Item2 = 80f;
+
+        Assert.That(mixture, Is.EquivalentTo(expectedList));
+    }
+}