Fix creature type limiter not considering upgrades of upgrades

lib/CCreatureHandler.cpp

@@ -337,12 +337,40 @@ void CCreature::addBonus(int val, BonusType type, BonusSubtypeID subtype)
 	}
 }
 
-bool CCreature::isMyUpgrade(const CCreature *anotherCre) const
+bool CCreature::isMyDirectUpgrade(const CCreature *anotherCre) const
 {
-	//TODO upgrade of upgrade?
 	return vstd::contains(upgrades, anotherCre->getId());
 }
 
+bool CCreature::isMyDirectOrIndirectUpgrade(const CCreature *anotherCre) const
+{
+	std::set<CreatureID> foundUpgrades;
+	std::vector<CreatureID> upgradesToTest;
+
+	upgradesToTest.push_back(getId());
+
+	while (!upgradesToTest.empty())
+	{
+		CreatureID testedID = upgradesToTest.back();
+		const CCreature * testedPtr = testedID.toCreature();
+
+		upgradesToTest.pop_back();
+
+		for (const auto & upgrade : testedPtr->upgrades)
+		{
+			if (upgrade == anotherCre->getId())
+				return true;
+
+			if (foundUpgrades.count(upgrade))
+				continue;
+
+			upgradesToTest.push_back(upgrade);
+			foundUpgrades.insert(upgrade);
+		}
+	}
+	return false;
+}
+
 std::string CCreature::nodeName() const
 {
 	return "\"" + getNamePluralTextID() + "\"";

lib/CCreatureHandler.h

@@ -163,7 +163,13 @@ public:
 	static CCreature::CreatureQuantityId getQuantityID(const int & quantity);
 	static std::string getQuantityRangeStringForId(const CCreature::CreatureQuantityId & quantityId);
 	static int estimateCreatureCount(ui32 countID); //reverse version of above function, returns middle of range
-	bool isMyUpgrade(const CCreature *anotherCre) const;
+
+	/// Returns true if this creature can be directly upgraded to target
+	bool isMyDirectUpgrade(const CCreature * target) const;
+
+	/// Returns true if this creature can be upgraded to target
+	/// Performs full search through potential upgrades of upgrades
+	bool isMyDirectOrIndirectUpgrade(const CCreature *target) const;
 
 	void addBonus(int val, BonusType type);
 	void addBonus(int val, BonusType type, BonusSubtypeID subtype);

lib/bonuses/Limiters.cpp

@@ -105,7 +105,7 @@ ILimiter::EDecision CCreatureTypeLimiter::limit(const BonusLimitationContext &co
 	if(!c)
 		return ILimiter::EDecision::NOT_APPLICABLE;
 	
-	auto accept =  c->getId() == creatureID || (includeUpgrades && creatureID.toCreature()->isMyUpgrade(c));
+	auto accept =  c->getId() == creatureID || (includeUpgrades && creatureID.toCreature()->isMyDirectOrIndirectUpgrade(c));
 	return accept ? ILimiter::EDecision::ACCEPT : ILimiter::EDecision::DISCARD;
 	//drop bonus if it's not our creature and (we don`t check upgrades or its not our upgrade)
 }

lib/entities/hero/CHeroHandler.cpp

@@ -140,35 +140,16 @@ void CHeroHandler::loadHeroSkills(CHero * hero, const JsonNode & node) const
 }
 
 /// creates standard H3 hero specialty for creatures
-static std::vector<std::shared_ptr<Bonus>> createCreatureSpecialty(CreatureID baseCreatureID)
+static std::vector<std::shared_ptr<Bonus>> createCreatureSpecialty(CreatureID cid)
 {
 	std::vector<std::shared_ptr<Bonus>> result;
-	std::set<CreatureID> targets;
-	targets.insert(baseCreatureID);
 
-	// go through entire upgrade chain and collect all creatures to which baseCreatureID can be upgraded
-	for (;;)
-	{
-		std::set<CreatureID> oldTargets = targets;
-
-		for(const auto & upgradeSourceID : oldTargets)
-		{
-			const CCreature * upgradeSource = upgradeSourceID.toCreature();
-			targets.insert(upgradeSource->upgrades.begin(), upgradeSource->upgrades.end());
-		}
-
-		if (oldTargets.size() == targets.size())
-			break;
-	}
-
-	for(CreatureID cid : targets)
-	{
 	const auto & specCreature = *cid.toCreature();
 	int stepSize = specCreature.getLevel() ? specCreature.getLevel() : 5;
 
 	{
 		auto bonus = std::make_shared<Bonus>();
-			bonus->limiter.reset(new CCreatureTypeLimiter(specCreature, false));
+		bonus->limiter.reset(new CCreatureTypeLimiter(specCreature, true));
 		bonus->type = BonusType::STACKS_SPEED;
 		bonus->val = 1;
 		result.push_back(bonus);
@@ -179,7 +160,7 @@ static std::vector<std::shared_ptr<Bonus>> createCreatureSpecialty(CreatureID ba
 		bonus->type = BonusType::PRIMARY_SKILL;
 		bonus->subtype = BonusSubtypeID(PrimarySkill::ATTACK);
 		bonus->val = 0;
-			bonus->limiter.reset(new CCreatureTypeLimiter(specCreature, false));
+		bonus->limiter.reset(new CCreatureTypeLimiter(specCreature, true));
 		bonus->updater.reset(new GrowsWithLevelUpdater(specCreature.getAttack(false), stepSize));
 		result.push_back(bonus);
 	}
@@ -189,11 +170,10 @@ static std::vector<std::shared_ptr<Bonus>> createCreatureSpecialty(CreatureID ba
 		bonus->type = BonusType::PRIMARY_SKILL;
 		bonus->subtype = BonusSubtypeID(PrimarySkill::DEFENSE);
 		bonus->val = 0;
-			bonus->limiter.reset(new CCreatureTypeLimiter(specCreature, false));
+		bonus->limiter.reset(new CCreatureTypeLimiter(specCreature, true));
 		bonus->updater.reset(new GrowsWithLevelUpdater(specCreature.getDefense(false), stepSize));
 		result.push_back(bonus);
 	}
-	}
 
 	return result;
 }

lib/mapObjects/CGCreature.cpp

@@ -529,7 +529,7 @@ void CGCreature::battleFinished(const CGHeroInstance *hero, const BattleResult &
 		const CCreature * cre = getCreature();
 		for(i = stacks.begin(); i != stacks.end(); i++)
 		{
-			if(cre->isMyUpgrade(i->second->getCreature()))
+			if(cre->isMyDirectUpgrade(i->second->getCreature()))
 			{
 				cb->changeStackType(StackLocation(id, i->first), cre); //un-upgrade creatures
 			}