Moved all affected stacks calculations to mechanics

lib/spells/BattleSpellMechanics.cpp

@@ -62,7 +62,7 @@ void AntimagicMechanics::applyBattle(BattleInfo * battle, const BattleSpellCast
 }
 
 ///ChainLightningMechanics
-std::vector<const CStack *> ChainLightningMechanics::getAffectedStacks(const CBattleInfoCallback * cb, SpellTargetingContext & ctx) const
+std::vector<const CStack *> ChainLightningMechanics::calculateAffectedStacks(const CBattleInfoCallback * cb, const SpellTargetingContext & ctx) const
 {
 	std::vector<const CStack *> res;
 

lib/spells/BattleSpellMechanics.h

@@ -44,7 +44,8 @@ class DLL_LINKAGE ChainLightningMechanics : public DefaultSpellMechanics
 {
 public:
 	ChainLightningMechanics(CSpell * s): DefaultSpellMechanics(s){};
-	std::vector<const CStack *> getAffectedStacks(const CBattleInfoCallback * cb, SpellTargetingContext & ctx) const override;
+protected:
+	std::vector<const CStack *> calculateAffectedStacks(const CBattleInfoCallback * cb, const SpellTargetingContext & ctx) const override;
 };
 
 class DLL_LINKAGE CloneMechanics : public DefaultSpellMechanics

lib/spells/CDefaultSpellMechanics.cpp

@@ -638,6 +638,13 @@ std::vector<BattleHex> DefaultSpellMechanics::rangeInHexes(BattleHex centralHex,
 }
 
 std::vector<const CStack *> DefaultSpellMechanics::getAffectedStacks(const CBattleInfoCallback * cb, SpellTargetingContext & ctx) const
+{
+	std::vector<const CStack *> attackedCres = calculateAffectedStacks(cb, ctx);
+	handleImmunities(cb, ctx, attackedCres);
+	return attackedCres;
+}
+
+std::vector<const CStack *> DefaultSpellMechanics::calculateAffectedStacks(const CBattleInfoCallback* cb, const SpellTargetingContext& ctx) const
 {
 	std::set<const CStack* > attackedCres;//std::set to exclude multiple occurrences of two hex creatures
 
@@ -751,6 +758,19 @@ void DefaultSpellMechanics::doDispell(BattleInfo * battle, const BattleSpellCast
 	}
 }
 
+void DefaultSpellMechanics::handleImmunities(const CBattleInfoCallback * cb, const SpellTargetingContext & ctx, std::vector<const CStack*> & stacks) const
+{
+	//now handle immunities
+	auto predicate = [&, this](const CStack * s)->bool
+	{
+		bool hitDirectly = ctx.ti.alwaysHitDirectly && s->coversPos(ctx.destination);
+		bool notImmune = (ESpellCastProblem::OK == owner->isImmuneByStack(ctx.caster, s));
+
+		return !(hitDirectly || notImmune);
+	};
+	vstd::erase_if(stacks, predicate);
+}
+
 void DefaultSpellMechanics::handleResistance(const SpellCastEnvironment * env, std::vector<const CStack* >& attackedCres, BattleSpellCast& sc) const
 {
 	//checking if creatures resist

lib/spells/CDefaultSpellMechanics.h

@@ -42,7 +42,7 @@ public:
 	DefaultSpellMechanics(CSpell * s): ISpellMechanics(s){};
 
 	std::vector<BattleHex> rangeInHexes(BattleHex centralHex, ui8 schoolLvl, ui8 side, bool * outDroppedHexes = nullptr) const override;
-	std::vector<const CStack *> getAffectedStacks(const CBattleInfoCallback * cb, SpellTargetingContext & ctx) const override;
+	std::vector<const CStack *> getAffectedStacks(const CBattleInfoCallback * cb, SpellTargetingContext & ctx) const override final;
 
 	ESpellCastProblem::ESpellCastProblem canBeCast(const CBattleInfoCallback * cb, const ISpellCaster * caster) const override;
 	ESpellCastProblem::ESpellCastProblem canBeCast(const CBattleInfoCallback * cb, const SpellTargetingContext & ctx) const override;
@@ -60,6 +60,10 @@ public:
 protected:
 	virtual void applyBattleEffects(const SpellCastEnvironment * env, const BattleSpellCastParameters & parameters, SpellCastContext & ctx) const;
 
+	virtual std::vector<const CStack *> calculateAffectedStacks(const CBattleInfoCallback * cb, const SpellTargetingContext & ctx) const;
+	virtual void handleImmunities(const CBattleInfoCallback * cb, const SpellTargetingContext & ctx, std::vector<const CStack *> & stacks) const;
+
+protected:
 	void doDispell(BattleInfo * battle, const BattleSpellCast * packet, const CSelector & selector) const;
 private:
 	void castNormal(const SpellCastEnvironment * env, const BattleSpellCastParameters & parameters, std::vector <const CStack*> & reflected) const;

lib/spells/CSpellHandler.cpp

@@ -214,20 +214,7 @@ std::vector<BattleHex> CSpell::rangeInHexes(BattleHex centralHex, ui8 schoolLvl,
 std::vector<const CStack *> CSpell::getAffectedStacks(const CBattleInfoCallback * cb, ECastingMode::ECastingMode mode, const ISpellCaster * caster, int spellLvl, BattleHex destination) const
 {
 	SpellTargetingContext ctx(this, mode, caster, spellLvl, destination);
-
-	std::vector<const CStack *> attackedCres = mechanics->getAffectedStacks(cb, ctx);
-
-	//now handle immunities
-	auto predicate = [&, this](const CStack * s)->bool
-	{
-		bool hitDirectly = ctx.ti.alwaysHitDirectly && s->coversPos(destination);
-		bool notImmune = (ESpellCastProblem::OK == isImmuneByStack(caster, s));
-
-		return !(hitDirectly || notImmune);
-	};
-	vstd::erase_if(attackedCres, predicate);
-
-	return attackedCres;
+	return mechanics->getAffectedStacks(cb, ctx);;
 }
 
 CSpell::ETargetType CSpell::getTargetType() const