From 46ac70a73490fa635a4826378f75c7852f8b4828 Mon Sep 17 00:00:00 2001
From: Vera Aguilera Puerto <6766154+Zumorica@users.noreply.github.com>
Date: Wed, 23 Mar 2022 14:05:10 +0100
Subject: [PATCH] Fix and refactor SolutionContainerSystem.RemoveEachReagent
 (#7245)

---
 .../Botany/Components/PlantHolderComponent.cs |  6 +--
 .../EntitySystems/SolutionContainerSystem.cs  | 40 +++++++++----------
 .../Chemistry/Components/Solution.cs          | 27 ++++++++-----
 3 files changed, 39 insertions(+), 34 deletions(-)

diff --git a/Content.Server/Botany/Components/PlantHolderComponent.cs b/Content.Server/Botany/Components/PlantHolderComponent.cs
index 6c5bf394a4..1022edc164 100644
--- a/Content.Server/Botany/Components/PlantHolderComponent.cs
+++ b/Content.Server/Botany/Components/PlantHolderComponent.cs
@@ -569,10 +569,10 @@ namespace Content.Server.Botany.Components
             else
             {
                 var amt = FixedPoint2.New(1);
-                foreach (var reagent in solutionSystem.RemoveEachReagent(Owner, solution, amt))
+                foreach (var (reagentId, quantity) in solutionSystem.RemoveEachReagent(Owner, solution, amt))
                 {
-                    var reagentProto = _prototypeManager.Index<ReagentPrototype>(reagent);
-                    reagentProto.ReactionPlant(Owner, new Solution.ReagentQuantity(reagent, amt), solution);
+                    var reagentProto = _prototypeManager.Index<ReagentPrototype>(reagentId);
+                    reagentProto.ReactionPlant(Owner, new Solution.ReagentQuantity(reagentId, quantity), solution);
                 }
             }
 
diff --git a/Content.Server/Chemistry/EntitySystems/SolutionContainerSystem.cs b/Content.Server/Chemistry/EntitySystems/SolutionContainerSystem.cs
index cd2ab59e8f..5d90eb098d 100644
--- a/Content.Server/Chemistry/EntitySystems/SolutionContainerSystem.cs
+++ b/Content.Server/Chemistry/EntitySystems/SolutionContainerSystem.cs
@@ -290,35 +290,33 @@ public sealed partial class SolutionContainerSystem : EntitySystem
         return solutionsMgr.Solutions[name];
     }
 
-    public string[] RemoveEachReagent(EntityUid uid, Solution solution, FixedPoint2 quantity)
+    /// <summary>
+    ///     Removes an amount from all reagents in a solution, adding it to a new solution.
+    /// </summary>
+    /// <param name="uid">The entity containing the solution.</param>
+    /// <param name="solution">The solution to remove reagents from.</param>
+    /// <param name="quantity">The amount to remove from every reagent in the solution.</param>
+    /// <returns>A new solution containing every removed reagent from the original solution.</returns>
+    public Solution RemoveEachReagent(EntityUid uid, Solution solution, FixedPoint2 quantity)
     {
-        var removedReagent = new string[solution.Contents.Count];
         if (quantity <= 0)
-            return Array.Empty<string>();
+            return new Solution();
 
-        var pos = 0;
-        for (var i = 0; i < solution.Contents.Count; i++)
+        var removedSolution = new Solution();
+
+        // RemoveReagent does a RemoveSwap, meaning we don't have to copy the list if we iterate it backwards.
+        for (var i = solution.Contents.Count-1; i >= 0; i--)
         {
-            var (reagentId, curQuantity) = solution.Contents[i];
-            removedReagent[pos++] = reagentId;
-            if (!_prototypeManager.TryIndex(reagentId, out ReagentPrototype? proto))
-                proto = new ReagentPrototype();
+            var (reagentId, _) = solution.Contents[i];
 
-            var newQuantity = curQuantity - quantity;
-            if (newQuantity <= 0)
-            {
-                solution.Contents.RemoveSwap(i);
-                solution.TotalVolume -= curQuantity;
-            }
-            else
-            {
-                solution.Contents[i] = new Solution.ReagentQuantity(reagentId, newQuantity);
-                solution.TotalVolume -= quantity;
-            }
+            var removedQuantity = solution.RemoveReagent(reagentId, quantity);
+
+            if(removedQuantity > 0)
+                removedSolution.AddReagent(reagentId, removedQuantity);
         }
 
         UpdateChemicals(uid, solution);
-        return removedReagent;
+        return removedSolution;
     }
 
     public FixedPoint2 GetReagentQuantity(EntityUid owner, string reagentId)
diff --git a/Content.Shared/Chemistry/Components/Solution.cs b/Content.Shared/Chemistry/Components/Solution.cs
index 73fbd9d40a..696c297669 100644
--- a/Content.Shared/Chemistry/Components/Solution.cs
+++ b/Content.Shared/Chemistry/Components/Solution.cs
@@ -167,34 +167,41 @@ namespace Content.Shared.Chemistry.Components
             return FixedPoint2.New(0);
         }
 
-        public void RemoveReagent(string reagentId, FixedPoint2 quantity)
+        /// <summary>
+        ///     Attempts to remove an amount of reagent from the solution.
+        /// </summary>
+        /// <param name="reagentId">The reagent to be removed.</param>
+        /// <param name="quantity">The amount of reagent to remove.</param>
+        /// <returns>How much reagent was actually removed. Zero if the reagent is not present on the solution.</returns>
+        public FixedPoint2 RemoveReagent(string reagentId, FixedPoint2 quantity)
         {
             if(quantity <= 0)
-                return;
+                return FixedPoint2.Zero;
 
             for (var i = 0; i < Contents.Count; i++)
             {
                 var reagent = Contents[i];
+
                 if(reagent.ReagentId != reagentId)
                     continue;
-                if (!IoCManager.Resolve<IPrototypeManager>().TryIndex(reagentId, out ReagentPrototype? proto))
-                    proto = new ReagentPrototype();
 
                 var curQuantity = reagent.Quantity;
                 var newQuantity = curQuantity - quantity;
+
                 if (newQuantity <= 0)
                 {
                     Contents.RemoveSwap(i);
                     TotalVolume -= curQuantity;
-                }
-                else
-                {
-                    Contents[i] = new ReagentQuantity(reagentId, newQuantity);
-                    TotalVolume -= quantity;
+                    return curQuantity;
                 }
 
-                return;
+                Contents[i] = new ReagentQuantity(reagentId, newQuantity);
+                TotalVolume -= quantity;
+                return quantity;
             }
+
+            // Reagent is not on the solution...
+            return FixedPoint2.Zero;
         }
 
         /// <summary>