Roll first for hero class, and then - for actual hero for tavern

lib/gameState/TavernHeroesPool.cpp

@@ -16,7 +16,7 @@ VCMI_LIB_NAMESPACE_BEGIN
 
 TavernHeroesPool::~TavernHeroesPool()
 {
-	for(auto ptr : heroesPool) // clean hero pool
+	for(const auto & ptr : heroesPool) // clean hero pool
 		delete ptr.second;
 }
 

server/HeroPoolProcessor.cpp

@@ -21,10 +21,13 @@
 #include "../lib/gameState/CGameState.h"
 #include "../lib/gameState/TavernHeroesPool.h"
 
-HeroPoolProcessor::HeroPoolProcessor() = default;
+HeroPoolProcessor::HeroPoolProcessor()
+	: gameHandler(nullptr)
+{
+}
 
-HeroPoolProcessor::HeroPoolProcessor(CGameHandler * gameHandler):
-	gameHandler(gameHandler)
+HeroPoolProcessor::HeroPoolProcessor(CGameHandler * gameHandler)
+	: gameHandler(gameHandler)
 {
 }
 
@@ -172,69 +175,115 @@ bool HeroPoolProcessor::hireHero(const CGObjectInstance *obj, const HeroTypeID &
 	return true;
 }
 
-CGHeroInstance * HeroPoolProcessor::pickHeroFor(bool isNative,
-													 const PlayerColor & player,
-													 const FactionID & factionID,
-													 CRandomGenerator & rand,
-													 const CHeroClass * bannedClass) const
+std::set<const CHeroClass *> HeroPoolProcessor::findAvailableClassesFor(const PlayerColor & player) const
 {
-	if(player >= PlayerColor::PLAYER_LIMIT)
-	{
-		logGlobal->error("Cannot pick hero for player %d. Wrong owner!", player.getStr());
-		return nullptr;
-	}
+	std::set<const CHeroClass *> result;
 
 	const auto & hpool = gameHandler->gameState()->hpool;
+	FactionID factionID = gameHandler->getPlayerSettings(player)->castle;
 
-	if(isNative)
+	for(auto & elem : hpool->unusedHeroesFromPool())
 	{
-		std::vector<CGHeroInstance *> pool;
-
-		for(auto & elem : hpool->unusedHeroesFromPool())
-		{
-			//get all available heroes
-			bool heroAvailable = hpool->isHeroAvailableFor(elem.first, player);
-			bool heroClassNative = elem.second->type->heroClass->faction == factionID;
+		bool heroAvailable = hpool->isHeroAvailableFor(elem.first, player);
+		bool heroClassBanned = elem.second->type->heroClass->selectionProbability[factionID] == 0;
 
-			if(heroAvailable && heroClassNative)
-				pool.push_back(elem.second);
-		}
+		if(heroAvailable && !heroClassBanned)
+			result.insert(elem.second->type->heroClass);
+	}
 
-		if(!pool.empty())
-			return *RandomGeneratorUtil::nextItem(pool, rand);
+	return result;
+}
 
-		logGlobal->error("Cannot pick native hero for %s. Picking any...", player.getStr());
-	}
+std::set<CGHeroInstance *> HeroPoolProcessor::findAvailableHeroesFor(const PlayerColor & player, const CHeroClass * heroClass) const
+{
+	std::set<CGHeroInstance *> result;
 
-	std::vector<CGHeroInstance *> pool;
-	int totalWeight = 0;
+	const auto & hpool = gameHandler->gameState()->hpool;
 
 	for(auto & elem : hpool->unusedHeroesFromPool())
 	{
 		bool heroAvailable = hpool->isHeroAvailableFor(elem.first, player);
-		bool heroClassBanned = bannedClass && elem.second->type->heroClass == bannedClass;
+		bool heroClassMatches = elem.second->type->heroClass == heroClass;
 
-		if(heroAvailable && !heroClassBanned)
-		{
-			pool.push_back(elem.second);
-			totalWeight += elem.second->type->heroClass->selectionProbability[factionID]; //total weight
-		}
+		if(heroAvailable && heroClassMatches)
+			result.insert(elem.second);
 	}
-	if(pool.empty() || totalWeight == 0)
+
+	return result;
+}
+
+const CHeroClass * HeroPoolProcessor::pickClassFor(bool isNative, const PlayerColor & player, const FactionID & factionID, CRandomGenerator & rand) const
+{
+	if(player >= PlayerColor::PLAYER_LIMIT)
+	{
+		logGlobal->error("Cannot pick hero for player %d. Wrong owner!", player.getStr());
+		return nullptr;
+	}
+
+	const auto & hpool = gameHandler->gameState()->hpool;
+	const auto & currentTavern = hpool->getHeroesFor(player);
+
+	std::set<const CHeroClass *> potentialClasses = findAvailableClassesFor(player);
+	std::set<const CHeroClass *> possibleClasses;
+
+	if(potentialClasses.empty())
 	{
 		logGlobal->error("There are no heroes available for player %s!", player.getStr());
 		return nullptr;
 	}
 
+	for(const auto & heroClass : potentialClasses)
+	{
+		if (isNative && heroClass->faction != factionID)
+			continue;
+
+		bool hasSameClass = vstd::contains_if(currentTavern, [&](const CGHeroInstance * hero){
+			return hero->type->heroClass == heroClass;
+		});
+
+		if (hasSameClass)
+			continue;
+
+		possibleClasses.insert(heroClass);
+	}
+
+	if (possibleClasses.empty())
+	{
+		logGlobal->error("Cannot pick native hero for %s. Picking any...", player.getStr());
+		possibleClasses = potentialClasses;
+	}
+
+	int totalWeight = 0;
+	for(const auto & heroClass : possibleClasses)
+		totalWeight += heroClass->selectionProbability.at(factionID);
+
 	int roll = rand.nextInt(totalWeight - 1);
-	for(auto & elem : pool)
+	for(const auto & heroClass : possibleClasses)
 	{
-		roll -= elem->type->heroClass->selectionProbability[factionID];
+		roll -= heroClass->selectionProbability.at(factionID);
 		if(roll < 0)
-		{
-			return elem;
-		}
+			return heroClass;
 	}
 
-	return pool.back();
+	return *possibleClasses.rbegin();
+}
+
+CGHeroInstance * HeroPoolProcessor::pickHeroFor(bool isNative,
+													 const PlayerColor & player,
+													 const FactionID & factionID,
+													 CRandomGenerator & rand,
+													 const CHeroClass * bannedClass) const
+{
+	const CHeroClass * heroClass = pickClassFor(isNative, player, factionID, rand);
+
+	if(!heroClass)
+		return nullptr;
+
+	std::set<CGHeroInstance *> possibleHeroes = findAvailableHeroesFor(player, heroClass);
+
+	assert(!possibleHeroes.empty());
+	if(possibleHeroes.empty())
+		return nullptr;
+
+	return *RandomGeneratorUtil::nextItem(possibleHeroes, rand);
 }

server/HeroPoolProcessor.h

@@ -31,6 +31,11 @@ class HeroPoolProcessor : boost::noncopyable
 	void clearHeroFromSlot(const PlayerColor & color, TavernHeroSlot slot);
 	void selectNewHeroForSlot(const PlayerColor & color, TavernHeroSlot slot, bool needNativeHero, bool giveStartingArmy);
 
+	std::set<const CHeroClass *> findAvailableClassesFor(const PlayerColor & player) const;
+	std::set<CGHeroInstance *> findAvailableHeroesFor(const PlayerColor & player, const CHeroClass * heroClass) const;
+
+	const CHeroClass * pickClassFor(bool isNative, const PlayerColor & player, const FactionID & faction, CRandomGenerator & rand) const;
+
 	CGHeroInstance * pickHeroFor(bool isNative, const PlayerColor & player, const FactionID & faction, CRandomGenerator & rand, const CHeroClass * bannedClass) const;
 public:
 	HeroPoolProcessor();