Rework and simplify logic of attaching heroes to bonus system

lib/battle/BattleInfo.cpp

@@ -358,13 +358,13 @@ std::unique_ptr<BattleInfo> BattleInfo::setupBattle(IGameInfoCallback *cb, const
 
 	if (currentBattle->townID.hasValue())
 	{
-		if (currentBattle->getTown()->fortificationsLevel().citadelHealth != 0)
+		if (currentBattle->getDefendedTown()->fortificationsLevel().citadelHealth != 0)
 			currentBattle->generateNewStack(currentBattle->nextUnitId(), CStackBasicDescriptor(CreatureID::ARROW_TOWERS, 1), BattleSide::DEFENDER, SlotID::ARROW_TOWERS_SLOT, BattleHex::CASTLE_CENTRAL_TOWER);
 
-		if (currentBattle->getTown()->fortificationsLevel().upperTowerHealth != 0)
+		if (currentBattle->getDefendedTown()->fortificationsLevel().upperTowerHealth != 0)
 			currentBattle->generateNewStack(currentBattle->nextUnitId(), CStackBasicDescriptor(CreatureID::ARROW_TOWERS, 1), BattleSide::DEFENDER, SlotID::ARROW_TOWERS_SLOT, BattleHex::CASTLE_UPPER_TOWER);
 
-		if (currentBattle->getTown()->fortificationsLevel().lowerTowerHealth != 0)
+		if (currentBattle->getDefendedTown()->fortificationsLevel().lowerTowerHealth != 0)
 			currentBattle->generateNewStack(currentBattle->nextUnitId(), CStackBasicDescriptor(CreatureID::ARROW_TOWERS, 1), BattleSide::DEFENDER, SlotID::ARROW_TOWERS_SLOT, BattleHex::CASTLE_BOTTOM_TOWER);
 
 		//Moat generating is done on server
@@ -568,13 +568,6 @@ const CGHeroInstance * BattleInfo::getSideHero(BattleSide side) const
 	return getSide(side).getHero();
 }
 
-const CGTownInstance * BattleInfo::getTown() const
-{
-	if (townID.hasValue())
-		return cb->getTown(townID);
-	return nullptr;
-}
-
 uint8_t BattleInfo::getTacticDist() const
 {
 	return tacticDistance;

lib/battle/BattleInfo.h

@@ -101,8 +101,6 @@ public:
 	const CArmedInstance * getSideArmy(BattleSide side) const override;
 	const CGHeroInstance * getSideHero(BattleSide side) const override;
 
-	const CGTownInstance * getTown() const;
-
 	ui8 getTacticDist() const override;
 	BattleSide getTacticsSide() const override;
 

lib/callback/IGameEventCallback.h

@@ -102,7 +102,6 @@ public:
 	virtual void startBattle(const CArmedInstance *army1, const CArmedInstance *army2, int3 tile, const CGHeroInstance *hero1, const CGHeroInstance *hero2, const BattleLayout & layout, const CGTownInstance *town)=0; //use hero=nullptr for no hero
 	virtual void startBattle(const CArmedInstance *army1, const CArmedInstance *army2)=0; //if any of armies is hero, hero will be used, visitable tile of second obj is place of battle
 	virtual bool moveHero(ObjectInstanceID hid, int3 dst, EMovementMode moveMove, bool transit = false, PlayerColor asker = PlayerColor::NEUTRAL)=0;
-	virtual bool swapGarrisonOnSiege(ObjectInstanceID tid)=0;
 	virtual void giveHeroBonus(GiveBonus * bonus)=0;
 	virtual void setMovePoints(SetMovePoints * smp)=0;
 	virtual void setMovePoints(ObjectInstanceID hid, int val, ChangeValueMode mode)=0;

lib/gameState/GameStatePackVisitor.cpp

@@ -385,7 +385,6 @@ void GameStatePackVisitor::visitRemoveObject(RemoveObject & pack)
 		auto beatenHero = dynamic_cast<CGHeroInstance*>(obj);
 		assert(beatenHero);
 
-		auto * siegeNode = beatenHero->whereShouldBeAttachedOnSiege(gs);
 		vstd::erase_if(beatenHero->artifactsInBackpack, [](const ArtSlotInfo& asi)
 		{
 			return asi.getArt()->getTypeId() == ArtifactID::GRAIL;
@@ -400,14 +399,6 @@ void GameStatePackVisitor::visitRemoveObject(RemoveObject & pack)
 
 			beatenHero->setVisitedTown(nullptr, false);
 		}
-		beatenHero->detachFromBonusSystem(gs);
-		beatenHero->tempOwner = PlayerColor::NEUTRAL; //no one owns beaten hero
-
-		// FIXME: workaround:
-		// hero should be attached to siegeNode after battle
-		// however this code might also be called on dismissing hero while in town
-		if (siegeNode && vstd::contains(beatenHero->getParentNodes(), siegeNode))
-			beatenHero->detachFrom(*siegeNode);
 
 		//If hero on Boat is removed, the Boat disappears
 		if(beatenHero->inBoat())
@@ -417,12 +408,13 @@ void GameStatePackVisitor::visitRemoveObject(RemoveObject & pack)
 			gs.getMap().eraseObject(boat->id);
 		}
 
+		beatenHero->detachFromBonusSystem(gs);
+		beatenHero->tempOwner = PlayerColor::NEUTRAL; //no one owns beaten hero
 		auto beatenObject = gs.getMap().eraseObject(obj->id);
 
 		//return hero to the pool, so he may reappear in tavern
 		gs.heroesPool->addHeroToPool(beatenHero->getHeroTypeID());
 		gs.getMap().addToHeroPool(std::dynamic_pointer_cast<CGHeroInstance>(beatenObject));
-
 		return;
 	}
 
@@ -1173,6 +1165,18 @@ void GameStatePackVisitor::visitBattleStart(BattleStart & pack)
 	pack.info->battleID = gs.nextBattleID;
 	pack.info->localInit();
 
+	if (pack.info->getDefendedTown() && pack.info->getSideHero(BattleSide::DEFENDER))
+	{
+		CGTownInstance * town = gs.getTown(pack.info->townID);
+		CGHeroInstance * hero = gs.getHero(pack.info->getSideHero(BattleSide::DEFENDER)->id);
+
+		if (town->getVisitingHero() == hero)
+		{
+			hero->detachFrom(town->townAndVis);
+			hero->attachTo(*town);
+		}
+	}
+
 	gs.currentBattles.push_back(std::move(pack.info));
 	gs.nextBattleID = BattleID(gs.nextBattleID.getNum() + 1);
 }
@@ -1232,17 +1236,6 @@ void GameStatePackVisitor::visitBattleUpdateGateState(BattleUpdateGateState & pa
 		gs.getBattle(pack.battleID)->si.gateState = pack.state;
 }
 
-void GameStatePackVisitor::visitBattleCancelled(BattleCancelled & pack)
-{
-	auto currentBattle = boost::range::find_if(gs.currentBattles, [&](const auto & battle)
-	{
-		return battle->battleID == pack.battleID;
-	});
-
-	assert(currentBattle != gs.currentBattles.end());
-	gs.currentBattles.erase(currentBattle);
-}
-
 void GameStatePackVisitor::visitBattleResultAccepted(BattleResultAccepted & pack)
 {
 	// Remove any "until next battle" bonuses
@@ -1348,8 +1341,44 @@ void GameStatePackVisitor::visitBattleUnitsChanged(BattleUnitsChanged & pack)
 	pack.visitTyped(battleVisitor);
 }
 
+void GameStatePackVisitor::restorePreBattleState(BattleID battleID)
+{
+	auto battleIter = boost::range::find_if(gs.currentBattles, [&](const auto & battle)
+	{
+		return battle->battleID == battleID;
+	});
+
+	const auto & currentBattle = **battleIter;
+
+	if (currentBattle.getDefendedTown() && currentBattle.getSideHero(BattleSide::DEFENDER))
+	{
+		CGTownInstance * town = gs.getTown(currentBattle.townID);
+		CGHeroInstance * hero = gs.getHero(currentBattle.getSideHero(BattleSide::DEFENDER)->id);
+
+		if (town->getVisitingHero() == hero)
+		{
+			hero->detachFrom(*town);
+			hero->attachTo(town->townAndVis);
+		}
+	}
+}
+
+void GameStatePackVisitor::visitBattleCancelled(BattleCancelled & pack)
+{
+	restorePreBattleState(pack.battleID);
+
+	auto battleIter = boost::range::find_if(gs.currentBattles, [&](const auto & battle)
+	{
+		return battle->battleID == pack.battleID;
+	});
+
+	assert(battleIter != gs.currentBattles.end());
+	gs.currentBattles.erase(battleIter);
+}
+
 void GameStatePackVisitor::visitBattleResultsApplied(BattleResultsApplied & pack)
 {
+	restorePreBattleState(pack.battleID);
 	pack.learnedSpells.visit(*this);
 
 	for(auto & discharging : pack.dischargingArtifacts)

lib/gameState/GameStatePackVisitor.h

@@ -17,6 +17,7 @@ class CGameState;
 
 class GameStatePackVisitor final : public ICPackVisitor
 {
+	void restorePreBattleState(BattleID battleID);
 private:
 	CGameState & gs;
 

lib/mapObjects/CGHeroInstance.cpp

@@ -1346,25 +1346,6 @@ void CGHeroInstance::detachFromBonusSystem(CGameState & gs)
 	}
 }
 
-CBonusSystemNode * CGHeroInstance::whereShouldBeAttachedOnSiege(const bool isBattleOutsideTown) const
-{
-	if(!getVisitedTown())
-		return nullptr;
-
-	if (isBattleOutsideTown)
-		return const_cast<CTownAndVisitingHero *>(&getVisitedTown()->townAndVis);
-
-	return const_cast<CGTownInstance*>(getVisitedTown());
-}
-
-CBonusSystemNode * CGHeroInstance::whereShouldBeAttachedOnSiege(CGameState & gs)
-{
-	if(getVisitedTown())
-		return whereShouldBeAttachedOnSiege(getVisitedTown()->isBattleOutsideTown(this));
-
-	return &CArmedInstance::whereShouldBeAttached(gs);
-}
-
 CBonusSystemNode & CGHeroInstance::whereShouldBeAttached(CGameState & gs)
 {
 	if(visitedTown.hasValue())

lib/mapObjects/CGHeroInstance.h

@@ -276,9 +276,6 @@ public:
 	///IConstBonusProvider
 	const IBonusBearer* getBonusBearer() const override;
 
-	CBonusSystemNode * whereShouldBeAttachedOnSiege(const bool isBattleOutsideTown) const;
-	CBonusSystemNode * whereShouldBeAttachedOnSiege(CGameState & gs);
-
 	///spells::Caster
 	int32_t getCasterUnitId() const override;
 	int32_t getSpellSchoolLevel(const spells::Spell * spell, SpellSchool * outSelectedSchool = nullptr) const override;

lib/mapObjects/CGTownInstance.cpp

@@ -307,19 +307,12 @@ void CGTownInstance::onHeroVisit(IGameEventCallback & gameEvents, const CGHeroIn
 		if(armedGarrison() || getVisitingHero())
 		{
 			const CGHeroInstance * defendingHero = getVisitingHero() ? getVisitingHero() : getGarrisonHero();
-			const CArmedInstance * defendingArmy = defendingHero ? (CArmedInstance *)defendingHero : this;
+			const CArmedInstance * defendingArmy = defendingHero ? static_cast<const CArmedInstance*>(defendingHero) : this;
 			const bool isBattleOutside = isBattleOutsideTown(defendingHero);
 
-			if(!isBattleOutside && getVisitingHero() && defendingHero == getVisitingHero())
-			{
-				//we have two approaches to merge armies: mergeGarrisonOnSiege() and used in the CGameHandler::garrisonSwap(ObjectInstanceID tid)
-				auto * nodeSiege = defendingHero->whereShouldBeAttachedOnSiege(isBattleOutside);
-
-				if(nodeSiege == (CBonusSystemNode *)this)
-					gameEvents.swapGarrisonOnSiege(this->id);
+			if(!isBattleOutside && defendingHero == getVisitingHero())
+				mergeGarrisonOnSiege(gameEvents);
 
-				const_cast<CGHeroInstance *>(defendingHero)->setVisitedTown(this, false); //hack to return visitor from garrison after battle
-			}
 			gameEvents.startBattle(h, defendingArmy, getSightCenter(), h, defendingHero, BattleLayout::createDefaultLayout(*cb, h, defendingArmy), (isBattleOutside ? nullptr : this));
 		}
 		else

server/CGameHandler.cpp

@@ -2548,33 +2548,6 @@ void CGameHandler::moveArmy(const CArmedInstance *src, const CArmedInstance *dst
 	}
 }
 
-bool CGameHandler::swapGarrisonOnSiege(ObjectInstanceID tid)
-{
-	const CGTownInstance * town = gameInfo().getTown(tid);
-
-	if(!town->getGarrisonHero() == !town->getVisitingHero())
-		return false;
-
-	SetHeroesInTown intown;
-	intown.tid = tid;
-
-	if(town->getGarrisonHero()) //garrison -> vising
-	{
-		intown.garrison = ObjectInstanceID();
-		intown.visiting = town->getGarrisonHero()->id;
-	}
-	else //visiting -> garrison
-	{
-		if(town->armedGarrison())
-			town->mergeGarrisonOnSiege(*this);
-
-		intown.visiting = ObjectInstanceID();
-		intown.garrison = town->getVisitingHero()->id;
-	}
-	sendAndApply(intown);
-	return true;
-}
-
 bool CGameHandler::garrisonSwap(ObjectInstanceID tid)
 {
 	const CGTownInstance * town = gameInfo().getTown(tid);

server/CGameHandler.h

@@ -224,7 +224,6 @@ public:
 	//void lootArtifacts (TArtHolder source, TArtHolder dest, std::vector<ui32> &arts); //after battle - move al arts to winer
 	bool buySecSkill( const IMarket *m, const CGHeroInstance *h, SecondarySkill skill);
 	bool garrisonSwap(ObjectInstanceID tid);
-	bool swapGarrisonOnSiege(ObjectInstanceID tid) override;
 	bool upgradeCreature( ObjectInstanceID objid, SlotID pos, CreatureID upgID );
 	bool recruitCreatures(ObjectInstanceID objid, ObjectInstanceID dst, CreatureID crid, int32_t cram, int32_t level, PlayerColor player);
 	bool buildStructure(ObjectInstanceID tid, BuildingID bid, bool force=false);//force - for events: no cost, no checkings

server/battles/BattleResultProcessor.cpp

@@ -322,14 +322,6 @@ void BattleResultProcessor::endBattleConfirm(const CBattleInfoCallback & battle)
 	const auto winnerHero = battle.battleGetFightingHero(finishingBattle->winnerSide);
 	const auto loserHero = battle.battleGetFightingHero(CBattleInfoEssentials::otherSide(finishingBattle->winnerSide));
 
-	if(battleResult->winner == BattleSide::DEFENDER
-	   && winnerHero
-	   && winnerHero->getVisitedTown()
-	   && !winnerHero->isGarrisoned()
-	   && winnerHero->getVisitedTown()->getGarrisonHero() == winnerHero)
-	{
-		gameHandler->swapGarrisonOnSiege(winnerHero->getVisitedTown()->id); //return defending visitor from garrison to its rightful place
-	}
 	//give exp
 	if(!finishingBattle->isDraw() && battleResult->exp[finishingBattle->winnerSide])
 	{

test/mock/mock_IGameEventCallback.h

@@ -72,7 +72,6 @@ public:
 	void startBattle(const CArmedInstance *army1, const CArmedInstance *army2, int3 tile, const CGHeroInstance *hero1, const CGHeroInstance *hero2, const BattleLayout & layout, const CGTownInstance *town) override {} //use hero=nullptr for no hero
 	void startBattle(const CArmedInstance *army1, const CArmedInstance *army2) override {}
 	bool moveHero(ObjectInstanceID hid, int3 dst, EMovementMode movementMode, bool transit, PlayerColor asker) override {return false;}
-	bool swapGarrisonOnSiege(ObjectInstanceID tid) override {return false;}
 	void giveHeroBonus(GiveBonus * bonus) override {}
 	void setMovePoints(SetMovePoints * smp) override {}
 	void setMovePoints(ObjectInstanceID hid, int val, ChangeValueMode mode) override {};