Correctly show results of CreatureTerrainLimiter outside of combat

lib/CCreatureSet.cpp

@@ -858,6 +858,11 @@ TerrainId CStackInstance::getNativeTerrain() const
 	return getFactionID().toEntity(LIBRARY)->getNativeTerrain();
 }
 
+TerrainId CStackInstance::getCurrentTerrain() const
+{
+	return armyObj->getCurrentTerrain();
+}
+
 void CStackInstance::deserializationFix()
 {
 	const CArmedInstance *armyBackup = _armyObj;

lib/CCreatureSet.h

@@ -141,6 +141,7 @@ public:
 
 	int32_t getInitiative(int turn = 0) const final;
 	TerrainId getNativeTerrain() const final;
+	TerrainId getCurrentTerrain() const;
 };
 
 class DLL_LINKAGE CCommanderInstance : public CStackInstance

lib/CStack.cpp

@@ -76,7 +76,6 @@ void CStack::localInit(BattleInfo * battleInfo)
 		attachTo(*army);
 		attachToSource(*typeID.toCreature());
 	}
-	nativeTerrain = getNativeTerrain(); //save nativeTerrain in the variable on the battle start to avoid dead lock
 	CUnitState::localInit(this); //it causes execution of the CStack::isOnNativeTerrain where nativeTerrain will be considered
 	position = initialPosition;
 }
@@ -316,14 +315,13 @@ bool CStack::canBeHealed() const
 
 bool CStack::isOnNativeTerrain() const
 {
-	//this code is called from CreatureTerrainLimiter::limit on battle start
+	auto nativeTerrain = getNativeTerrain();
-	auto res = nativeTerrain == ETerrainId::ANY_TERRAIN || nativeTerrain == battle->getTerrainType();
+	return nativeTerrain == ETerrainId::ANY_TERRAIN || getCurrentTerrain() == nativeTerrain;
-	return res;
 }
 
-bool CStack::isOnTerrain(TerrainId terrain) const
+TerrainId CStack::getCurrentTerrain() const
 {
-	return battle->getTerrainType() == terrain;
+	return battle->getTerrainType();
 }
 
 const CCreature * CStack::unitType() const

lib/CStack.h

@@ -28,7 +28,6 @@ class DLL_LINKAGE CStack final : public CBonusSystemNode, public battle::CUnitSt
 private:
 	ui32 ID = -1; //unique ID of stack
 	CreatureID typeID;
-	TerrainId nativeTerrain; //tmp variable to save native terrain value on battle init
 	ui32 baseAmount = -1;
 
 	PlayerColor owner; //owner - player color (255 for neutrals)
@@ -56,7 +55,7 @@ public:
 
 	bool canBeHealed() const; //for first aid tent - only harmed stacks that are not war machines
 	bool isOnNativeTerrain() const;
-	bool isOnTerrain(TerrainId terrain) const;
+	TerrainId getCurrentTerrain() const;
 
 	ui32 level() const;
 	si32 magicResistance() const override; //include aura of resistance

lib/bonuses/Limiters.cpp

@@ -280,18 +280,52 @@ CreatureTerrainLimiter::CreatureTerrainLimiter(TerrainId terrain):
 
 ILimiter::EDecision CreatureTerrainLimiter::limit(const BonusLimitationContext &context) const
 {
-	const CStack *stack = retrieveStackBattle(&context.node);
+	if (context.node.getNodeType() != CBonusSystemNode::STACK_BATTLE && context.node.getNodeType() != CBonusSystemNode::STACK_INSTANCE)
-	if(stack)
+		return ILimiter::EDecision::DISCARD;
+
+	if (terrainType == ETerrainId::NATIVE_TERRAIN)
 	{
-		if (terrainType == ETerrainId::NATIVE_TERRAIN && stack->isOnNativeTerrain())//terrainType not specified = native
+		auto selector = Selector::type()(BonusType::TERRAIN_NATIVE);
+
+		if(context.alreadyAccepted.getFirst(selector))
 			return ILimiter::EDecision::ACCEPT;
 
-		if(terrainType != ETerrainId::NATIVE_TERRAIN && stack->isOnTerrain(terrainType))
+		if(context.stillUndecided.getFirst(selector))
-			return ILimiter::EDecision::ACCEPT;
+			return ILimiter::EDecision::NOT_SURE;
 
+		// TODO: CStack and CStackInstance need some common base type that represents any stack
+		// Closest existing class is ACreature, however it is also used as base for CCreature, which is not a stack
+		if (context.node.getNodeType() == CBonusSystemNode::STACK_BATTLE)
+		{
+			const auto * unit = dynamic_cast<const CStack *>(&context.node);
+			auto unitNativeTerrain = unit->getFactionID().toEntity(LIBRARY)->getNativeTerrain();
+			if (unit->getCurrentTerrain() == unitNativeTerrain)
+				return ILimiter::EDecision::ACCEPT;
+		}
+		else
+		{
+			const auto * unit = dynamic_cast<const CStackInstance *>(&context.node);
+			auto unitNativeTerrain = unit->getFactionID().toEntity(LIBRARY)->getNativeTerrain();
+			if (unit->getCurrentTerrain() == unitNativeTerrain)
+				return ILimiter::EDecision::ACCEPT;
+		}
+	}
+	else
+	{
+		if (context.node.getNodeType() == CBonusSystemNode::STACK_BATTLE)
+		{
+			const auto * unit = dynamic_cast<const CStack *>(&context.node);
+			if (unit->getCurrentTerrain() == terrainType)
+				return ILimiter::EDecision::ACCEPT;
+		}
+		else
+		{
+			const auto * unit = dynamic_cast<const CStackInstance*>(&context.node);
+			if (unit->getCurrentTerrain() == terrainType)
+				return ILimiter::EDecision::ACCEPT;
+		}
 	}
 	return ILimiter::EDecision::DISCARD;
-	//TODO neutral creatues
 }
 
 std::string CreatureTerrainLimiter::toString() const

lib/mapObjects/CArmedInstance.cpp

@@ -17,6 +17,7 @@
 #include "../entities/faction/CTown.h"
 #include "../entities/faction/CTownHandler.h"
 #include "../gameState/CGameState.h"
+#include "../mapping/CMapDefines.h"
 #include "../texts/CGeneralTextHandler.h"
 
 VCMI_LIB_NAMESPACE_BEGIN
@@ -163,4 +164,9 @@ void CArmedInstance::serializeJsonOptions(JsonSerializeFormat & handler)
 	CCreatureSet::serializeJson(handler, "army", 7);
 }
 
+TerrainId CArmedInstance::getCurrentTerrain() const
+{
+	return cb->getTile(anchorPos())->getTerrainID();
+}
+
 VCMI_LIB_NAMESPACE_END

lib/mapObjects/CArmedInstance.h

@@ -48,6 +48,8 @@ public:
 	{
 		return this->tempOwner;
 	}
+
+	TerrainId getCurrentTerrain() const;
 	
 	void serializeJsonOptions(JsonSerializeFormat & handler) override;