From 92c4bcbf68f853552cddb6085ad66584067d9876 Mon Sep 17 00:00:00 2001
From: Ivan Savenko <saven.ivan@gmail.com>
Date: Wed, 7 May 2025 15:55:48 +0300
Subject: [PATCH] Unify checks for creature limiter

---
 lib/CCreatureSet.cpp         | 32 ++++++++++++++++++++++++++++++++
 lib/CCreatureSet.h           |  5 +++++
 lib/mapObjects/CQuest.cpp    | 24 +-----------------------
 lib/mapping/MapFormatH3M.cpp |  1 +
 lib/rewardable/Limiter.cpp   | 25 ++-----------------------
 5 files changed, 41 insertions(+), 46 deletions(-)

diff --git a/lib/CCreatureSet.cpp b/lib/CCreatureSet.cpp
index c393fdcaf..11f128e79 100644
--- a/lib/CCreatureSet.cpp
+++ b/lib/CCreatureSet.cpp
@@ -612,6 +612,38 @@ bool CCreatureSet::canBeMergedWith(const CCreatureSet &cs, bool allowMergingStac
 	}
 }
 
+bool CCreatureSet::hasUnits(const std::vector<CStackBasicDescriptor> & units, bool requireLastStack) const
+{
+	bool foundExtraCreatures = false;
+	int testedSlots = 0;
+	for(const auto & reqStack : units)
+	{
+		size_t count = 0;
+		for(const auto & slot : Slots())
+		{
+			const auto & heroStack = slot.second;
+			if (heroStack->getType() == reqStack.getType())
+			{
+				count += heroStack->getCount();
+				testedSlots += 1;
+			}
+		}
+		if (count > reqStack.getCount())
+			foundExtraCreatures = true;
+
+		if (count < reqStack.getCount()) //not enough creatures of this kind
+			return false;
+	}
+
+	if (requireLastStack)
+	{
+		if (!foundExtraCreatures && testedSlots >= Slots().size())
+			return false;
+	}
+
+	return true;
+}
+
 bool CCreatureSet::hasStackAtSlot(const SlotID & slot) const
 {
 	return vstd::contains(stacks, slot);
diff --git a/lib/CCreatureSet.h b/lib/CCreatureSet.h
index 893660df2..f3bdd08c2 100644
--- a/lib/CCreatureSet.h
+++ b/lib/CCreatureSet.h
@@ -333,6 +333,11 @@ public:
 	bool contains(const CStackInstance *stack) const;
 	bool canBeMergedWith(const CCreatureSet &cs, bool allowMergingStacks = true) const;
 
+	/// Returns true if this creature set contains all listed units
+	/// If requireLastStack is true, then this function will also
+	/// require presence of any unit other than requested (or more units than requested)
+	bool hasUnits(const std::vector<CStackBasicDescriptor> & units, bool requireLastStack = true) const;
+
 	template <typename Handler> void serialize(Handler &h)
 	{
 		h & stacks;
diff --git a/lib/mapObjects/CQuest.cpp b/lib/mapObjects/CQuest.cpp
index e5039dae2..fdb956879 100644
--- a/lib/mapObjects/CQuest.cpp
+++ b/lib/mapObjects/CQuest.cpp
@@ -101,29 +101,7 @@ const std::string & CQuest::missionState(int state)
 
 bool CQuest::checkMissionArmy(const CQuest * q, const CCreatureSet * army)
 {
-	std::vector<CStackBasicDescriptor>::const_iterator cre;
-	TSlots::const_iterator it;
-	ui32 count = 0;
-	ui32 slotsCount = 0;
-	bool hasExtraCreatures = false;
-	for(cre = q->mission.creatures.begin(); cre != q->mission.creatures.end(); ++cre)
-	{
-		for(count = 0, it = army->Slots().begin(); it != army->Slots().end(); ++it)
-		{
-			if(it->second->getType() == cre->getType())
-			{
-				count += it->second->getCount();
-				slotsCount++;
-			}
-		}
-
-		if(static_cast<TQuantity>(count) < cre->getCount()) //not enough creatures of this kind
-			return false;
-
-		hasExtraCreatures |= static_cast<TQuantity>(count) > cre->getCount();
-	}
-
-	return hasExtraCreatures || slotsCount < army->Slots().size();
+	return army->hasUnits(q->mission.creatures, true);
 }
 
 bool CQuest::checkQuest(const CGHeroInstance * h) const
diff --git a/lib/mapping/MapFormatH3M.cpp b/lib/mapping/MapFormatH3M.cpp
index b658582c9..ec44d54e9 100644
--- a/lib/mapping/MapFormatH3M.cpp
+++ b/lib/mapping/MapFormatH3M.cpp
@@ -2429,6 +2429,7 @@ EQuestMission CMapLoaderH3M::readQuest(IQuestObject * guard, const int3 & positi
 		{
 			size_t typeNumber = reader->readUInt8();
 			guard->getQuest().mission.creatures.resize(typeNumber);
+			guard->getQuest().mission.hasExtraCreatures = true;
 			for(size_t hh = 0; hh < typeNumber; ++hh)
 			{
 				guard->getQuest().mission.creatures[hh].setType(reader->readCreature().toCreature());
diff --git a/lib/rewardable/Limiter.cpp b/lib/rewardable/Limiter.cpp
index e12a1d2c0..429f2ced1 100644
--- a/lib/rewardable/Limiter.cpp
+++ b/lib/rewardable/Limiter.cpp
@@ -86,30 +86,9 @@ bool Rewardable::Limiter::heroAllowed(const CGHeroInstance * hero) const
 			return false;
 	}
 
-	bool foundExtraCreatures = false;
-	int testedSlots = 0;
-	for(const auto & reqStack : creatures)
+	if (!creatures.empty())
 	{
-		size_t count = 0;
-		for(const auto & slot : hero->Slots())
-		{
-			const auto & heroStack = slot.second;
-			if (heroStack->getType() == reqStack.getType())
-			{
-				count += heroStack->getCount();
-				testedSlots += 1;
-			}
-		}
-		if (count > reqStack.getCount())
-			foundExtraCreatures = true;
-
-		if (count < reqStack.getCount()) //not enough creatures of this kind
-			return false;
-	}
-
-	if (hasExtraCreatures)
-	{
-		if (!foundExtraCreatures && testedSlots >= hero->Slots().size())
+		if (!hero->hasUnits(creatures, hasExtraCreatures))
 			return false;
 	}