Core/Units: implement PowerTypeFlags::RegenAffectedByHaste

6 files changed, 68 insertions, 21 deletions

diff --git a/src/server/game/DataStores/DBCEnums.h b/src/server/game/DataStores/DBCEnums.h
index a5d88d81f50..0f5e9e3d599 100644
--- a/src/server/game/DataStores/DBCEnums.h
+++ b/src/server/game/DataStores/DBCEnums.h
@@ -1856,7 +1856,7 @@ enum class PowerTypeFlags : int16
     NotSetToDefaultOnResurrect    = 0x0040, // NYI
     IsUsedByNPCs                  = 0x0080,
     ContinueRegenWhileFatigued    = 0x0200, // NYI
-    RegenAffectedByHaste          = 0x0400, // NYI
+    RegenAffectedByHaste          = 0x0400,
     SetToMaxOnLevelUp             = 0x1000,
     SetToMaxOnInitialLogIn        = 0x2000, // NYI
     AllowCostModsForPlayers       = 0x4000  // NYI
diff --git a/src/server/game/Entities/Player/Player.cpp b/src/server/game/Entities/Player/Player.cpp
index 38e2ef148fb..13d91494f1f 100644
--- a/src/server/game/Entities/Player/Player.cpp
+++ b/src/server/game/Entities/Player/Player.cpp
@@ -5335,6 +5335,10 @@ void Player::UpdateRating(CombatRating cr)
                     ApplyAttackTimePercentMod(OFF_ATTACK, oldVal, false);
                     ApplyAttackTimePercentMod(BASE_ATTACK, newVal, true);
                     ApplyAttackTimePercentMod(OFF_ATTACK, newVal, true);
+
+                    // Item and aura mods increase the haste of all three combat ratings so we only have to pick one
+                    ApplyHasteRegenMod(oldVal, false);
+                    ApplyHasteRegenMod(newVal, true);
                     break;
                 case CR_HASTE_RANGED:
                     ApplyAttackTimePercentMod(RANGED_ATTACK, oldVal, false);
diff --git a/src/server/game/Entities/Unit/StatSystem.cpp b/src/server/game/Entities/Unit/StatSystem.cpp
index aab089781e6..ed331937999 100644
--- a/src/server/game/Entities/Unit/StatSystem.cpp
+++ b/src/server/game/Entities/Unit/StatSystem.cpp
@@ -839,30 +839,38 @@ void Player::UpdatePowerRegen(Powers powerType)
             SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenInterruptedFlatModifier, powerIndex), baseRegen + CalculatePct(spiritRegen, modManaRegenInterrupt));
             break;
         }
-        case POWER_FOCUS:
-        case POWER_ENERGY:
-        {
-            PowerTypeEntry const* powerTypeEntry = sDB2Manager.GetPowerTypeEntry(powerType);
-            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenFlatModifier, powerIndex), powerTypeEntry->RegenPeace * powerRegenModPct - powerTypeEntry->RegenPeace + powerRegenMod);
-            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenInterruptedFlatModifier, powerIndex), powerTypeEntry->RegenCombat * powerRegenModPct - powerTypeEntry->RegenCombat + powerRegenMod);
-            break;
-        }
-        case POWER_RUNIC_POWER:
-        case POWER_RAGE:
-            // Butchery and Anger Management
-            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenInterruptedFlatModifier, powerIndex), powerRegenMod);
-            break;
-        case POWER_RUNE_BLOOD:
-        case POWER_RUNE_FROST:
-        case POWER_RUNE_UNHOLY:
+        default:
         {
+            // Classic Only - Death Knight Runes use the flags of the POWER_RUNES
+            if (powerType == POWER_RUNE_BLOOD || powerType == POWER_RUNE_FROST || powerType == POWER_RUNE_UNHOLY)
+                powerType = POWER_RUNES;
+
             PowerTypeEntry const* powerTypeEntry = sDB2Manager.GetPowerTypeEntry(powerType);
-            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenFlatModifier, powerIndex), float(1 * IN_MILLISECONDS) / float(RUNE_BASE_COOLDOWN) - powerTypeEntry->RegenPeace);
-            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenInterruptedFlatModifier, powerIndex), float(1 * IN_MILLISECONDS) / float(RUNE_BASE_COOLDOWN) - powerTypeEntry->RegenCombat);
+            // Base Regen
+            float peaceRegen = powerTypeEntry->RegenPeace;
+            float combatRegen = powerTypeEntry->RegenCombat;
+
+            // Haste Regen
+            if (powerTypeEntry->GetFlags().HasFlag(PowerTypeFlags::RegenAffectedByHaste) && G3D::fuzzyNe(m_unitData->ModHaste, 0.0f))
+            {
+                peaceRegen /= m_unitData->ModHaste;
+                combatRegen /= m_unitData->ModHaste;
+            }
+
+            peaceRegen *= powerRegenModPct;
+            combatRegen *= powerRegenModPct;
+
+            // Subtract the base value to get the proper offset
+            peaceRegen -= powerTypeEntry->RegenPeace;
+            combatRegen -= powerTypeEntry->RegenCombat;
+
+            peaceRegen += powerRegenMod;
+            combatRegen += powerRegenMod;
+
+            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenFlatModifier, powerIndex), peaceRegen);
+            SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::PowerRegenInterruptedFlatModifier, powerIndex), combatRegen);
             break;
         }
-        default:
-            break;
     }
 }
 
diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index 4746b6e845f..bf3b0cb4778 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -10502,6 +10502,38 @@ void Unit::ApplyCastTimePercentMod(float val, bool apply)
     }
 }
 
+void Unit::ApplyHasteRegenMod(float val, bool apply)
+{
+    if (val > 0.f)
+        ApplyPercentModUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::ModHasteRegen), val, !apply);
+    else
+        ApplyPercentModUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::ModHasteRegen), -val, apply);
+
+    if (IsPlayer())
+    {
+        for (uint8 powerType = POWER_MANA; powerType != MAX_POWERS; ++powerType)
+        {
+            uint32 powerIndex = GetPowerIndex(static_cast<Powers>(powerType));
+            if (powerIndex == MAX_POWERS)
+                continue;
+
+            PowerTypeEntry const* powerTypeEntry = sDB2Manager.GetPowerTypeEntry(static_cast<Powers>(powerType));
+            if (!powerTypeEntry)
+                continue;
+
+            bool regenAffectedByHaste = powerTypeEntry->GetFlags().HasFlag(PowerTypeFlags::RegenAffectedByHaste);
+
+            // Classic Only - Death Knight Runes use the flags of the POWER_RUNES
+            if (powerType == POWER_RUNE_BLOOD || powerType == POWER_RUNE_FROST || powerType == POWER_RUNE_UNHOLY)
+                if (PowerTypeEntry const* powerTypeEntry = sDB2Manager.GetPowerTypeEntry(POWER_RUNES))
+                    regenAffectedByHaste = powerTypeEntry->GetFlags().HasFlag(PowerTypeFlags::RegenAffectedByHaste);
+
+            if (regenAffectedByHaste)
+                ToPlayer()->UpdatePowerRegen(static_cast<Powers>(powerType));
+        }
+    }
+}
+
 void Unit::UpdateAuraForGroup()
 {
     if (Player* player = ToPlayer())
diff --git a/src/server/game/Entities/Unit/Unit.h b/src/server/game/Entities/Unit/Unit.h
index 257fca6ed24..5bdc465b23d 100644
--- a/src/server/game/Entities/Unit/Unit.h
+++ b/src/server/game/Entities/Unit/Unit.h
@@ -819,6 +819,7 @@ class TC_GAME_API Unit : public WorldObject
         void UpdateAttackTimeField(WeaponAttackType att);
         void ApplyAttackTimePercentMod(WeaponAttackType att, float val, bool apply);
         void ApplyCastTimePercentMod(float val, bool apply);
+        void ApplyHasteRegenMod(float val, bool apply);
 
         void SetModCastingSpeed(float castingSpeed) { SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::ModCastingSpeed), castingSpeed); }
         void SetModSpellHaste(float spellHaste) { SetUpdateFieldValue(m_values.ModifyValue(&Unit::m_unitData).ModifyValue(&UF::UnitData::ModSpellHaste), spellHaste); }
diff --git a/src/server/game/Spells/Auras/SpellAuraEffects.cpp b/src/server/game/Spells/Auras/SpellAuraEffects.cpp
index 26248007f18..805c5c9ffa3 100644
--- a/src/server/game/Spells/Auras/SpellAuraEffects.cpp
+++ b/src/server/game/Spells/Auras/SpellAuraEffects.cpp
@@ -4436,9 +4436,11 @@ void AuraEffect::HandleModMeleeSpeedPct(AuraApplication const* aurApp, uint8 mod
     {
         target->ApplyAttackTimePercentMod(BASE_ATTACK, float(spellGroupVal), !apply);
         target->ApplyAttackTimePercentMod(OFF_ATTACK, float(spellGroupVal), !apply);
+        target->ApplyHasteRegenMod(float(spellGroupVal), !apply);
     }
     target->ApplyAttackTimePercentMod(BASE_ATTACK, float(GetAmount()), apply);
     target->ApplyAttackTimePercentMod(OFF_ATTACK,  float(GetAmount()), apply);
+    target->ApplyHasteRegenMod(float(GetAmount()), apply);
 }
 
 void AuraEffect::HandleAuraModRangedHaste(AuraApplication const* aurApp, uint8 mode, bool apply) const