Move getAffectedStacks to mechanics classes

lib/CSpellHandler.cpp

@@ -31,7 +31,6 @@
 namespace SpellConfig
 {
 	static const std::string LEVEL_NAMES[] = {"none", "basic", "advanced", "expert"};
-
 }
 
 ///CSpell::LevelInfo
@@ -176,116 +175,18 @@ std::vector<BattleHex> CSpell::rangeInHexes(BattleHex centralHex, ui8 schoolLvl,
 
 std::set<const CStack* > CSpell::getAffectedStacks(const CBattleInfoCallback * cb, ECastingMode::ECastingMode mode, PlayerColor casterColor, int spellLvl, BattleHex destination, const CGHeroInstance * caster) const
 {
-	std::set<const CStack* > attackedCres;//std::set to exclude multiple occurrences of two hex creatures
-	
-	const ui8 attackerSide = cb->playerToSide(casterColor) == 1;
-	const auto attackedHexes = rangeInHexes(destination, spellLvl, attackerSide);
-	
-	const CSpell::TargetInfo ti(this, spellLvl, mode);
-	
-	
-	//TODO: more generic solution for mass spells
-	if (id == SpellID::CHAIN_LIGHTNING)
-	{
-		std::set<BattleHex> possibleHexes;
-		for (auto stack : cb->battleGetAllStacks())
-		{
-			if (stack->isValidTarget())
-			{
-				for (auto hex : stack->getHexes())
-				{
-					possibleHexes.insert (hex);
-				}
-			}
-		}
-		int targetsOnLevel[4] = {4, 4, 5, 5};
-
-		BattleHex lightningHex =  destination;
-		for (int i = 0; i < targetsOnLevel[spellLvl]; ++i)
-		{
-			auto stack = cb->battleGetStackByPos (lightningHex, true);
-			if (!stack)
-				break;
-			attackedCres.insert (stack);
-			for (auto hex : stack->getHexes())
-			{
-				possibleHexes.erase (hex); //can't hit same place twice
-			}
-			if (possibleHexes.empty()) //not enough targets
-				break;
-			lightningHex = BattleHex::getClosestTile (stack->attackerOwned, destination, possibleHexes);
-		}
-	}
-	else if (getLevelInfo(spellLvl).range.size() > 1) //custom many-hex range
-	{
-		for(BattleHex hex : attackedHexes)
-		{
-			if(const CStack * st = cb->battleGetStackByPos(hex, ti.onlyAlive))
-			{
-				attackedCres.insert(st);
-			}
-		}
-	}
-	else if(getTargetType() == CSpell::CREATURE)
-	{
-		auto predicate = [=](const CStack * s){
-			const bool positiveToAlly = isPositive() && s->owner == casterColor;
-			const bool negativeToEnemy = isNegative() && s->owner != casterColor;
-			const bool validTarget = s->isValidTarget(!ti.onlyAlive); //todo: this should be handled by spell class
-	
-			//for single target spells select stacks covering destination tile
-			const bool rangeCovers = ti.massive || s->coversPos(destination);
-			//handle smart targeting
-			const bool positivenessFlag = !ti.smart || isNeutral() || positiveToAlly || negativeToEnemy;
-			
-			return rangeCovers  && positivenessFlag && validTarget;		
-		};
-		
-		TStacks stacks = cb->battleGetStacksIf(predicate);
-		
-		if (ti.massive)
-		{
-			//for massive spells add all targets
-			for (auto stack : stacks)
-				attackedCres.insert(stack);
+	ISpellMechanics::SpellTargetingContext ctx(this, cb,mode,casterColor,spellLvl,destination);
 
-		}
-		else
-		{
-			//for single target spells we must select one target. Alive stack is preferred (issue #1763)
-			for(auto stack : stacks)
-			{
-				if(stack->alive())
-				{
-					attackedCres.insert(stack);
-					break;
-				}				
-			}	
-			
-			if(attackedCres.empty() && !stacks.empty())
-			{
-				attackedCres.insert(stacks.front());
-			}						
-		}
-	}
-	else //custom range from attackedHexes
-	{
-		for(BattleHex hex : attackedHexes)
-		{
-			if(const CStack * st = cb->battleGetStackByPos(hex, ti.onlyAlive))
-				attackedCres.insert(st);
-		}
-	}	
+	std::set<const CStack* > attackedCres = mechanics->getAffectedStacks(ctx);
 	
 	//now handle immunities		
 	auto predicate = [&, this](const CStack * s)->bool
 	{
-		bool hitDirectly = ti.alwaysHitDirectly && s->coversPos(destination);
+		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;

lib/SpellMechanics.cpp

@@ -208,7 +208,76 @@ std::vector<BattleHex> DefaultSpellMechanics::rangeInHexes(BattleHex centralHex,
 
 std::set<const CStack *> DefaultSpellMechanics::getAffectedStacks(SpellTargetingContext & ctx) const
 {
+	std::set<const CStack* > attackedCres;//std::set to exclude multiple occurrences of two hex creatures
 	
+	const ui8 attackerSide = ctx.cb->playerToSide(ctx.casterColor) == 1;
+	const auto attackedHexes = rangeInHexes(ctx.destination, ctx.schoolLvl, attackerSide);
+
+	const CSpell::TargetInfo ti(owner, ctx.schoolLvl, ctx.mode);
+	
+	//TODO: more generic solution for mass spells
+	if (owner->getLevelInfo(ctx.schoolLvl).range.size() > 1) //custom many-hex range
+	{
+		for(BattleHex hex : attackedHexes)
+		{
+			if(const CStack * st = ctx.cb->battleGetStackByPos(hex, ti.onlyAlive))
+			{
+				attackedCres.insert(st);
+			}
+		}
+	}
+	else if(ti.type == CSpell::CREATURE)
+	{
+		auto predicate = [=](const CStack * s){
+			const bool positiveToAlly = owner->isPositive() && s->owner == ctx.casterColor;
+			const bool negativeToEnemy = owner->isNegative() && s->owner != ctx.casterColor;
+			const bool validTarget = s->isValidTarget(!ti.onlyAlive); //todo: this should be handled by spell class
+	
+			//for single target spells select stacks covering destination tile
+			const bool rangeCovers = ti.massive || s->coversPos(ctx.destination);
+			//handle smart targeting
+			const bool positivenessFlag = !ti.smart || owner->isNeutral() || positiveToAlly || negativeToEnemy;
+			
+			return rangeCovers && positivenessFlag && validTarget;		
+		};
+		
+		TStacks stacks = ctx.cb->battleGetStacksIf(predicate);
+		
+		if (ti.massive)
+		{
+			//for massive spells add all targets
+			for (auto stack : stacks)
+				attackedCres.insert(stack);
+
+		}
+		else
+		{
+			//for single target spells we must select one target. Alive stack is preferred (issue #1763)
+			for(auto stack : stacks)
+			{
+				if(stack->alive())
+				{
+					attackedCres.insert(stack);
+					break;
+				}				
+			}	
+			
+			if(attackedCres.empty() && !stacks.empty())
+			{
+				attackedCres.insert(stacks.front());
+			}						
+		}
+	}
+	else //custom range from attackedHexes
+	{
+		for(BattleHex hex : attackedHexes)
+		{
+			if(const CStack * st = ctx.cb->battleGetStackByPos(hex, ti.onlyAlive))
+				attackedCres.insert(st);
+		}
+	}	
+	
+	return attackedCres;
 }
 
 
@@ -218,6 +287,7 @@ ESpellCastProblem::ESpellCastProblem DefaultSpellMechanics::isImmuneByStack(cons
 	return owner->isImmuneBy(obj);
 }
 
+///WallMechanics
 std::vector<BattleHex> WallMechanics::rangeInHexes(BattleHex centralHex, ui8 schoolLvl, ui8 side, bool* outDroppedHexes) const
 {
 	using namespace SRSLPraserHelpers;
@@ -256,7 +326,42 @@ std::vector<BattleHex> WallMechanics::rangeInHexes(BattleHex centralHex, ui8 sch
 	return ret;	
 }
 
+///ChainLightningMechanics
+std::set<const CStack *> ChainLightningMechanics::getAffectedStacks(SpellTargetingContext & ctx) const
+{
+	std::set<const CStack* > attackedCres;
+	
+	std::set<BattleHex> possibleHexes;
+	for(auto stack : ctx.cb->battleGetAllStacks())
+	{
+		if(stack->isValidTarget())
+		{
+			for(auto hex : stack->getHexes())
+			{
+				possibleHexes.insert (hex);
+			}
+		}
+	}
+	int targetsOnLevel[4] = {4, 4, 5, 5};
 
+	BattleHex lightningHex = ctx.destination;
+	for(int i = 0; i < targetsOnLevel[ctx.schoolLvl]; ++i)
+	{
+		auto stack = ctx.cb->battleGetStackByPos(lightningHex, true);
+		if(!stack)
+			break;
+		attackedCres.insert (stack);
+		for(auto hex : stack->getHexes())
+		{
+			possibleHexes.erase(hex); //can't hit same place twice
+		}
+		if(possibleHexes.empty()) //not enough targets
+			break;
+		lightningHex = BattleHex::getClosestTile(stack->attackerOwned, ctx.destination, possibleHexes);
+	}	
+		
+	return attackedCres;
+}
 
 ///CloneMechanics
 ESpellCastProblem::ESpellCastProblem CloneMechanics::isImmuneByStack(const CGHeroInstance* caster, const CStack * obj) const

lib/SpellMechanics.h

@@ -11,18 +11,25 @@
 #pragma once
 
 #include "CSpellHandler.h"
+#include "BattleHex.h"
 
 class DLL_LINKAGE ISpellMechanics
 {
 public:
 	
-	struct SpellTargetingContext
+	struct DLL_LINKAGE SpellTargetingContext
 	{
-		CBattleInfoCallback * cb;
-		
+		const CBattleInfoCallback * cb;		
 		CSpell::TargetInfo ti;
-		
+		ECastingMode::ECastingMode mode;
+		BattleHex destination;
+		PlayerColor casterColor;
 		int schoolLvl;
+		
+		SpellTargetingContext(const CSpell * s, const CBattleInfoCallback * c, ECastingMode::ECastingMode m, PlayerColor cc, int lvl, BattleHex dest)
+			: cb(c), ti(s,lvl, m), mode(m), destination(dest)
+		{};
+		
 	};
 	
 public:
@@ -77,7 +84,7 @@ class ChainLightningMechanics: public DefaultSpellMechanics
 {
 public:
 	ChainLightningMechanics(CSpell * s): DefaultSpellMechanics(s){};	
-	
+	std::set<const CStack *> getAffectedStacks(SpellTargetingContext & ctx) const override;
 };
 
 class CloneMechanics: public DefaultSpellMechanics