Merge pull request #4996 from MichalZr6/hill_fort_fix_for_AI

Enable AI to use all Hill Fort upgrades, including alternative ones

AI/Nullkiller/AIGateway.cpp

@@ -19,6 +19,7 @@
 #include "../../lib/CConfigHandler.h"
 #include "../../lib/IGameSettings.h"
 #include "../../lib/gameState/CGameState.h"
+#include "../../lib/gameState/UpgradeInfo.h"
 #include "../../lib/serializer/CTypeList.h"
 #include "../../lib/networkPacks/PacksForClient.h"
 #include "../../lib/networkPacks/PacksForClientBattle.h"
@@ -788,14 +789,30 @@ bool AIGateway::makePossibleUpgrades(const CArmedInstance * obj)
 	{
 		if(const CStackInstance * s = obj->getStackPtr(SlotID(i)))
 		{
-			UpgradeInfo ui;
-			myCb->fillUpgradeInfo(obj, SlotID(i), ui);
-			if(ui.oldID != CreatureID::NONE && nullkiller->getFreeResources().canAfford(ui.cost[0] * s->count))
+			UpgradeInfo upgradeInfo(s->getId());
+			do
 			{
-				myCb->upgradeCreature(obj, SlotID(i), ui.newID[0]);
-				upgraded = true;
-				logAi->debug("Upgraded %d %s to %s", s->count, ui.oldID.toCreature()->getNamePluralTranslated(), ui.newID[0].toCreature()->getNamePluralTranslated());
+				myCb->fillUpgradeInfo(obj, SlotID(i), upgradeInfo);
+
+				if(upgradeInfo.hasUpgrades())
+				{
+					// creature at given slot might have alternative upgrades, pick best one
+					CreatureID upgID = *vstd::maxElementByFun(upgradeInfo.getAvailableUpgrades(), [](const CreatureID & id)
+						{
+							return id.toCreature()->getAIValue();
+						});
+					if(nullkiller->getFreeResources().canAfford(upgradeInfo.getUpgradeCostsFor(upgID) * s->count))
+					{
+						myCb->upgradeCreature(obj, SlotID(i), upgID);
+						upgraded = true;
+						logAi->debug("Upgraded %d %s to %s", s->count, upgradeInfo.oldID.toCreature()->getNamePluralTranslated(), 
+							upgradeInfo.getUpgrade().toCreature()->getNamePluralTranslated());
+					}
+					else
+						break;
+				}
 			}
+			while(upgradeInfo.hasUpgrades());
 		}
 	}
 

AI/VCAI/VCAI.cpp

@@ -22,6 +22,7 @@
 #include "../../lib/CConfigHandler.h"
 #include "../../lib/IGameSettings.h"
 #include "../../lib/gameState/CGameState.h"
+#include "../../lib/gameState/UpgradeInfo.h"
 #include "../../lib/bonuses/Limiters.h"
 #include "../../lib/bonuses/Updaters.h"
 #include "../../lib/bonuses/Propagators.h"
@@ -754,12 +755,29 @@ void makePossibleUpgrades(const CArmedInstance * obj)
 	{
 		if(const CStackInstance * s = obj->getStackPtr(SlotID(i)))
 		{
-			UpgradeInfo ui;
-			cb->fillUpgradeInfo(obj, SlotID(i), ui);
-			if(ui.oldID != CreatureID::NONE && cb->getResourceAmount().canAfford(ui.cost[0] * s->count))
+			UpgradeInfo upgradeInfo(s->getId());
+			do
 			{
-				cb->upgradeCreature(obj, SlotID(i), ui.newID[0]);
+				cb->fillUpgradeInfo(obj, SlotID(i), upgradeInfo);
+
+				if(upgradeInfo.hasUpgrades())
+				{
+					// creature at given slot might have alternative upgrades, pick best one
+					CreatureID upgID = *vstd::maxElementByFun(upgradeInfo.getAvailableUpgrades(), [](const CreatureID & id)
+						{
+							return id.toCreature()->getAIValue();
+						});
+					if(cb->getResourceAmount().canAfford(upgradeInfo.getUpgradeCostsFor(upgID) * s->count))
+					{
+						cb->upgradeCreature(obj, SlotID(i), upgID);
+						logAi->debug("Upgraded %d %s to %s", s->count, upgradeInfo.oldID.toCreature()->getNamePluralTranslated(), 
+							upgradeInfo.getUpgrade().toCreature()->getNamePluralTranslated());
+					}
+					else
+						break;
+				}
 			}
+			while(upgradeInfo.hasUpgrades());
 		}
 	}
 }

client/CPlayerInterface.h

@@ -24,7 +24,7 @@ class CCreature;
 struct CGPath;
 class CCreatureSet;
 class CGObjectInstance;
-struct UpgradeInfo;
+class UpgradeInfo;
 class ConditionalWait;
 struct CPathsInfo;
 

client/widgets/CGarrisonInt.cpp

@@ -33,6 +33,7 @@
 #include "../../lib/mapObjects/CGHeroInstance.h"
 #include "../../lib/networkPacks/ArtifactLocation.h"
 #include "../../lib/gameState/CGameState.h"
+#include "../../lib/gameState/UpgradeInfo.h"
 
 void CGarrisonSlot::setHighlight(bool on)
 {
@@ -162,10 +163,10 @@ std::function<void()> CGarrisonSlot::getDismiss() const
 /// @return Whether the view should be refreshed
 bool CGarrisonSlot::viewInfo()
 {
-	UpgradeInfo pom;
+	UpgradeInfo pom(ID.getNum());
 	LOCPLINT->cb->fillUpgradeInfo(getObj(), ID, pom);
 
-	bool canUpgrade = getObj()->tempOwner == LOCPLINT->playerID && pom.oldID != CreatureID::NONE; //upgrade is possible
+	bool canUpgrade = getObj()->tempOwner == LOCPLINT->playerID && pom.canUpgrade(); //upgrade is possible
 	std::function<void(CreatureID)> upgr = nullptr;
 	auto dism = getDismiss();
 	if(canUpgrade) upgr = [=] (CreatureID newID) { LOCPLINT->cb->upgradeCreature(getObj(), ID, newID); };

client/windows/CCreatureWindow.cpp

@@ -35,6 +35,7 @@
 #include "../../lib/IGameSettings.h"
 #include "../../lib/entities/hero/CHeroHandler.h"
 #include "../../lib/gameState/CGameState.h"
+#include "../../lib/gameState/UpgradeInfo.h"
 #include "../../lib/networkPacks/ArtifactLocation.h"
 #include "../../lib/texts/CGeneralTextHandler.h"
 #include "../../lib/texts/TextOperations.h"
@@ -57,6 +58,10 @@ public:
 	};
 	struct StackUpgradeInfo
 	{
+		StackUpgradeInfo() = delete;
+		StackUpgradeInfo(const UpgradeInfo & upgradeInfo)
+			: info(upgradeInfo)
+		{ }
 		UpgradeInfo info;
 		std::function<void(CreatureID)> callback;
 	};
@@ -355,15 +360,15 @@ CStackWindow::ButtonsSection::ButtonsSection(CStackWindow * owner, int yOffset)
 		// besides - should commander really be upgradeable?
 
 		auto & upgradeInfo = parent->info->upgradeInfo.value();
-		const size_t buttonsToCreate = std::min<size_t>(upgradeInfo.info.newID.size(), upgrade.size());
+		const size_t buttonsToCreate = std::min<size_t>(upgradeInfo.info.size(), upgrade.size());
 
 		for(size_t buttonIndex = 0; buttonIndex < buttonsToCreate; buttonIndex++)
 		{
-			TResources totalCost = upgradeInfo.info.cost[buttonIndex] * parent->info->creatureCount;
+			TResources totalCost = upgradeInfo.info.getAvailableUpgradeCosts().at(buttonIndex) * parent->info->creatureCount;
 
 			auto onUpgrade = [=]()
 			{
-				upgradeInfo.callback(upgradeInfo.info.newID[buttonIndex]);
+				upgradeInfo.callback(upgradeInfo.info.getAvailableUpgrades().at(buttonIndex));
 				parent->close();
 			};
 			auto onClick = [=]()
@@ -385,7 +390,7 @@ CStackWindow::ButtonsSection::ButtonsSection(CStackWindow * owner, int yOffset)
 			};
 			auto upgradeBtn = std::make_shared<CButton>(Point(221 + (int)buttonIndex * 40, 5), AnimationPath::builtin("stackWindow/upgradeButton"), CGI->generaltexth->zelp[446], onClick);
 
-			upgradeBtn->setOverlay(std::make_shared<CAnimImage>(AnimationPath::builtin("CPRSMALL"), VLC->creh->objects[upgradeInfo.info.newID[buttonIndex]]->getIconIndex()));
+			upgradeBtn->setOverlay(std::make_shared<CAnimImage>(AnimationPath::builtin("CPRSMALL"), VLC->creh->objects[upgradeInfo.info.getAvailableUpgrades()[buttonIndex]]->getIconIndex()));
 
 			if(buttonsToCreate == 1) // single upgrade available
 				upgradeBtn->assignedKey = EShortcut::RECRUITMENT_UPGRADE;
@@ -763,9 +768,8 @@ CStackWindow::CStackWindow(const CStackInstance * stack, std::function<void()> d
 	info->creature = stack->getCreature();
 	info->creatureCount = stack->count;
 
-	info->upgradeInfo = std::make_optional(UnitView::StackUpgradeInfo());
+	info->upgradeInfo = std::make_optional(UnitView::StackUpgradeInfo(upgradeInfo));
 	info->dismissInfo = std::make_optional(UnitView::StackDismissInfo());
-	info->upgradeInfo->info = upgradeInfo;
 	info->upgradeInfo->callback = callback;
 	info->dismissInfo->callback = dismiss;
 	info->owner = dynamic_cast<const CGHeroInstance *> (stack->armyObj);

client/windows/CCreatureWindow.h

@@ -19,7 +19,7 @@ VCMI_LIB_NAMESPACE_BEGIN
 class CCommanderInstance;
 class CStackInstance;
 class CStack;
-struct UpgradeInfo;
+class UpgradeInfo;
 
 VCMI_LIB_NAMESPACE_END
 

client/windows/GUIClasses.cpp

@@ -53,6 +53,7 @@
 #include "../lib/gameState/CGameState.h"
 #include "../lib/gameState/SThievesGuildInfo.h"
 #include "../lib/gameState/TavernHeroesPool.h"
+#include "../lib/gameState/UpgradeInfo.h"
 #include "../lib/texts/CGeneralTextHandler.h"
 #include "../lib/IGameSettings.h"
 #include "ConditionalWait.h"
@@ -1153,12 +1154,15 @@ void CHillFortWindow::updateGarrisons()
 		State newState = getState(SlotID(i));
 		if(newState != State::EMPTY)
 		{
-			UpgradeInfo info;
-			LOCPLINT->cb->fillUpgradeInfo(hero, SlotID(i), info);
-			if(info.newID.size())//we have upgrades here - update costs
+			if(const CStackInstance * s = hero->getStackPtr(SlotID(i)))
 			{
-				costs[i] = info.cost.back() * hero->getStackCount(SlotID(i));
-				totalSum += costs[i];
+				UpgradeInfo info(s->getCreature()->getId());
+				LOCPLINT->cb->fillUpgradeInfo(hero, SlotID(i), info);
+				if(info.canUpgrade())	//we have upgrades here - update costs
+				{
+					costs[i] = info.getUpgradeCosts() * hero->getStackCount(SlotID(i));
+					totalSum += costs[i];
+				}
 			}
 		}
 
@@ -1264,9 +1268,12 @@ void CHillFortWindow::makeDeal(SlotID slot)
 			{
 				if(slot.getNum() == i || ( slot.getNum() == slotsCount && currState[i] == State::MAKE_UPGRADE ))//this is activated slot or "upgrade all"
 				{
-					UpgradeInfo info;
-					LOCPLINT->cb->fillUpgradeInfo(hero, SlotID(i), info);
-					LOCPLINT->cb->upgradeCreature(hero, SlotID(i), info.newID.back());
+					if(const CStackInstance * s = hero->getStackPtr(SlotID(i)))
+					{
+						UpgradeInfo info(s->getCreatureID());
+						LOCPLINT->cb->fillUpgradeInfo(hero, SlotID(i), info);
+						LOCPLINT->cb->upgradeCreature(hero, SlotID(i), info.getUpgrade());
+					}
 				}
 			}
 			break;
@@ -1295,18 +1302,15 @@ CHillFortWindow::State CHillFortWindow::getState(SlotID slot)
 	if(hero->slotEmpty(slot))
 		return State::EMPTY;
 
-	UpgradeInfo info;
+	UpgradeInfo info(hero->getStackPtr(slot)->getCreatureID());
 	LOCPLINT->cb->fillUpgradeInfo(hero, slot, info);
-	if (info.newID.empty())
-	{
-		// Hill Fort may limit level of upgradeable creatures, e.g. mini Hill Fort from HOTA
-		if (hero->getCreature(slot)->hasUpgrades())
-			return State::UNAVAILABLE;
+	if(info.hasUpgrades() && !info.canUpgrade())
+		return State::UNAVAILABLE;  // Hill Fort may limit level of upgradeable creatures, e.g. mini Hill Fort from HOTA
 
+	if(!info.hasUpgrades())
 		return State::ALREADY_UPGRADED;
-	}
 
-	if(!(info.cost.back() * hero->getStackCount(slot)).canBeAfforded(myRes))
+	if(!(info.getUpgradeCosts() * hero->getStackCount(slot)).canBeAfforded(myRes))
 		return State::UNAFFORDABLE;
 
 	return State::MAKE_UPGRADE;

lib/CGameInfoCallback.cpp

@@ -183,7 +183,7 @@ const IMarket * CGameInfoCallback::getMarket(ObjectInstanceID objid) const
 		return nullptr;
 }
 
-void CGameInfoCallback::fillUpgradeInfo(const CArmedInstance *obj, SlotID stackPos, UpgradeInfo &out) const
+void CGameInfoCallback::fillUpgradeInfo(const CArmedInstance *obj, SlotID stackPos, UpgradeInfo & out) const
 {
 	//boost::shared_lock<boost::shared_mutex> lock(*gs->mx);
 	ERROR_RET_IF(!canGetFullInfo(obj), "Cannot get info about not owned object!");

lib/CGameInfoCallback.h

@@ -33,7 +33,7 @@ struct CPathsInfo;
 struct InfoAboutHero;
 struct InfoAboutTown;
 
-struct UpgradeInfo;
+class UpgradeInfo;
 struct SThievesGuildInfo;
 class CMapHeader;
 struct TeamState;
@@ -172,7 +172,7 @@ public:
 
 
 	//armed object
-	virtual void fillUpgradeInfo(const CArmedInstance *obj, SlotID stackPos, UpgradeInfo &out)const;
+	virtual void fillUpgradeInfo(const CArmedInstance *obj, SlotID stackPos, UpgradeInfo &out) const;
 
 	//hero
 	const CGHeroInstance * getHero(ObjectInstanceID objid) const override;

lib/CMakeLists.txt

@@ -109,6 +109,7 @@ set(lib_MAIN_SRCS
 	gameState/RumorState.cpp
 	gameState/TavernHeroesPool.cpp
 	gameState/GameStatistics.cpp
+	gameState/UpgradeInfo.cpp
 
 	mapObjectConstructors/AObjectTypeHandler.cpp
 	mapObjectConstructors/CBankInstanceConstructor.cpp

lib/gameState/CGameState.cpp

@@ -52,6 +52,7 @@
 #include "../rmg/CMapGenerator.h"
 #include "../serializer/CMemorySerializer.h"
 #include "../spells/CSpellHandler.h"
+#include "UpgradeInfo.h"
 
 #include <vstd/RNG.h>
 
@@ -1088,10 +1089,11 @@ void CGameState::fillUpgradeInfo(const CArmedInstance *obj, SlotID stackPos, Upg
 	out = fillUpgradeInfo(obj->getStack(stackPos));
 }
 
-UpgradeInfo CGameState::fillUpgradeInfo(const CStackInstance &stack) const
+UpgradeInfo CGameState::fillUpgradeInfo(const CStackInstance & stack) const
 {
-	UpgradeInfo ret;
 	const CCreature *base = stack.getCreature();
+	
+	UpgradeInfo ret(base->getId());
 
 	if (stack.armyObj->ID == Obj::HERO)
 	{
@@ -1117,12 +1119,6 @@ UpgradeInfo CGameState::fillUpgradeInfo(const CStackInstance &stack) const
 		town->fillUpgradeInfo(ret, stack);
 	}
 
-	if(!ret.newID.empty())
-		ret.oldID = base->getId();
-
-	for (ResourceSet &cost : ret.cost)
-		cost.positive(); //upgrade cost can't be negative, ignore missing resources
-
 	return ret;
 }
 

lib/gameState/CGameState.h

@@ -38,16 +38,8 @@ class TavernHeroesPool;
 struct SThievesGuildInfo;
 class CRandomGenerator;
 class GameSettings;
-
-struct UpgradeInfo
-{
-	CreatureID oldID; //creature to be upgraded
-	std::vector<CreatureID> newID; //possible upgrades
-	std::vector<ResourceSet> cost; // cost[upgrade_serial] -> set of pairs<resource_ID,resource_amount>; cost is for single unit (not entire stack)
-	UpgradeInfo(){oldID = CreatureID::NONE;};
-};
-
 class BattleInfo;
+class UpgradeInfo;
 
 DLL_LINKAGE std::ostream & operator<<(std::ostream & os, const EVictoryLossCheckResult & victoryLossCheckResult);
 

lib/gameState/UpgradeInfo.cpp

@@ -0,0 +1,36 @@
+/*
+ * UpgradeInfo.cpp, part of VCMI engine
+ *
+ * Authors: listed in file AUTHORS in main folder
+ *
+ * License: GNU General Public License v2.0 or later
+ * Full text of license available in license.txt file, in main folder
+ *
+ */
+#include "StdInc.h"
+#include "UpgradeInfo.h"
+#include "CCreatureHandler.h"
+
+VCMI_LIB_NAMESPACE_BEGIN
+
+void UpgradeInfo::addUpgrade(const CreatureID & upgradeID, const Creature * creature, int costPercentageModifier)
+{
+	isAvailable = costPercentageModifier >= 0;
+
+	upgradesIDs.push_back(upgradeID);
+
+	ResourceSet upgradeCost = (upgradeID.toCreature()->getFullRecruitCost() - creature->getFullRecruitCost()) * costPercentageModifier / 100;
+	upgradeCost.positive(); //upgrade cost can't be negative, ignore missing resources
+	upgradesCosts.push_back(std::move(upgradeCost));
+
+	// sort from highest ID to smallest
+	size_t pos = upgradesIDs.size() - 1;
+	while(pos > 0 && upgradesIDs[pos] > upgradesIDs[pos - 1])
+	{
+		std::swap(upgradesIDs[pos], upgradesIDs[pos - 1]);
+		std::swap(upgradesCosts[pos], upgradesCosts[pos - 1]);
+		--pos;
+	}
+}
+
+VCMI_LIB_NAMESPACE_END

lib/gameState/UpgradeInfo.h

@@ -0,0 +1,82 @@
+/*
+ * UpgradeInfo.h, part of VCMI engine
+ *
+ * Authors: listed in file AUTHORS in main folder
+ *
+ * License: GNU General Public License v2.0 or later
+ * Full text of license available in license.txt file, in main folder
+ *
+ */
+#pragma once
+
+#include "../lib/constants/EntityIdentifiers.h"
+#include "../lib/ResourceSet.h"
+
+VCMI_LIB_NAMESPACE_BEGIN
+
+class DLL_LINKAGE UpgradeInfo
+{
+public:
+	UpgradeInfo() = delete;
+	UpgradeInfo(CreatureID base)
+		: oldID(base), isAvailable(true)
+	{
+	}
+
+	CreatureID oldID; //creature to be upgraded
+
+	const std::vector<CreatureID> & getAvailableUpgrades() const
+	{
+		return upgradesIDs;
+	}
+
+	const CreatureID & getUpgrade() const
+	{
+		return upgradesIDs.back();
+	}
+
+	const ResourceSet & getUpgradeCostsFor(CreatureID id) const
+	{
+		auto idIt = std::find(upgradesIDs.begin(), upgradesIDs.end(), id);
+
+		assert(idIt != upgradesIDs.end());
+
+		return upgradesCosts[std::distance(upgradesIDs.begin(), idIt)];
+	}
+
+	const std::vector<ResourceSet> & getAvailableUpgradeCosts() const
+	{
+		return upgradesCosts;
+	}
+
+	const ResourceSet & getUpgradeCosts() const
+	{
+		return upgradesCosts.back();
+	}
+
+	bool canUpgrade() const
+	{
+		return !upgradesIDs.empty() && isAvailable;
+	}
+
+	bool hasUpgrades() const
+	{
+		return !upgradesIDs.empty();
+	}
+
+	// Adds a new upgrade and ensures alignment and sorted order
+	void addUpgrade(const CreatureID & upgradeID, const Creature * creature, int costPercentageModifier = 100);
+
+	auto size() const
+	{
+		return upgradesIDs.size();
+	}
+
+private:
+	std::vector<CreatureID> upgradesIDs; //possible upgrades
+	std::vector<ResourceSet> upgradesCosts; // cost[upgrade_serial] -> set of pairs<resource_ID,resource_amount>; cost is for single unit (not entire stack)
+	bool isAvailable;		// flag for unavailableUpgrades like in miniHillFort from HoTA
+};
+
+VCMI_LIB_NAMESPACE_END
+

lib/mapObjects/CGHeroInstance.cpp

@@ -25,6 +25,7 @@
 #include "../CSkillHandler.h"
 #include "../IGameCallback.h"
 #include "../gameState/CGameState.h"
+#include "../gameState/UpgradeInfo.h"
 #include "../CCreatureHandler.h"
 #include "../mapping/CMap.h"
 #include "../StartInfo.h"
@@ -1847,7 +1848,7 @@ bool CGHeroInstance::isMissionCritical() const
 	return false;
 }
 
-void CGHeroInstance::fillUpgradeInfo(UpgradeInfo & info, const CStackInstance &stack) const
+void CGHeroInstance::fillUpgradeInfo(UpgradeInfo & info, const CStackInstance & stack) const
 {
 	TConstBonusListPtr lista = getBonuses(Selector::typeSubtype(BonusType::SPECIAL_UPGRADE, BonusSubtypeID(stack.getId())));
 	for(const auto & it : *lista)
@@ -1855,8 +1856,7 @@ void CGHeroInstance::fillUpgradeInfo(UpgradeInfo & info, const CStackInstance &s
 		auto nid = CreatureID(it->additionalInfo[0]);
 		if (nid != stack.getId()) //in very specific case the upgrade is available by default (?)
 		{
-			info.newID.push_back(nid);
-			info.cost.push_back(nid.toCreature()->getFullRecruitCost() - stack.getType()->getFullRecruitCost());
+			info.addUpgrade(nid, stack.getType());
 		}
 	}
 }

lib/mapObjects/CGHeroInstance.h

@@ -23,6 +23,7 @@ class CHero;
 class CGBoat;
 class CGTownInstance;
 class CMap;
+class UpgradeInfo;
 struct TerrainTile;
 struct TurnInfo;
 

lib/mapObjects/CGTownInstance.cpp

@@ -20,6 +20,7 @@
 #include "../texts/CGeneralTextHandler.h"
 #include "../IGameCallback.h"
 #include "../gameState/CGameState.h"
+#include "../gameState/UpgradeInfo.h"
 #include "../mapping/CMap.h"
 #include "../CPlayerState.h"
 #include "../StartInfo.h"
@@ -1234,8 +1235,7 @@ void CGTownInstance::fillUpgradeInfo(UpgradeInfo & info, const CStackInstance &s
 			{
 				if(vstd::contains(stack.getCreature()->upgrades, upgrID)) //possible upgrade
 				{
-					info.newID.push_back(upgrID);
-					info.cost.push_back(upgrID.toCreature()->getFullRecruitCost() - stack.getType()->getFullRecruitCost());
+					info.addUpgrade(upgrID, stack.getType());
 				}
 			}
 		}

lib/mapObjects/IObjectInterface.h

@@ -23,7 +23,7 @@ class RNG;
 }
 
 struct BattleResult;
-struct UpgradeInfo;
+class UpgradeInfo;
 class BoatId;
 class CGObjectInstance;
 class CStackInstance;

lib/mapObjects/MiscObjects.cpp

@@ -32,6 +32,7 @@
 #include "../networkPacks/PacksForClient.h"
 #include "../networkPacks/PacksForClientBattle.h"
 #include "../networkPacks/StackLocation.h"
+#include "../lib/gameState/UpgradeInfo.h"
 
 #include <vstd/RNG.h>
 
@@ -1323,13 +1324,9 @@ void HillFort::fillUpgradeInfo(UpgradeInfo & info, const CStackInstance &stack)
 
 	int costModifier = upgradeCostPercentage[index];
 
-	if (costModifier < 0)
-		return; // upgrade not allowed
-
 	for(const auto & nid : stack.getCreature()->upgrades)
 	{
-		info.newID.push_back(nid);
-		info.cost.push_back((nid.toCreature()->getFullRecruitCost() - stack.getType()->getFullRecruitCost()) * costModifier / 100);
+		info.addUpgrade(nid, stack.getType(), costModifier);
 	}
 }
 

lib/mapObjects/MiscObjects.h

@@ -16,6 +16,7 @@
 VCMI_LIB_NAMESPACE_BEGIN
 
 class CMap;
+class UpgradeInfo;
 
 // This one teleport-specific, but has to be available everywhere in callbacks and netpacks
 // For now it's will be there till teleports code refactored and moved into own file

server/CGameHandler.cpp

@@ -53,6 +53,7 @@
 #include "../lib/filesystem/Filesystem.h"
 
 #include "../lib/gameState/CGameState.h"
+#include "../lib/gameState/UpgradeInfo.h"
 
 #include "../lib/mapping/CMap.h"
 #include "../lib/mapping/CMapService.h"
@@ -2414,19 +2415,18 @@ bool CGameHandler::upgradeCreature(ObjectInstanceID objid, SlotID pos, CreatureI
 	{
 		COMPLAIN_RET("Cannot upgrade, no stack at slot " + std::to_string(pos));
 	}
-	UpgradeInfo ui;
-	fillUpgradeInfo(obj, pos, ui);
+	UpgradeInfo upgradeInfo(obj->getStackPtr(pos)->getId());
+	fillUpgradeInfo(obj, pos, upgradeInfo);
 	PlayerColor player = obj->tempOwner;
 	const PlayerState *p = getPlayerState(player);
 	int crQuantity = obj->stacks.at(pos)->count;
-	int newIDpos= vstd::find_pos(ui.newID, upgID);//get position of new id in UpgradeInfo
 
 	//check if upgrade is possible
-	if ((ui.oldID == CreatureID::NONE || newIDpos == -1) && complain("That upgrade is not possible!"))
+	if (!upgradeInfo.hasUpgrades() && complain("That upgrade is not possible!"))
 	{
 		return false;
 	}
-	TResources totalCost = ui.cost.at(newIDpos) * crQuantity;
+	TResources totalCost = upgradeInfo.getUpgradeCostsFor(upgID) * crQuantity;
 
 	//check if player has enough resources
 	if (!p->resources.canAfford(totalCost))