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 TStacks; +typedef std::function TStackFilter; class CBattleInfoEssentials; @@ -168,7 +169,15 @@ public: ETerrainType battleTerrainType() const; BFieldType battleGetBattlefieldType() const; std::vector > battleGetAllObstacles(boost::optional 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;