Move logic from EvenHealthChangeEntityEffectSystem to the damage system API (#41684)

* add two methods

* move stuff to damage system api

* use TryIndex

* simplify

* minor fix

* add helper functions

* fix

* remove random new line

* simplify

* remove unnecessary lines

* rename to GetDamage

* Got it working

* make more clear

* why backwards

* value should be the amount to heal

* fix

* fix all dumb fixedpoint2 edge cases I hope

* One more thing

* fix

* make it more simple

* ops it was backwards

* valueHeal can't be more than remaining

* add all keys beforehand and no need to check and add them inside the loop

* break for loop in case remaining is zero

* comment was wrong

* optimized, works

* remove random spaces

---------

Co-authored-by: Princess Cheeseballs <66055347+Pronana@users.noreply.github.com>

Content.Server/Body/Systems/RespiratorSystem.cs

@@ -17,8 +17,8 @@ using Content.Shared.Database;
 using Content.Shared.EntityConditions;
 using Content.Shared.EntityConditions.Conditions.Body;
 using Content.Shared.EntityEffects;
-using Content.Shared.EntityEffects.Effects;
 using Content.Shared.EntityEffects.Effects.Body;
+using Content.Shared.EntityEffects.Effects.Damage;
 using Content.Shared.Mobs.Systems;
 using JetBrains.Annotations;
 using Robust.Shared.Prototypes;

Content.Shared/Damage/Systems/DamageableSystem.API.cs

@@ -1,3 +1,5 @@
+using System.Linq;
+using System.Net.Sockets;
 using Content.Shared.Damage.Components;
 using Content.Shared.Damage.Prototypes;
 using Content.Shared.FixedPoint;
@@ -184,6 +186,179 @@ public sealed partial class DamageableSystem
         return damageDone;
     }
 
+    /// <summary>
+    /// Will reduce the damage on the entity exactly by <see cref="amount"/> as close as equally distributed among all damage types the entity has.
+    /// If one of the damage types of the entity is too low. it will heal that completly and distribute the excess healing among the other damage types.
+    /// If the <see cref="amount"/> is larger than the total damage of the entity then it just clears all damage.
+    /// </summary>
+    /// <param name="ent">entity to be healed</param>
+    /// <param name="amount">how much to heal. value has to be negative to heal</param>
+    /// <param name="group">from which group to heal. if null, heal from all groups</param>
+    /// <param name="origin">who did the healing</param>
+    public DamageSpecifier HealEvenly(
+        Entity<DamageableComponent?> ent,
+        FixedPoint2 amount,
+        ProtoId<DamageGroupPrototype>? group = null,
+        EntityUid? origin = null)
+    {
+        var damageChange = new DamageSpecifier();
+
+        if (!_damageableQuery.Resolve(ent, ref ent.Comp, false) || amount >= 0)
+            return damageChange;
+
+        // Get our total damage, or heal if we're below a certain amount.
+        if (!TryGetDamageGreaterThan((ent, ent.Comp), -amount, out var damage, group))
+            return ChangeDamage(ent, -damage, true, false, origin);
+
+        // make sure damageChange has the same damage types as damage
+        damageChange.DamageDict.EnsureCapacity(damage.DamageDict.Count);
+        foreach (var type in damage.DamageDict.Keys)
+        {
+            damageChange.DamageDict.Add(type, FixedPoint2.Zero);
+        }
+
+        var remaining = -amount;
+        var keys = damage.DamageDict.Keys.ToList();
+
+        while (remaining > 0)
+        {
+            var count = keys.Count;
+            // We do this to ensure that we always round up when dividing to avoid excess loops.
+            // We already have logic to prevent healing more than we have.
+            var maxHeal = count == 1 ? remaining : (remaining + FixedPoint2.Epsilon * (count - 1)) / count;
+
+            // Iterate backwards since we're removing items.
+            for (var i = count - 1; i >= 0; i--)
+            {
+                var type = keys[i];
+                // This is the amount we're trying to heal, capped by maxHeal
+                var heal = damage.DamageDict[type] + damageChange.DamageDict[type];
+
+                // Don't go above max, if we don't go above max
+                if (heal > maxHeal)
+                    heal = maxHeal;
+                // If we're not above max, we will heal it fully and don't need to enumerate anymore!
+                else
+                    keys.RemoveAt(i);
+
+                if (heal >= remaining)
+                {
+                    // Don't remove more than we can remove. Prevents us from healing more than we'd expect...
+                    damageChange.DamageDict[type] -= remaining;
+                    remaining = FixedPoint2.Zero;
+                    break;
+                }
+
+                remaining -= heal;
+                damageChange.DamageDict[type] -= heal;
+            }
+        }
+
+        return ChangeDamage(ent, damageChange, true, false, origin);
+    }
+
+    /// <summary>
+    /// Will reduce the damage on the entity exactly by <see cref="amount"/> distributed by weight among all damage types the entity has.
+    /// (the weight is how much damage of the type there is)
+    /// If the <see cref="amount"/> is larger than the total damage of the entity then it just clears all damage.
+    /// </summary>
+    /// <param name="ent">entity to be healed</param>
+    /// <param name="amount">how much to heal. value has to be negative to heal</param>
+    /// <param name="group">from which group to heal. if null, heal from all groups</param>
+    /// <param name="origin">who did the healing</param>
+    public DamageSpecifier HealDistributed(
+        Entity<DamageableComponent?> ent,
+        FixedPoint2 amount,
+        ProtoId<DamageGroupPrototype>? group = null,
+        EntityUid? origin = null)
+    {
+        var damageChange = new DamageSpecifier();
+
+        if (!_damageableQuery.Resolve(ent, ref ent.Comp, false) || amount >= 0)
+            return damageChange;
+
+        // Get our total damage, or heal if we're below a certain amount.
+        if (!TryGetDamageGreaterThan((ent, ent.Comp), -amount, out var damage, group))
+            return ChangeDamage(ent, -damage, true, false, origin);
+
+        // make sure damageChange has the same damage types as damageEntity
+        damageChange.DamageDict.EnsureCapacity(damage.DamageDict.Count);
+        var total = damage.GetTotal();
+
+        // heal weighted by the damage of that type
+        foreach (var (type, value) in damage.DamageDict)
+        {
+            damageChange.DamageDict.Add(type, value / total * amount);
+        }
+
+        return ChangeDamage(ent, damageChange, true, false, origin);
+    }
+
+    /// <summary>
+    /// Tries to get damage from an entity with an optional group specifier.
+    /// </summary>
+    /// <param name="ent">Entity we're checking the damage on</param>
+    /// <param name="amount">Amount we want the damage to be greater than ideally</param>
+    /// <param name="damage">Damage specifier we're returning with</param>
+    /// <param name="group">An optional group, note that if it fails to index it will just use all damage.</param>
+    /// <returns>True if the total damage is greater than the specified amount</returns>
+    public bool TryGetDamageGreaterThan(Entity<DamageableComponent> ent,
+        FixedPoint2 amount,
+        out DamageSpecifier damage,
+        ProtoId<DamageGroupPrototype>? group = null)
+    {
+        // get the damage should be healed (either all or only from one group)
+        damage = group == null ? GetDamage(ent) : GetDamage(ent, group.Value);
+
+        // If trying to heal more than the total damage of damageEntity just heal everything
+        return damage.GetTotal() > amount;
+    }
+
+    /// <summary>
+    /// Returns a <see cref="DamageSpecifier"/> with all positive damage of the entity from the group specified
+    /// </summary>
+    /// <param name="ent">entity with damage</param>
+    /// <param name="group">group of damage to get values from</param>
+    /// <returns></returns>
+    public DamageSpecifier GetDamage(Entity<DamageableComponent> ent, ProtoId<DamageGroupPrototype> group)
+    {
+        // No damage if no group exists...
+        if (!_prototypeManager.Resolve(group, out var groupProto))
+            return new DamageSpecifier();
+
+        var damage = new DamageSpecifier();
+        damage.DamageDict.EnsureCapacity(groupProto.DamageTypes.Count);
+
+        foreach (var damageId in groupProto.DamageTypes)
+        {
+            if (!ent.Comp.Damage.DamageDict.TryGetValue(damageId, out var value))
+                continue;
+            if (value > FixedPoint2.Zero)
+                damage.DamageDict.Add(damageId, value);
+        }
+
+        return damage;
+    }
+
+    /// <summary>
+    /// Returns a <see cref="DamageSpecifier"/> with all positive damage of the entity
+    /// </summary>
+    /// <param name="ent">entity with damage</param>
+    /// <returns></returns>
+    public DamageSpecifier GetDamage(Entity<DamageableComponent> ent)
+    {
+        var damage = new DamageSpecifier();
+        damage.DamageDict.EnsureCapacity(ent.Comp.Damage.DamageDict.Count);
+
+        foreach (var (damageId, value) in ent.Comp.Damage.DamageDict)
+        {
+            if (value > FixedPoint2.Zero)
+                damage.DamageDict.Add(damageId, value);
+        }
+
+        return damage;
+    }
+
     /// <summary>
     /// Applies the two universal "All" modifiers, if set.
     /// Individual damage source modifiers are set in their respective code.

Content.Shared/EntityEffects/Effects/Damage/DistributedHealthChangeEntityEffectSystem.cs

@@ -0,0 +1,91 @@
+using Content.Shared.Damage.Components;
+using Content.Shared.Damage.Prototypes;
+using Content.Shared.Damage.Systems;
+using Content.Shared.FixedPoint;
+using Content.Shared.Localizations;
+using Robust.Shared.Prototypes;
+
+namespace Content.Shared.EntityEffects.Effects.Damage;
+
+/// <summary>
+/// Heal the damage types in a damage group by up to a specified total on this entity.
+/// The amount healed per type is weighted by the amount of damage for that type scaling linearly.
+/// Total adjustment is modified by scale.
+/// </summary>
+/// <inheritdoc cref="EntityEffectSystem{T,TEffect}"/>
+public sealed partial class DistributedHealthChangeEntityEffectSystem : EntityEffectSystem<DamageableComponent, DistributedHealthChange>
+{
+    [Dependency] private readonly DamageableSystem _damageable = default!;
+
+    protected override void Effect(Entity<DamageableComponent> entity, ref EntityEffectEvent<DistributedHealthChange> args)
+    {
+        foreach (var (group, amount) in args.Effect.Damage)
+        {
+            _damageable.HealDistributed(entity.AsNullable(), amount * args.Scale, group);
+        }
+    }
+}
+
+/// <inheritdoc cref="EntityEffect"/>
+public sealed partial class DistributedHealthChange : EntityEffectBase<DistributedHealthChange>
+{
+    /// <summary>
+    /// Damage to heal, collected into entire damage groups.
+    /// </summary>
+    [DataField(required: true)]
+    public Dictionary<ProtoId<DamageGroupPrototype>, FixedPoint2> Damage = new();
+
+    /// <summary>
+    /// Should this effect ignore damage modifiers?
+    /// </summary>
+    [DataField]
+    public bool IgnoreResistances = true;
+
+    public override string EntityEffectGuidebookText(IPrototypeManager prototype, IEntitySystemManager entSys)
+    {
+        var damages = new List<string>();
+        var heals = false;
+        var deals = false;
+
+        var damagableSystem = entSys.GetEntitySystem<DamageableSystem>();
+        var universalReagentDamageModifier = damagableSystem.UniversalReagentDamageModifier;
+        var universalReagentHealModifier = damagableSystem.UniversalReagentHealModifier;
+
+        foreach (var (group, amount) in Damage)
+        {
+            var groupProto = prototype.Index(group);
+
+            var sign = FixedPoint2.Sign(amount);
+            float mod;
+
+            switch (sign)
+            {
+                case < 0:
+                    heals = true;
+                    mod = universalReagentHealModifier;
+                    break;
+                case > 0:
+                    deals = true;
+                    mod = universalReagentDamageModifier;
+                    break;
+                default:
+                    continue; // Don't need to show damage types of 0...
+            }
+
+            damages.Add(
+                Loc.GetString("health-change-display",
+                    ("kind", groupProto.LocalizedName),
+                    ("amount", MathF.Abs(amount.Float() * mod)),
+                    ("deltasign", sign)
+                ));
+        }
+
+        // We use health change since in practice it's not even and distributed is a mouthful.
+        // Also because healing groups not using even or distributed healing should be kill.
+        var healsordeals = heals ? deals ? "both" : "heals" : deals ? "deals" : "none";
+        return Loc.GetString("entity-effect-guidebook-health-change",
+            ("chance", Probability),
+            ("changes", ContentLocalizationManager.FormatList(damages)),
+            ("healsordeals", healsordeals));
+    }
+}

Content.Shared/EntityEffects/Effects/Damage/EvenHealthChangeEntityEffectSystem.cs

@@ -1,54 +1,27 @@
-using Content.Shared.Damage;
-using Content.Shared.Damage.Components;
+using Content.Shared.Damage.Components;
 using Content.Shared.Damage.Prototypes;
 using Content.Shared.Damage.Systems;
 using Content.Shared.FixedPoint;
 using Content.Shared.Localizations;
 using Robust.Shared.Prototypes;
-using Robust.Shared.Utility;
 
-namespace Content.Shared.EntityEffects.Effects;
+namespace Content.Shared.EntityEffects.Effects.Damage;
 
 /// <summary>
-/// Evenly adjust the damage types in a damage group by up to a specified total on this entity.
+/// Evenly heal the damage types in a damage group by up to a specified total on this entity.
 /// Total adjustment is modified by scale.
 /// </summary>
 /// <inheritdoc cref="EntityEffectSystem{T,TEffect}"/>
 public sealed partial class EvenHealthChangeEntityEffectSystem : EntityEffectSystem<DamageableComponent, EvenHealthChange>
 {
     [Dependency] private readonly DamageableSystem _damageable = default!;
-    [Dependency] private readonly IPrototypeManager _proto = default!;
 
     protected override void Effect(Entity<DamageableComponent> entity, ref EntityEffectEvent<EvenHealthChange> args)
     {
-        var damageSpec = new DamageSpecifier();
-
         foreach (var (group, amount) in args.Effect.Damage)
         {
-            var groupProto = _proto.Index(group);
-            var groupDamage = new Dictionary<string, FixedPoint2>();
-            foreach (var damageId in groupProto.DamageTypes)
-            {
-                var damageAmount = entity.Comp.Damage.DamageDict.GetValueOrDefault(damageId);
-                if (damageAmount != FixedPoint2.Zero)
-                    groupDamage.Add(damageId, damageAmount);
-            }
-
-            var sum = groupDamage.Values.Sum();
-            foreach (var (damageId, damageAmount) in groupDamage)
-            {
-                var existing = damageSpec.DamageDict.GetOrNew(damageId);
-                damageSpec.DamageDict[damageId] = existing + damageAmount / sum * amount;
-            }
+            _damageable.HealEvenly(entity.AsNullable(), amount * args.Scale, group);
         }
-
-        damageSpec *= args.Scale;
-
-        _damageable.TryChangeDamage(
-            entity.AsNullable(),
-            damageSpec,
-            args.Effect.IgnoreResistances,
-            interruptsDoAfters: false);
     }
 }
 

Content.Shared/EntityEffects/Effects/Damage/HealthChangeEntityEffectSystem.cs

@@ -1,11 +1,12 @@
 using Content.Shared.Damage;
 using Content.Shared.Damage.Components;
 using Content.Shared.Damage.Prototypes;
+using Content.Shared.Damage.Systems;
 using Content.Shared.FixedPoint;
 using Content.Shared.Localizations;
 using Robust.Shared.Prototypes;
 
-namespace Content.Shared.EntityEffects.Effects;
+namespace Content.Shared.EntityEffects.Effects.Damage;
 
 /// <summary>
 /// Adjust the damages on this entity by specified amounts.
@@ -14,7 +15,7 @@ namespace Content.Shared.EntityEffects.Effects;
 /// <inheritdoc cref="EntityEffectSystem{T,TEffect}"/>
 public sealed partial class HealthChangeEntityEffectSystem : EntityEffectSystem<DamageableComponent, HealthChange>
 {
-    [Dependency] private readonly Damage.Systems.DamageableSystem _damageable = default!;
+    [Dependency] private readonly DamageableSystem _damageable = default!;
 
     protected override void Effect(Entity<DamageableComponent> entity, ref EntityEffectEvent<HealthChange> args)
     {
@@ -50,10 +51,10 @@ public sealed partial class HealthChange : EntityEffectBase<HealthChange>
 
             var damageSpec = new DamageSpecifier(Damage);
 
-            var universalReagentDamageModifier = entSys.GetEntitySystem<Damage.Systems.DamageableSystem>().UniversalReagentDamageModifier;
-            var universalReagentHealModifier = entSys.GetEntitySystem<Damage.Systems.DamageableSystem>().UniversalReagentHealModifier;
+            var universalReagentDamageModifier = entSys.GetEntitySystem<DamageableSystem>().UniversalReagentDamageModifier;
+            var universalReagentHealModifier = entSys.GetEntitySystem<DamageableSystem>().UniversalReagentHealModifier;
 
-            damageSpec = entSys.GetEntitySystem<Damage.Systems.DamageableSystem>().ApplyUniversalAllModifiers(damageSpec);
+            damageSpec = entSys.GetEntitySystem<DamageableSystem>().ApplyUniversalAllModifiers(damageSpec);
 
             foreach (var (kind, amount) in damageSpec.DamageDict)
             {