Add ReagentWhitelist to Injector component and system (#28262)

* Add reagent whitelist to Injector component and system

* Apply suggested Changes

* Remove LINQ function for performance

* Update Content.Server/Chemistry/EntitySystems/InjectorSystem.cs

Co-authored-by: Ed <96445749+TheShuEd@users.noreply.github.com>
---------

Co-authored-by: Ed <96445749+TheShuEd@users.noreply.github.com>

diff --git a/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs b/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs
index d94faa8123291b6b0ffd24f8ae03cc9060781183..c5c45daa5bd5818d07ee5eb280221b792b5f9131 100644 (file)
--- a/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs
+++ b/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs
@@ -327,8 +327,23 @@ public sealed class InjectorSystem : SharedInjectorSystem
             return false;
         }
 
+        var applicableTargetSolution = targetSolution.Comp.Solution;
+        // If a whitelist exists, remove all non-whitelisted reagents from the target solution temporarily
+        var temporarilyRemovedSolution = new Solution();
+        if (injector.Comp.ReagentWhitelist is { } reagentWhitelist)
+        {
+            string[] reagentPrototypeWhitelistArray = new string[reagentWhitelist.Count];
+            var i = 0;
+            foreach (var reagent in reagentWhitelist)
+            {
+                reagentPrototypeWhitelistArray[i] = reagent;
+                ++i;
+            }
+            temporarilyRemovedSolution = applicableTargetSolution.SplitSolutionWithout(applicableTargetSolution.Volume, reagentPrototypeWhitelistArray);
+        }
+
         // Get transfer amount. May be smaller than _transferAmount if not enough room, also make sure there's room in the injector
-        var realTransferAmount = FixedPoint2.Min(injector.Comp.TransferAmount, targetSolution.Comp.Solution.Volume,
+        var realTransferAmount = FixedPoint2.Min(injector.Comp.TransferAmount, applicableTargetSolution.Volume,
             solution.AvailableVolume);
 
         if (realTransferAmount <= 0)
@@ -350,6 +365,9 @@ public sealed class InjectorSystem : SharedInjectorSystem
         // Move units from attackSolution to targetSolution
         var removedSolution = SolutionContainers.Draw(target.Owner, targetSolution, realTransferAmount);
 
+        // Add back non-whitelisted reagents to the target solution
+        applicableTargetSolution.AddSolution(temporarilyRemovedSolution, null);
+
         if (!SolutionContainers.TryAddSolution(soln.Value, removedSolution))
         {
             return false;

diff --git a/Content.Shared/Chemistry/Components/InjectorComponent.cs b/Content.Shared/Chemistry/Components/InjectorComponent.cs
index 1f2716356c50fa91f81e9c6e6cf698565d049fea..17a65ef1c179cf93f129c6c898ffd872b9758721 100644 (file)
--- a/Content.Shared/Chemistry/Components/InjectorComponent.cs
+++ b/Content.Shared/Chemistry/Components/InjectorComponent.cs
@@ -1,7 +1,9 @@
 using Content.Shared.Chemistry.EntitySystems;
+using Content.Shared.Chemistry.Reagent;
 using Content.Shared.DoAfter;
 using Content.Shared.FixedPoint;
 using Robust.Shared.GameStates;
+using Robust.Shared.Prototypes;
 using Robust.Shared.Serialization;
 
 namespace Content.Shared.Chemistry.Components;
@@ -88,6 +90,14 @@ public sealed partial class InjectorComponent : Component
     [DataField]
     public InjectorToggleMode ToggleState = InjectorToggleMode.Draw;
 
+    /// <summary>
+    /// Reagents that are allowed to be within this injector.
+    /// If a solution has both allowed and non-allowed reagents, only allowed reagents will be drawn into this injector.
+    /// A null ReagentWhitelist indicates all reagents are allowed.
+    /// </summary>
+    [DataField]
+    public List<ProtoId<ReagentPrototype>>? ReagentWhitelist = null;
+
     #region Arguments for injection doafter
 
     /// <inheritdoc cref=DoAfterArgs.NeedHand>

diff --git a/Content.Shared/Chemistry/Components/Solution.cs b/Content.Shared/Chemistry/Components/Solution.cs
index fc25781005f8c6c5aa2921b72828e96c6ca82bbf..c65ba0e80ea3e8f3aaab1e19917f8ec55232f909 100644 (file)
--- a/Content.Shared/Chemistry/Components/Solution.cs
+++ b/Content.Shared/Chemistry/Components/Solution.cs
@@ -612,7 +612,7 @@ namespace Content.Shared.Chemistry.Components
         }
 
         /// <summary>
-        /// Splits a solution without the specified reagent prototypes.
+        /// Splits a solution with only the specified reagent prototypes.
         /// </summary>
         public Solution SplitSolutionWithOnly(FixedPoint2 toTake, params string[] includedPrototypes)
         {