- Pre-distribute max number of Prisons

- Minor refactor for artifact and hero pool management

lib/rmg/CMapGenerator.cpp

@@ -31,7 +31,7 @@ VCMI_LIB_NAMESPACE_BEGIN
 
 CMapGenerator::CMapGenerator(CMapGenOptions& mapGenOptions, int RandomSeed) :
 	mapGenOptions(mapGenOptions), randomSeed(RandomSeed),
-	prisonsRemaining(0), monolithIndex(0)
+	allowedPrisons(0), monolithIndex(0)
 {
 	loadConfig();
 	rand.setSeed(this->randomSeed);
@@ -99,13 +99,13 @@ const CMapGenOptions& CMapGenerator::getMapGenOptions() const
 
 void CMapGenerator::initPrisonsRemaining()
 {
-	prisonsRemaining = 0;
+	allowedPrisons = 0;
 	for (auto isAllowed : map->map().allowedHeroes)
 	{
 		if (isAllowed)
-			prisonsRemaining++;
+			allowedPrisons++;
 	}
-	prisonsRemaining = std::max<int> (0, prisonsRemaining - 16 * mapGenOptions.getPlayerCount()); //so at least 16 heroes will be available for every player
+	allowedPrisons = std::max<int> (0, allowedPrisons - 16 * mapGenOptions.getPlayerCount()); //so at least 16 heroes will be available for every player
 }
 
 void CMapGenerator::initQuestArtsRemaining()
@@ -368,7 +368,7 @@ int CMapGenerator::getNextMonlithIndex()
 
 int CMapGenerator::getPrisonsRemaning() const
 {
-	return prisonsRemaining;
+	return allowedPrisons;
 }
 
 std::shared_ptr<CZonePlacer> CMapGenerator::getZonePlacer() const
@@ -376,22 +376,34 @@ std::shared_ptr<CZonePlacer> CMapGenerator::getZonePlacer() const
 	return placer;
 }
 
-void CMapGenerator::decreasePrisonsRemaining()
+const std::vector<ArtifactID> & CMapGenerator::getAllPossibleQuestArtifacts() const
 {
-	prisonsRemaining = std::max (0, prisonsRemaining - 1);
+	return questArtifacts;
 }
 
-const std::vector<ArtifactID> & CMapGenerator::getQuestArtsRemaning() const
+const std::vector<HeroTypeID>& CMapGenerator::getAllPossibleHeroes() const
 {
-	return questArtifacts;
+	//Skip heroes that were banned, including the ones placed in prisons
+	std::vector<HeroTypeID> ret;
+	for (int j = 0; j < map->map().allowedHeroes.size(); j++)
+	{
+		if (map->map().allowedHeroes[j])
+			ret.push_back(HeroTypeID(j));
+	}
+	return ret;
 }
 
 void CMapGenerator::banQuestArt(const ArtifactID & id)
 {
 	map->map().allowedArtifact[id] = false;
-	vstd::erase_if_present(questArtifacts, id);
 }
 
+void CMapGenerator::banHero(const HeroTypeID & id)
+{
+	map->map().allowedHeroes[id] = false;
+}
+
+
 Zone * CMapGenerator::getZoneWater() const
 {
 	for(auto & z : map->getZones())

lib/rmg/CMapGenerator.h

@@ -63,9 +63,10 @@ public:
 	int getNextMonlithIndex();
 	int getPrisonsRemaning() const;
 	std::shared_ptr<CZonePlacer> getZonePlacer() const;
-	void decreasePrisonsRemaining();
-	const std::vector<ArtifactID> & getQuestArtsRemaning() const;
+	const std::vector<ArtifactID> & getAllPossibleQuestArtifacts() const;
+	const std::vector<HeroTypeID>& getAllPossibleHeroes() const;
 	void banQuestArt(const ArtifactID & id);
+	void banHero(const HeroTypeID& id);
 
 	Zone * getZoneWater() const;
 	void addWaterTreasuresInfo();
@@ -83,8 +84,7 @@ private:
 	
 	//std::pair<Zones::key_type, Zones::mapped_type> zoneWater;
 
-	int prisonsRemaining;
-	//int questArtsRemaining;
+	int allowedPrisons;
 	int monolithIndex;
 	std::vector<ArtifactID> questArtifacts; //TODO: Protect with mutex
 

lib/rmg/ObjectDistributor.cpp

@@ -130,7 +130,7 @@ void ObjectDistributor::distributeSeerHuts()
 
 	RandomGeneratorUtil::randomShuffle(zones, generator.rand);
 
-	const auto & possibleQuestArts = generator.getQuestArtsRemaning();
+	const auto & possibleQuestArts = generator.getAllPossibleQuestArtifacts();
 	size_t availableArts = possibleQuestArts.size();
 	auto artIt = possibleQuestArts.begin();
 	for (int i = zones.size() - 1; i >= 0 ; i--)
@@ -149,4 +149,25 @@ void ObjectDistributor::distributeSeerHuts()
 	}
 }
 
+void ObjectDistributor::distributePrisons()
+{
+	//Copy by value to random shuffle
+	const auto & zoneMap = map.getZones();
+	RmgMap::ZoneVector zones(zoneMap.begin(), zoneMap.end());
+
+	RandomGeneratorUtil::randomShuffle(zones, generator.rand);
+
+	size_t allowedPrisons = generator.getPrisonsRemaning();
+	for (int i = zones.size() - 1; i >= 0; i--)
+	{
+		auto zone = zones[i].second;
+		auto * tp = zone->getModificator<TreasurePlacer>();
+		if (tp)
+		{
+			tp->setMaxPrisons(std::ceil(float(allowedPrisons) / (i + 1)));
+			allowedPrisons -= tp->getMaxPrisons();
+		}
+	}
+}
+
 VCMI_LIB_NAMESPACE_END

lib/rmg/ObjectDistributor.h

@@ -22,6 +22,7 @@ class ObjectDistributor : public Modificator
 {
 	void distributeLimitedObjects();
 	void distributeSeerHuts();
+	void distributePrisons();
 
 public:
 	MODIFICATOR(ObjectDistributor);

lib/rmg/QuestArtifactPlacer.cpp

@@ -65,7 +65,6 @@ void QuestArtifactPlacer::findZonesForQuestArts()
 	}
 
 	logGlobal->info("Number of nearby zones suitable for quest artifacts: %d", questArtZones.size());
-	logGlobal->info("Number of possible quest artifacts remaining: %d", generator.getQuestArtsRemaning().size());
 }
 
 void QuestArtifactPlacer::placeQuestArtifacts(CRandomGenerator * rand)

lib/rmg/TreasurePlacer.cpp

@@ -96,26 +96,28 @@ void TreasurePlacer::addAllPossibleObjects()
 	//prisons
 	//levels 1, 5, 10, 20, 30
 	static int prisonsLevels = std::min(generator.getConfig().prisonExperience.size(), generator.getConfig().prisonValues.size());
-	for(int i = 0; i < prisonsLevels; i++)
+	
+	size_t prisonsLeft = getMaxPrisons();
+	for(int i = prisonsLevels - 1; i >= 0 ;i--)
 	{
+		oi.value = generator.getConfig().prisonValues[i];
+		if (oi.value > zone.getMaxTreasureValue())
+		{
+			continue;
+		}
+
 		oi.generateObject = [i, this]() -> CGObjectInstance *
 		{
-			std::vector<ui32> possibleHeroes;
-			for(int j = 0; j < map.map().allowedHeroes.size(); j++)
-			{
-				if(map.map().allowedHeroes[j])
-					possibleHeroes.push_back(j);
-			}
-			
-			auto hid = *RandomGeneratorUtil::nextItem(possibleHeroes, generator.rand);
+			auto possibleHeroes = generator.getAllPossibleHeroes();
+			HeroTypeID hid = *RandomGeneratorUtil::nextItem(possibleHeroes, generator.rand);
+
 			auto factory = VLC->objtypeh->getHandlerFor(Obj::PRISON, 0);
 			auto * obj = dynamic_cast<CGHeroInstance *>(factory->create());
 
 			obj->subID = hid; //will be initialized later
 			obj->exp = generator.getConfig().prisonExperience[i];
 			obj->setOwner(PlayerColor::NEUTRAL);
-			map.map().allowedHeroes[hid] = false; //ban this hero
-			generator.decreasePrisonsRemaining();
+			generator.banHero(hid);
 			obj->appearance = VLC->objtypeh->getHandlerFor(Obj::PRISON, 0)->getTemplates(zone.getTerrainType()).front(); //can't init template with hero subID
 			
 			return obj;
@@ -123,7 +125,10 @@ void TreasurePlacer::addAllPossibleObjects()
 		oi.setTemplate(Obj::PRISON, 0, zone.getTerrainType());
 		oi.value = generator.getConfig().prisonValues[i];
 		oi.probability = 30;
-		oi.maxPerZone = generator.getPrisonsRemaning() / 5; //probably not perfect, but we can't generate more prisons than hereos.
+		
+		//Distribute all allowed prisons, starting from the most valuable
+		oi.maxPerZone = (std::ceil((float)prisonsLeft / (i + 1)));
+		prisonsLeft -= oi.maxPerZone;
 		addObjectToRandomPool(oi);
 	}
 	
@@ -523,6 +528,16 @@ size_t TreasurePlacer::getPossibleObjectsSize() const
 	return possibleObjects.size();
 }
 
+void TreasurePlacer::setMaxPrisons(size_t count)
+{
+	maxPrisons = count;
+}
+
+size_t TreasurePlacer::getMaxPrisons() const
+{
+	return maxPrisons;
+}
+
 bool TreasurePlacer::isGuardNeededForTreasure(int value)
 {
 	return zone.getType() != ETemplateZoneType::WATER && value > minGuardedValue;

lib/rmg/TreasurePlacer.h

@@ -48,6 +48,8 @@ public:
 	void addAllPossibleObjects(); //add objects, including zone-specific, to possibleObjects
 
 	size_t getPossibleObjectsSize() const;
+	void setMaxPrisons(size_t count);
+	size_t getMaxPrisons() const;
 	
 protected:
 	bool isGuardNeededForTreasure(int value);
@@ -63,6 +65,8 @@ protected:
 	rmg::Area treasureArea;
 	rmg::Area treasureBlockArea;
 	rmg::Area guards;
+
+	size_t maxPrisons;
 };
 
 VCMI_LIB_NAMESPACE_END