Core/Spells: Expose number of targets selected for each spell effect to scripts to allow implementing spells that do something depending on number of targets hit

5 files changed, 79 insertions, 11 deletions

diff --git a/src/server/game/Spells/Spell.cpp b/src/server/game/Spells/Spell.cpp
index 97cd6f64d20..f4e058fad08 100644
--- a/src/server/game/Spells/Spell.cpp
+++ b/src/server/game/Spells/Spell.cpp
@@ -2276,6 +2276,30 @@ void Spell::AddDestTarget(SpellDestination const& dest, uint32 effIndex)
     m_destTargets[effIndex] = dest;
 }
 
+int64 Spell::GetUnitTargetCountForEffect(SpellEffIndex effect) const
+{
+    return std::count_if(m_UniqueTargetInfo.begin(), m_UniqueTargetInfo.end(), [effect](TargetInfo const& targetInfo)
+    {
+        return targetInfo.effectMask & (1 << effect);
+    });
+}
+
+int64 Spell::GetGameObjectTargetCountForEffect(SpellEffIndex effect) const
+{
+    return std::count_if(m_UniqueGOTargetInfo.begin(), m_UniqueGOTargetInfo.end(), [effect](GOTargetInfo const& targetInfo)
+    {
+        return targetInfo.effectMask & (1 << effect);
+    });
+}
+
+int64 Spell::GetItemTargetCountForEffect(SpellEffIndex effect) const
+{
+    return std::count_if(m_UniqueItemInfo.begin(), m_UniqueItemInfo.end(), [effect](ItemTargetInfo const& targetInfo)
+    {
+        return targetInfo.effectMask & (1 << effect);
+    });
+}
+
 void Spell::DoAllEffectOnTarget(TargetInfo* target)
 {
     if (!target || target->processed)
@@ -7454,9 +7478,9 @@ void Spell::DoAllEffectOnLaunchTarget(TargetInfo& targetInfo, float* multiplier)
 
                     if (m_caster->GetTypeId() == TYPEID_PLAYER)
                     {
-                        uint32 targetAmount = m_UniqueTargetInfo.size();
+                        int64 targetAmount = GetUnitTargetCountForEffect(SpellEffIndex(effect->EffectIndex));
                         if (targetAmount > 20)
-                            m_damage = m_damage * 20/targetAmount;
+                            m_damage = m_damage * 20 / targetAmount;
                     }
                 }
             }
diff --git a/src/server/game/Spells/Spell.h b/src/server/game/Spells/Spell.h
index d625bd95f40..a817e1dfc17 100644
--- a/src/server/game/Spells/Spell.h
+++ b/src/server/game/Spells/Spell.h
@@ -699,6 +699,10 @@ class TC_GAME_API Spell
 
         int32 GetTimer() const { return m_timer; }
 
+        int64 GetUnitTargetCountForEffect(SpellEffIndex effect) const;
+        int64 GetGameObjectTargetCountForEffect(SpellEffIndex effect) const;
+        int64 GetItemTargetCountForEffect(SpellEffIndex effect) const;
+
     protected:
         bool HasGlobalCooldown() const;
         void TriggerGlobalCooldown();
diff --git a/src/server/game/Spells/SpellEffects.cpp b/src/server/game/Spells/SpellEffects.cpp
index 83c6b37b925..d44e84529c4 100644
--- a/src/server/game/Spells/SpellEffects.cpp
+++ b/src/server/game/Spells/SpellEffects.cpp
@@ -464,13 +464,8 @@ void Spell::EffectSchoolDMG(SpellEffIndex effIndex)
         // Meteor like spells (divided damage to targets)
         if (m_spellInfo->HasAttribute(SPELL_ATTR0_CU_SHARE_DAMAGE))
         {
-            uint32 count = std::count_if(m_UniqueTargetInfo.begin(), m_UniqueTargetInfo.end(), [effIndex](TargetInfo const& targetInfo)
-            {
-                return targetInfo.effectMask & (1 << effIndex);
-            });
-
             // divide to all targets
-            if (count)
+            if (int64 count = GetUnitTargetCountForEffect(effIndex))
                 damage /= count;
         }
 
diff --git a/src/server/game/Spells/SpellScript.cpp b/src/server/game/Spells/SpellScript.cpp
index 12337bd0c67..294028d8882 100644
--- a/src/server/game/Spells/SpellScript.cpp
+++ b/src/server/game/Spells/SpellScript.cpp
@@ -413,6 +413,15 @@ bool SpellScript::IsInCheckCastHook() const
 {
     return m_currentScriptState == SPELL_SCRIPT_HOOK_CHECK_CAST;
 }
+
+bool SpellScript::IsAfterTargetSelectionPhase() const
+{
+    return IsInHitPhase()
+        || m_currentScriptState == SPELL_SCRIPT_HOOK_ON_CAST
+        || m_currentScriptState == SPELL_SCRIPT_HOOK_AFTER_CAST
+        || m_currentScriptState == SPELL_SCRIPT_HOOK_CALC_CRIT_CHANCE;
+}
+
 bool SpellScript::IsInTargetHook() const
 {
     switch (m_currentScriptState)
@@ -485,6 +494,39 @@ Item* SpellScript::GetExplTargetItem() const
     return m_spell->m_targets.GetItemTarget();
 }
 
+int64 SpellScript::GetUnitTargetCountForEffect(SpellEffIndex effect) const
+{
+    if (!IsAfterTargetSelectionPhase())
+    {
+        TC_LOG_ERROR("scripts", "Script: `%s` Spell: `%u`: function SpellScript::GetUnitTargetCountForEffect was called, but function has no effect in current hook! (spell has not selected targets yet)",
+            m_scriptName->c_str(), m_scriptSpellId);
+        return 0;
+    }
+    return m_spell->GetUnitTargetCountForEffect(effect);
+}
+
+int64 SpellScript::GetGameObjectTargetCountForEffect(SpellEffIndex effect) const
+{
+    if (!IsAfterTargetSelectionPhase())
+    {
+        TC_LOG_ERROR("scripts", "Script: `%s` Spell: `%u`: function SpellScript::GetGameObjectTargetCountForEffect was called, but function has no effect in current hook! (spell has not selected targets yet)",
+            m_scriptName->c_str(), m_scriptSpellId);
+        return 0;
+    }
+    return m_spell->GetGameObjectTargetCountForEffect(effect);
+}
+
+int64 SpellScript::GetItemTargetCountForEffect(SpellEffIndex effect) const
+{
+    if (!IsAfterTargetSelectionPhase())
+    {
+        TC_LOG_ERROR("scripts", "Script: `%s` Spell: `%u`: function SpellScript::GetItemTargetCountForEffect was called, but function has no effect in current hook! (spell has not selected targets yet)",
+            m_scriptName->c_str(), m_scriptSpellId);
+        return 0;
+    }
+    return m_spell->GetItemTargetCountForEffect(effect);
+}
+
 Unit* SpellScript::GetHitUnit() const
 {
     if (!IsInTargetHook())
diff --git a/src/server/game/Spells/SpellScript.h b/src/server/game/Spells/SpellScript.h
index 7325f75b3f6..fc665f0947a 100644
--- a/src/server/game/Spells/SpellScript.h
+++ b/src/server/game/Spells/SpellScript.h
@@ -181,9 +181,6 @@ enum SpellScriptHookType
 
 #define HOOK_SPELL_HIT_START SPELL_SCRIPT_HOOK_EFFECT_HIT
 #define HOOK_SPELL_HIT_END SPELL_SCRIPT_HOOK_AFTER_HIT + 1
-#define HOOK_SPELL_START SPELL_SCRIPT_HOOK_EFFECT
-#define HOOK_SPELL_END SPELL_SCRIPT_HOOK_CHECK_CAST + 1
-#define HOOK_SPELL_COUNT HOOK_SPELL_END - HOOK_SPELL_START
 
 class TC_GAME_API SpellScript : public _SpellScript
 {
@@ -320,6 +317,7 @@ class TC_GAME_API SpellScript : public _SpellScript
         void _PrepareScriptCall(SpellScriptHookType hookType);
         void _FinishScriptCall();
         bool IsInCheckCastHook() const;
+        bool IsAfterTargetSelectionPhase() const;
         bool IsInTargetHook() const;
         bool IsInHitPhase() const;
         bool IsInEffectHook() const;
@@ -444,6 +442,11 @@ class TC_GAME_API SpellScript : public _SpellScript
         // returns: Item which was selected as an explicit spell target or NULL if there's no target
         Item* GetExplTargetItem() const;
 
+        // methods usable only after spell targets have been fully selected
+        int64 GetUnitTargetCountForEffect(SpellEffIndex effect) const;
+        int64 GetGameObjectTargetCountForEffect(SpellEffIndex effect) const;
+        int64 GetItemTargetCountForEffect(SpellEffIndex effect) const;
+
         // methods useable only during spell hit on target, or during spell launch on target:
         // returns: target of current effect if it was Unit otherwise NULL
         Unit* GetHitUnit() const;