Fixed battle replay

lib/NetPacks.h

@@ -1530,6 +1530,19 @@ struct DLL_LINKAGE BattleSetActiveStack : public CPackForClient
 	}
 };
 
+struct DLL_LINKAGE BattleCancelled: public CPackForClient
+{
+	void applyGs(CGameState * gs) const;
+
+	BattleID battleID = BattleID::NONE;
+
+	template <typename Handler> void serialize(Handler & h, const int version)
+	{
+		h & battleID;
+		assert(battleID != BattleID::NONE);
+	}
+};
+
 struct DLL_LINKAGE BattleResultAccepted : public CPackForClient
 {
 	void applyGs(CGameState * gs) const;

lib/NetPacksLib.cpp

@@ -2119,7 +2119,7 @@ void BattleStart::applyGs(CGameState * gs) const
 	info->battleID = gs->nextBattleID;
 	info->localInit();
 
-	vstd::next(gs->nextBattleID, 1);
+	gs->nextBattleID = vstd::next(gs->nextBattleID, 1);
 }
 
 void BattleNextRound::applyGs(CGameState * gs) const
@@ -2177,6 +2177,17 @@ void BattleUpdateGateState::applyGs(CGameState * gs) const
 		gs->getBattle(battleID)->si.gateState = state;
 }
 
+void BattleCancelled::applyGs(CGameState * gs) const
+{
+	auto currentBattle = boost::range::find_if(gs->currentBattles, [&](const auto & battle)
+	{
+		return battle->battleID == battleID;
+	});
+
+	assert(currentBattle != gs->currentBattles.end());
+	gs->currentBattles.erase(currentBattle);
+}
+
 void BattleResultAccepted::applyGs(CGameState * gs) const
 {
 	// Remove any "until next battle" bonuses

lib/registerTypes/RegisterTypes.h

@@ -286,6 +286,7 @@ void registerTypesClientPacks2(Serializer &s)
 	s.template registerType<CPackForClient, BattleSetActiveStack>();
 	s.template registerType<CPackForClient, BattleResult>();
 	s.template registerType<CPackForClient, BattleResultAccepted>();
+	s.template registerType<CPackForClient, BattleCancelled>();
 	s.template registerType<CPackForClient, BattleLogMessage>();
 	s.template registerType<CPackForClient, BattleStackMoved>();
 	s.template registerType<CPackForClient, BattleAttack>();

lib/spells/effects/Catapult.cpp

@@ -137,6 +137,7 @@ void Catapult::applyTargeted(ServerCallback * server, const Mechanics * m, const
 		attack.damageDealt = getRandomDamage(server);
 
 		CatapultAttack ca; //package for clients
+		ca.battleID = m->battle()->getBattle()->getBattleID();
 		ca.attacker = m->caster->getHeroCaster() ? -1 : m->caster->getCasterUnitId();
 		ca.attackedParts.push_back(attack);
 		server->apply(&ca);

server/TurnTimerHandler.cpp

@@ -254,12 +254,15 @@ void TurnTimerHandler::onBattleLoop(const BattleID & battleID, int waitTime)
 	std::lock_guard<std::recursive_mutex> guard(mx);
 	const auto * gs = gameHandler.gameState();
 	const auto * si = gameHandler.getStartInfo();
-	if(!si || !gs || !si->turnTimerInfo.isBattleEnabled())
+	if(!si || !gs)
 	{
 		assert(0);
 		return;
 	}
 	
+	if (!si->turnTimerInfo.isBattleEnabled())
+		return;
+
 	ui8 side = 0;
 	const CStack * stack = nullptr;
 	bool isTactisPhase = gs->getBattle(battleID)->battleTacticDist() > 0;
@@ -279,6 +282,7 @@ void TurnTimerHandler::onBattleLoop(const BattleID & battleID, int waitTime)
 		return;
 	
 	const auto * state = gameHandler.getPlayerState(player);
+	assert(state && state->status != EPlayerStatus::INGAME);
 	if(!state || state->status != EPlayerStatus::INGAME || !state->human)
 		return;
 	

server/battles/BattleActionProcessor.cpp

@@ -904,7 +904,7 @@ void BattleActionProcessor::makeAttack(const CBattleInfoCallback & battle, const
 		bat.flags |= BattleAttack::DEATH_BLOW;
 	}
 
-	const auto * owner = battle.battleGetFightingHero(attacker->unitOwner());
+	const auto * owner = battle.battleGetFightingHero(attacker->unitSide());
 	if(owner)
 	{
 		int chance = owner->valOfBonuses(BonusType::BONUS_DAMAGE_CHANCE, attacker->creatureIndex());

server/battles/BattleProcessor.cpp

@@ -51,9 +51,50 @@ void BattleProcessor::engageIntoBattle(PlayerColor player)
 	gameHandler->sendAndApply(&pb);
 }
 
+void BattleProcessor::restartBattlePrimary(const BattleID & battleID, const CArmedInstance *army1, const CArmedInstance *army2, int3 tile,
+								const CGHeroInstance *hero1, const CGHeroInstance *hero2, bool creatureBank,
+								const CGTownInstance *town)
+{
+	auto battle = gameHandler->gameState()->getBattle(battleID);
+
+	auto lastBattleQuery = std::dynamic_pointer_cast<CBattleQuery>(gameHandler->queries->topQuery(battle->sides[0].color));
+
+	assert(lastBattleQuery);
+
+	//existing battle query for retying auto-combat
+	if(lastBattleQuery)
+	{
+		const CGHeroInstance*heroes[2];
+		heroes[0] = hero1;
+		heroes[1] = hero2;
+
+		for(int i : {0, 1})
+		{
+			if(heroes[i])
+			{
+				SetMana restoreInitialMana;
+				restoreInitialMana.val = lastBattleQuery->initialHeroMana[i];
+				restoreInitialMana.hid = heroes[i]->id;
+				gameHandler->sendAndApply(&restoreInitialMana);
+			}
+		}
+
+		lastBattleQuery->result = std::nullopt;
+
+		assert(lastBattleQuery->belligerents[0] == battle->sides[0].armyObject);
+		assert(lastBattleQuery->belligerents[1] == battle->sides[1].armyObject);
+	}
+
+	BattleCancelled bc;
+	bc.battleID = battleID;
+	gameHandler->sendAndApply(&bc);
+
+	startBattlePrimary(army1, army2, tile, hero1, hero2, creatureBank, town);
+}
+
 void BattleProcessor::startBattlePrimary(const CArmedInstance *army1, const CArmedInstance *army2, int3 tile,
 								const CGHeroInstance *hero1, const CGHeroInstance *hero2, bool creatureBank,
-								const CGTownInstance *town) //use hero=nullptr for no hero
+								const CGTownInstance *town)
 {
 	assert(gameHandler->gameState()->getBattle(army1->getOwner()) == nullptr);
 	assert(gameHandler->gameState()->getBattle(army2->getOwner()) == nullptr);
@@ -61,10 +102,10 @@ void BattleProcessor::startBattlePrimary(const CArmedInstance *army1, const CArm
 	engageIntoBattle(army1->tempOwner);
 	engageIntoBattle(army2->tempOwner);
 
-	static const CArmedInstance *armies[2];
+	const CArmedInstance *armies[2];
 	armies[0] = army1;
 	armies[1] = army2;
-	static const CGHeroInstance*heroes[2];
+	const CGHeroInstance*heroes[2];
 	heroes[0] = hero1;
 	heroes[1] = hero2;
 
@@ -75,35 +116,21 @@ void BattleProcessor::startBattlePrimary(const CArmedInstance *army1, const CArm
 
 	auto lastBattleQuery = std::dynamic_pointer_cast<CBattleQuery>(gameHandler->queries->topQuery(battle->sides[0].color));
 
-	//existing battle query for retying auto-combat
-	if(lastBattleQuery)
+	if (lastBattleQuery)
+	{
+		lastBattleQuery->battleID = battleID;
+	}
+	else
 	{
+		auto newBattleQuery = std::make_shared<CBattleQuery>(gameHandler, battle);
+
+		// store initial mana to reset if battle has been restarted
 		for(int i : {0, 1})
-		{
 			if(heroes[i])
-			{
-				SetMana restoreInitialMana;
-				restoreInitialMana.val = lastBattleQuery->initialHeroMana[i];
-				restoreInitialMana.hid = heroes[i]->id;
-				gameHandler->sendAndApply(&restoreInitialMana);
-			}
-		}
+				newBattleQuery->initialHeroMana[i] = heroes[i]->mana;
 
-		lastBattleQuery->battleID = battle->getBattleID();
-		lastBattleQuery->result = std::nullopt;
-		lastBattleQuery->belligerents[0] = battle->sides[0].armyObject;
-		lastBattleQuery->belligerents[1] = battle->sides[1].armyObject;
-	}
-
-	auto nextBattleQuery = std::make_shared<CBattleQuery>(gameHandler, battle);
-	for(int i : {0, 1})
-	{
-		if(heroes[i])
-		{
-			nextBattleQuery->initialHeroMana[i] = heroes[i]->mana;
-		}
+		gameHandler->queries->addQuery(newBattleQuery);
 	}
-	gameHandler->queries->addQuery(nextBattleQuery);
 
 	flowProcessor->onBattleStarted(*battle);
 }

server/battles/BattleProcessor.h

@@ -63,6 +63,8 @@ public:
 	void startBattleI(const CArmedInstance *army1, const CArmedInstance *army2, int3 tile, bool creatureBank = false);
 	/// Starts battle between two armies (which can also be heroes) at position of 2nd object
 	void startBattleI(const CArmedInstance *army1, const CArmedInstance *army2, bool creatureBank = false);
+	/// Restart ongoing battle and end previous battle
+	void restartBattlePrimary(const BattleID & battleID, const CArmedInstance *army1, const CArmedInstance *army2, int3 tile, const CGHeroInstance *hero1, const CGHeroInstance *hero2, bool creatureBank = false, const CGTownInstance *town = nullptr);
 
 	/// Processing of incoming battle action netpack
 	bool makePlayerBattleAction(const BattleID & battleID, PlayerColor player, const BattleAction & ba);

server/battles/BattleResultProcessor.cpp

@@ -278,11 +278,12 @@ void BattleResultProcessor::endBattle(const CBattleInfoCallback & battle)
 	for(auto q : gameHandler->queries->allQueries())
 	{
 		auto otherBattleQuery = std::dynamic_pointer_cast<CBattleQuery>(q);
-		if(otherBattleQuery)
+		if(otherBattleQuery && otherBattleQuery->battleID == battle.getBattle()->getBattleID())
 			otherBattleQuery->result = battleQuery->result;
 	}
 
 	gameHandler->turnTimerHandler.onBattleEnd(battle.getBattle()->getBattleID());
+	gameHandler->sendAndApply(battleResult);
 
 	if (battleResult->queryID == QueryID::NONE)
 		endBattleConfirm(battle);

server/queries/BattleQueries.cpp

@@ -10,6 +10,7 @@
 #include "StdInc.h"
 #include "BattleQueries.h"
 #include "MapQueries.h"
+#include "QueriesProcessor.h"
 
 #include "../CGameHandler.h"
 #include "../battles/BattleProcessor.h"
@@ -18,6 +19,8 @@
 
 void CBattleQuery::notifyObjectAboutRemoval(const CObjectVisitQuery & objectVisit) const
 {
+	assert(result);
+
 	if(result)
 		objectVisit.visitedObject->battleFinished(objectVisit.visitingHero, *result);
 }
@@ -47,10 +50,21 @@ bool CBattleQuery::blocksPack(const CPack * pack) const
 
 void CBattleQuery::onRemoval(PlayerColor color)
 {
+	assert(result);
+
 	if(result)
 		gh->battles->battleAfterLevelUp(battleID, *result);
 }
 
+void CBattleQuery::onExposure(QueryPtr topQuery)
+{
+	// this method may be called in two cases:
+	// 1) when requesting battle replay (but before replay starts -> no valid result)
+	// 2) when aswering on levelup queries after accepting battle result -> valid result
+	if(result)
+		owner->popQuery(*this);
+}
+
 CBattleDialogQuery::CBattleDialogQuery(CGameHandler * owner, const IBattleInfo * bi):
 	CDialogQuery(owner),
 	bi(bi)
@@ -64,7 +78,8 @@ void CBattleDialogQuery::onRemoval(PlayerColor color)
 	assert(answer);
 	if(*answer == 1)
 	{
-		gh->startBattlePrimary(
+		gh->battles->restartBattlePrimary(
+			bi->getBattleID(),
 			bi->getSideArmy(0),
 			bi->getSideArmy(1),
 			bi->getLocation(),

server/queries/BattleQueries.h

@@ -31,6 +31,7 @@ public:
 	virtual void notifyObjectAboutRemoval(const CObjectVisitQuery &objectVisit) const override;
 	virtual bool blocksPack(const CPack *pack) const override;
 	virtual void onRemoval(PlayerColor color) override;
+	virtual void onExposure(QueryPtr topQuery) override;
 };
 
 class CBattleDialogQuery : public CDialogQuery