- Unban hero if related Prison is destroyed during map generation

- Move prison counter to PrisonHeroPlacer

lib/mapping/CMap.cpp

@@ -650,10 +650,17 @@ void CMap::banWaterHeroes()
 void CMap::banHero(const HeroTypeID & id)
 {
 	if (!vstd::contains(allowedHeroes, id))
-		logGlobal->warn("Attempt to ban hero %d, who is already not allowed", id.encode(id));
+		logGlobal->warn("Attempt to ban hero %s, who is already not allowed", id.encode(id));
 	allowedHeroes.erase(id);
 }
 
+void CMap::unbanHero(const HeroTypeID & id)
+{
+	if (vstd::contains(allowedHeroes, id))
+		logGlobal->warn("Attempt to unban hero %s, who is already allowed", id.encode(id));
+	allowedHeroes.insert(id);
+}
+
 void CMap::initTerrain()
 {
 	terrain.resize(boost::extents[levels()][width][height]);

lib/mapping/CMap.h

@@ -112,6 +112,7 @@ public:
 	void banWaterArtifacts();
 	void banWaterHeroes();
 	void banHero(const HeroTypeID& id);
+	void unbanHero(const HeroTypeID & id);
 	void banWaterSpells();
 	void banWaterSkills();
 	void banWaterContent();

lib/rmg/CMapGenerator.cpp

@@ -35,7 +35,7 @@ VCMI_LIB_NAMESPACE_BEGIN
 
 CMapGenerator::CMapGenerator(CMapGenOptions& mapGenOptions, int RandomSeed) :
 	mapGenOptions(mapGenOptions), randomSeed(RandomSeed),
-	allowedPrisons(0), monolithIndex(0)
+	monolithIndex(0)
 {
 	loadConfig();
 	rand.setSeed(this->randomSeed);
@@ -96,12 +96,6 @@ const CMapGenOptions& CMapGenerator::getMapGenOptions() const
 	return mapGenOptions;
 }
 
-void CMapGenerator::initPrisonsRemaining()
-{
-	allowedPrisons = map->getMap(this).allowedHeroes.size();
-	allowedPrisons = std::max<int> (0, allowedPrisons - 16 * mapGenOptions.getHumanOrCpuPlayerCount()); //so at least 16 heroes will be available for every player
-}
-
 void CMapGenerator::initQuestArtsRemaining()
 {
 	//TODO: Move to QuestArtifactPlacer?
@@ -122,7 +116,6 @@ std::unique_ptr<CMap> CMapGenerator::generate()
 		addHeaderInfo();
 		map->initTiles(*this, rand);
 		Load::Progress::step();
-		initPrisonsRemaining();
 		initQuestArtsRemaining();
 		genZones();
 		Load::Progress::step();
@@ -468,11 +461,6 @@ int CMapGenerator::getNextMonlithIndex()
 	}
 }
 
-int CMapGenerator::getPrisonsRemaning() const
-{
-	return allowedPrisons;
-}
-
 std::shared_ptr<CZonePlacer> CMapGenerator::getZonePlacer() const
 {
 	return placer;
@@ -514,6 +502,12 @@ void CMapGenerator::banHero(const HeroTypeID & id)
 	map->getMap(this).banHero(id);
 }
 
+void CMapGenerator::unbanHero(const HeroTypeID & id)
+{
+	map->getMap(this).unbanHero(id);
+}
+
+
 Zone * CMapGenerator::getZoneWater() const
 {
 	for(auto & z : map->getZones())

lib/rmg/CMapGenerator.h

@@ -66,6 +66,7 @@ public:
 	const std::vector<HeroTypeID> getAllPossibleHeroes() const;
 	void banQuestArt(const ArtifactID & id);
 	void banHero(const HeroTypeID& id);
+	void unbanHero(const HeroTypeID & id);
 
 	Zone * getZoneWater() const;
 	void addWaterTreasuresInfo();
@@ -82,7 +83,6 @@ private:
 	
 	std::vector<rmg::ZoneConnection> connectionsLeft;
 	
-	int allowedPrisons;
 	int monolithIndex;
 	std::vector<ArtifactID> questArtifacts;
 

lib/rmg/modificators/ObjectDistributor.cpp

@@ -15,6 +15,7 @@
 #include "../RmgMap.h"
 #include "../CMapGenerator.h"
 #include "TreasurePlacer.h"
+#include "PrisonHeroPlacer.h"
 #include "QuestArtifactPlacer.h"
 #include "TownPlacer.h"
 #include "TerrainPainter.h"
@@ -145,7 +146,18 @@ void ObjectDistributor::distributePrisons()
 
 	RandomGeneratorUtil::randomShuffle(zones, zone.getRand());
 
-	size_t allowedPrisons = generator.getPrisonsRemaning();
+	// TODO: Some shorthand for unique Modificator
+	PrisonHeroPlacer * prisonHeroPlacer = nullptr;
+	for(auto & z : map.getZones())
+	{
+		prisonHeroPlacer = z.second->getModificator<PrisonHeroPlacer>();
+		if (prisonHeroPlacer)
+		{
+			break;
+		}
+	}
+
+	size_t allowedPrisons = prisonHeroPlacer->getPrisonsRemaning();
 	for (int i = zones.size() - 1; i >= 0; i--)
 	{
 		auto zone = zones[i].second;

lib/rmg/modificators/PrisonHeroPlacer.cpp

@@ -28,17 +28,30 @@ void PrisonHeroPlacer::process()
 
 void PrisonHeroPlacer::init()
 {
+	// Reserve at least 16 heroes for each player
+	reservedHeroes = 16 * generator.getMapGenOptions().getHumanOrCpuPlayerCount();
 }
 
 void PrisonHeroPlacer::getAllowedHeroes()
 {
-    allowedHeroes = generator.getAllPossibleHeroes();
+	// TODO: Give each zone unique HeroPlacer with private hero list?
+
+	// Call that only once
+	if (allowedHeroes.empty())
+	{
+    	allowedHeroes = generator.getAllPossibleHeroes();
+	}
+}
+
+int PrisonHeroPlacer::getPrisonsRemaning() const
+{
+	return std::max<int>(allowedHeroes.size() - reservedHeroes, 0);
 }
 
 HeroTypeID PrisonHeroPlacer::drawRandomHero()
 {
 	RecursiveLock lock(externalAccessMutex);
-	if (!allowedHeroes.empty())
+	if (getPrisonsRemaning() > 0)
 	{
 		RandomGeneratorUtil::randomShuffle(allowedHeroes, zone.getRand());
         HeroTypeID ret = allowedHeroes.back();
@@ -53,4 +66,10 @@ HeroTypeID PrisonHeroPlacer::drawRandomHero()
 	}
 }
 
+void PrisonHeroPlacer::unbanHero(const HeroTypeID & hid)
+{
+	RecursiveLock lock(externalAccessMutex);
+	generator.unbanHero(hid);
+}
+
 VCMI_LIB_NAMESPACE_END

lib/rmg/modificators/PrisonHeroPlacer.h

@@ -25,16 +25,17 @@ public:
 	void process() override;
 	void init() override;
 
+	int getPrisonsRemaning() const;
 	HeroTypeID drawRandomHero();
+	void unbanHero(const HeroTypeID & hid);
 
 private:
     void getAllowedHeroes();
+	size_t reservedHeroes;
 
 protected:
 
     std::vector<HeroTypeID> allowedHeroes;
-
-	// TODO: Count allowed heroes?
 };
 
 VCMI_LIB_NAMESPACE_END

lib/rmg/modificators/TreasurePlacer.cpp

@@ -33,6 +33,12 @@
 
 VCMI_LIB_NAMESPACE_BEGIN
 
+ObjectInfo::ObjectInfo():
+	destroyObject([](){})
+{
+
+}
+
 void TreasurePlacer::process()
 {
 	addAllPossibleObjects();
@@ -109,17 +115,23 @@ void TreasurePlacer::addAllPossibleObjects()
 		size_t prisonsLeft = getMaxPrisons();
 		for (int i = prisonsLevels - 1; i >= 0; i--)
 		{
+			ObjectInfo oi; // Create new instance which will hold destructor operation
+
 			oi.value = generator.getConfig().prisonValues[i];
 			if (oi.value > zone.getMaxTreasureValue())
 			{
 				continue;
 			}
 
-			oi.generateObject = [i, this, prisonHeroPlacer]() -> CGObjectInstance*
+			oi.generateObject = [i, this, prisonHeroPlacer, &oi]() -> CGObjectInstance*
 			{
 				auto possibleHeroes = generator.getAllPossibleHeroes();
 
 				HeroTypeID hid = prisonHeroPlacer->drawRandomHero();
+				oi.destroyObject = [hid, prisonHeroPlacer]()
+				{
+					prisonHeroPlacer->unbanHero(hid);
+				};
 
 				auto factory = VLC->objtypeh->getHandlerFor(Obj::PRISON, 0);
 				auto* obj = dynamic_cast<CGHeroInstance*>(factory->create());
@@ -654,9 +666,8 @@ rmg::Object TreasurePlacer::constructTreasurePile(const std::vector<ObjectInfo*>
 		if(oi->templates.empty())
 		{
 			logGlobal->warn("Deleting randomized object with no templates: %s", object->getObjectName());
-			// Possible memory leak, but this is a weird case in first place
+			oi->destroyObject();
 			delete object;
-			// FIXME: We also lose randomized hero or quest artifact
 			continue;
 		}
 		
@@ -816,6 +827,7 @@ void TreasurePlacer::createTreasures(ObjectManager& manager)
 	{
 		for (auto* oi : treasurePile)
 		{
+			oi->destroyObject();
 			oi->maxPerZone++;
 		}
 	};

lib/rmg/modificators/TreasurePlacer.h

@@ -22,12 +22,15 @@ class CRandomGenerator;
 
 struct ObjectInfo
 {
+	ObjectInfo::ObjectInfo();
+
 	std::vector<std::shared_ptr<const ObjectTemplate>> templates;
 	ui32 value = 0;
 	ui16 probability = 0;
 	ui32 maxPerZone = 1;
 	//ui32 maxPerMap; //unused
 	std::function<CGObjectInstance *()> generateObject;
+	std::function<void()> destroyObject;
 	
 	void setTemplates(MapObjectID type, MapObjectSubID subtype, TerrainId terrain);
 };