From 0e93ec28c55f87e00afd228656c6edb4959d8235 Mon Sep 17 00:00:00 2001
From: AlexVinS <alexvins@users.noreply.github.com>
Date: Sat, 17 May 2014 13:37:12 +0400
Subject: [PATCH] Add new general method for battle stack access

---
 lib/CBattleCallback.cpp | 29 ++++++++++++++++++-----------
 lib/CBattleCallback.h   | 16 +++++++++++++++-
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/lib/CBattleCallback.cpp b/lib/CBattleCallback.cpp
index bc94982ae..36159d41b 100644
--- a/lib/CBattleCallback.cpp
+++ b/lib/CBattleCallback.cpp
@@ -177,29 +177,36 @@ bool CBattleInfoEssentials::battleHasNativeStack(ui8 side) const
 	return false;
 }
 
-TStacks CBattleInfoEssentials::battleGetAllStacks(bool includeTurrets /*= false*/) const /*returns all stacks, alive or dead or undead or mechanical :) */
+TStacks CBattleInfoEssentials::battleGetAllStacks(bool includeTurrets /*= false*/) const
+{
+	return battleGetStacksIf([](const CStack * s){return true;},includeTurrets);
+}
+
+TStacks CBattleInfoEssentials::battleGetStacksIf(TStackFilter predicate, bool includeTurrets /*= false*/) const
 {
 	TStacks ret;
 	RETURN_IF_NOT_BATTLE(ret);
-	boost::copy(getBattle()->stacks, std::back_inserter(ret));
-	if(!includeTurrets)
-		vstd::erase_if(ret, [](const CStack *stack) { return stack->type->idNumber == CreatureID::ARROW_TOWERS; });
-
+	
+	vstd::copy_if(getBattle()->stacks, std::back_inserter(ret), [=](const CStack * s){
+		return predicate(s) && (includeTurrets || !(s->type->idNumber == CreatureID::ARROW_TOWERS));
+	});
+	
 	return ret;
 }
 
+
 TStacks CBattleInfoEssentials::battleAliveStacks() const
 {
-	TStacks ret;
-	vstd::copy_if(battleGetAllStacks(), std::back_inserter(ret), [](const CStack *s){ return s->alive(); });
-	return ret;
+	return battleGetStacksIf([](const CStack * s){
+		return s->alive();
+	});
 }
 
 TStacks CBattleInfoEssentials::battleAliveStacks(ui8 side) const
 {
-	TStacks ret;
-	vstd::copy_if(battleGetAllStacks(), std::back_inserter(ret), [=](const CStack *s){ return s->alive()  &&  s->attackerOwned == !side; });
-	return ret;
+	return battleGetStacksIf([=](const CStack * s){
+		return s->alive() && s->attackerOwned == !side;
+	});
 }
 
 int CBattleInfoEssentials::battleGetMoatDmg() const
diff --git a/lib/CBattleCallback.h b/lib/CBattleCallback.h
index a780c71c1..3daac2587 100644
--- a/lib/CBattleCallback.h
+++ b/lib/CBattleCallback.h
@@ -31,6 +31,7 @@ namespace BattleSide
 }
 
 typedef std::vector<const CStack*> TStacks;
+typedef std::function<bool(const CStack *)> TStackFilter;
 
 class CBattleInfoEssentials;
 
@@ -168,7 +169,15 @@ public:
 	ETerrainType battleTerrainType() const;
 	BFieldType battleGetBattlefieldType() const;
 	std::vector<shared_ptr<const CObstacleInstance> > battleGetAllObstacles(boost::optional<BattlePerspective::BattlePerspective> perspective = boost::none) const; //returns all obstacles on the battlefield
-	TStacks battleGetAllStacks(bool includeTurrets = false) const; //returns all stacks, alive or dead or undead or mechanical :)
+    
+    /** @brief Main method for getting battle stacks
+     *
+     * @param predicate Functor that shall return true for desired stack
+     * @return filtered stacks
+     *
+     */                             	
+	TStacks battleGetStacksIf(TStackFilter predicate, bool includeTurrets = false) const;
+	
 	bool battleHasNativeStack(ui8 side) const;
 	int battleGetMoatDmg() const; //what dmg unit will suffer if ending turn in the moat
 	const CGTownInstance * battleGetDefendedTown() const; //returns defended town if current battle is a siege, nullptr instead
@@ -190,7 +199,12 @@ public:
 	si8 battleGetWallState(int partOfWall) const;
 
 	//helpers
+	///returns all stacks, alive or dead or undead or mechanical :)
+	TStacks battleGetAllStacks(bool includeTurrets = false) const;
+	
+	///returns all alive stacks excluding turrets
 	TStacks battleAliveStacks() const;
+	///returns all alive stacks from particular side excluding turrets
 	TStacks battleAliveStacks(ui8 side) const;
 	const CStack * battleGetStackByID(int ID, bool onlyAlive = true) const; //returns stack info by given ID
 	bool battleIsObstacleVisibleForSide(const CObstacleInstance & coi, BattlePerspective::BattlePerspective side) const;