[GEN][ZH] Fix AIGroup memory management and leaks (#1214)

Author: xezon

Core/GameEngine/Include/Common/GameDefines.h

@@ -29,6 +29,14 @@
 #define RETAIL_COMPATIBLE_XFER_SAVE (1) // Game is expected to be Xfer Save compatible with retail Generals 1.08, Zero Hour 1.04
 #endif
 
+// This is essentially synonymous for RETAIL_COMPATIBLE_CRC. There is a lot wrong with AIGroup, such as use-after-free, double-free, leaks,
+// but we cannot touch it much without breaking retail compatibility. Do not shy away from using massive hacks when fixing issues with AIGroup,
+// but put them behind this macro.
+
+#ifndef RETAIL_COMPATIBLE_AIGROUP
+#define RETAIL_COMPATIBLE_AIGROUP (1) // AIGroup logic is expected to be CRC compatible with retail Generals 1.08, Zero Hour 1.04
+#endif
+
 #ifndef ENABLE_GAMETEXT_SUBSTITUTES
 #define ENABLE_GAMETEXT_SUBSTITUTES (1) // The code can provide substitute texts when labels and strings are missing in the STR or CSF translation file
 #endif

Generals/Code/GameEngine/Include/Common/Player.h

@@ -608,7 +608,7 @@ class Player : public Snapshot
 	// add to the player's current selection this hotkey team.
 	void processAddTeamGameMessage(Int hotkeyNum, GameMessage *msg);
 
-	// returns an AIGroup object that is the currently selected group.
+	// fills an AIGroup object that is the currently selected group.
 	void getCurrentSelectionAsAIGroup(AIGroup *group);
 
 	// sets the currently selected group to be the given AIGroup

Generals/Code/GameEngine/Include/GameLogic/AI.h

@@ -37,6 +37,8 @@
 #include "Common/GameType.h"
 #include "GameLogic/Damage.h"
 #include "Common/STLTypedefs.h"
+#include "refcount.h"
+#include "ref_ptr.h"
 
 class AIGroup;
 class AttackPriorityInfo;
@@ -57,6 +59,12 @@ enum WeaponSetType CPP_11(: Int);
 enum WeaponLockType CPP_11(: Int);
 enum SpecialPowerType CPP_11(: Int);
 
+#if RETAIL_COMPATIBLE_AIGROUP
+typedef AIGroup* AIGroupPtr;
+#else
+typedef RefCountPtr<AIGroup> AIGroupPtr;
+#endif
+
 typedef std::vector<ObjectID> VecObjectID;
 typedef VecObjectID::iterator VecObjectIDIt;
 
@@ -265,7 +273,7 @@ class AI : public SubsystemInterface, public Snapshot
 	void loadPostProcess( void );
 
 	// AI Groups -----------------------------------------------------------------------------------------------
-	AIGroup *createGroup( void );					///< instantiate a new AI Group
+	AIGroupPtr createGroup( void ); ///< instantiate a new AI Group
 	void destroyGroup( AIGroup *group );	///< destroy the given AI Group
 	AIGroup *findGroup( UnsignedInt id );	///< return the AI Group with the given ID
 
@@ -851,6 +859,12 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 	void xfer( Xfer *xfer );
 	void loadPostProcess( void );
 
+#if !RETAIL_COMPATIBLE_AIGROUP
+	void Add_Ref() const { m_refCount.Add_Ref(); }
+	void Release_Ref() const { m_refCount.Release_Ref(destroy, this); }
+	UnsignedShort Num_Refs() const { return m_refCount.Num_Refs(); }
+#endif
+
 	void groupMoveToPosition( const Coord3D *pos, Bool addWaypoint, CommandSourceType cmdSource );
 	void groupMoveToAndEvacuate( const Coord3D *pos, CommandSourceType cmdSource );			///< move to given position(s)
 	void groupMoveToAndEvacuateAndExit( const Coord3D *pos, CommandSourceType cmdSource );			///< move to given position & unload transport.
@@ -938,13 +952,15 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 
 	void add( Object *obj );								///< add object to group
 
-	// returns true if remove destroyed the group for us.
+	// Returns true if the group was emptied.
 	Bool remove( Object *obj);
+
+	void removeAll( void );
 
 	// If the group contains any objects not owned by ownerPlayer, return TRUE.
 	Bool containsAnyObjectsNotOwnedByPlayer( const Player *ownerPlayer );
 
-	// Remove any objects that aren't owned by the player, and return true if the group was destroyed due to emptiness
+	// Removes any objects that aren't owned by the player, and returns true if the group was emptied.
 	Bool removeAnyObjectsNotOwnedByPlayer( const Player *ownerPlayer );
 
 	UnsignedInt getID( void );
@@ -974,6 +990,13 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 	friend class AI;
 	AIGroup( void );
 
+#if !RETAIL_COMPATIBLE_AIGROUP
+	static void destroy(AIGroup* self)
+	{
+		TheAI->destroyGroup(self);
+	}
+#endif
+
 	void recompute( void );									///< recompute various group info, such as speed, leader, etc
 
 	ListObjectPtr m_memberList;							///< the list of member Objects
@@ -982,6 +1005,10 @@ class AIGroup : public MemoryPoolObject, public Snapshot
 	Real m_speed;														///< maximum speed of group (slowest member)
 	Bool m_dirty;														///< "dirty bit" - if true then group speed, leader, needs recompute
 
+#if !RETAIL_COMPATIBLE_AIGROUP
+	RefCountValue<UnsignedShort> m_refCount; ///< the reference counter
+#endif
+
 	UnsignedInt m_id;												///< the unique ID of this group
 	Path *m_groundPath;											///< Group ground path.
 

Generals/Code/GameEngine/Include/GameLogic/Object.h

@@ -31,6 +31,7 @@
 #define _OBJECT_H_
 
 #include "Lib/BaseType.h"
+#include "ref_ptr.h"
 
 #include "Common/Geometry.h"
 #include "Common/Snapshot.h"
@@ -663,7 +664,11 @@ class Object : public Thing, public Snapshot
 
 	GeometryInfo	m_geometryInfo;
 
+#if RETAIL_COMPATIBLE_AIGROUP
 	AIGroup*			m_group;								///< if non-NULL, we are part of this group of agents
+#else
+	RefCountPtr<AIGroup> m_group; ///< if non-NULL, we are part of this group of agents
+#endif
 
 	// These will last for my lifetime.  I will reuse them and reset them.  The truly dynamic ones are in PartitionManager
 	SightingInfo	*m_partitionLastLook;		///< Where and for whom I last looked, so I can undo its effects when I stop

Generals/Code/GameEngine/Source/GameLogic/AI/AI.cpp

@@ -330,6 +330,8 @@ void AI::reset( void )
 		m_aiData = m_aiData->m_next;
 		delete cur;
 	}
+
+#if RETAIL_COMPATIBLE_AIGROUP
 	while (m_groupList.size())
 	{
 		AIGroup *groupToRemove = m_groupList.front();
@@ -342,6 +344,12 @@ void AI::reset( void )
 			m_groupList.pop_front(); // NULL group, just kill from list.  Shouldn't really happen, but just in case.
 		}
 	}
+#else
+	DEBUG_ASSERTCRASH(m_groupList.empty(), ("AI::m_groupList is expected empty already\n"));
+
+	m_groupList.clear(); // Clear just in case...
+#endif
+
 	m_nextGroupID = 0;
 	m_nextFormationID = NO_FORMATION_ID;
 	getNextFormationID(); // increment once past NO_FORMATION_ID.  jba.
@@ -440,14 +448,22 @@ void AI::parseAiDataDefinition( INI* ini )
 /**
  * Create a new AI Group
  */
-AIGroup *AI::createGroup( void )
+AIGroupPtr AI::createGroup( void )
 {
 	// create a new instance
+#if RETAIL_COMPATIBLE_AIGROUP
 	AIGroup *group = newInstance(AIGroup);
+#else
+	AIGroupPtr group = AIGroupPtr::Create_NoAddRef(newInstance(AIGroup));
+#endif
 
 	// add it to the list
 //	DEBUG_LOG(("***AIGROUP %x is being added to m_groupList.\n", group ));
+#if RETAIL_COMPATIBLE_AIGROUP
 	m_groupList.push_back( group );
+#else
+	m_groupList.push_back( group.Peek() );
+#endif
 
 	return group;
 }

Generals/Code/GameEngine/Source/GameLogic/AI/AIGroup.cpp

@@ -92,6 +92,9 @@ AIGroup::~AIGroup()
 {
 //	DEBUG_LOG(("***AIGROUP %x is being destructed.\n", this));
 	// disassociate each member from the group
+
+#if RETAIL_COMPATIBLE_AIGROUP
+
 	std::list<Object *>::iterator i;
 	for( i = m_memberList.begin(); i != m_memberList.end(); /* empty */ )
 	{
@@ -106,6 +109,13 @@ AIGroup::~AIGroup()
 			i = m_memberList.erase(i);
 		}
 	}
+
+#else
+
+	removeAll();
+
+#endif
+
 	if (m_groundPath) {
 		deleteInstance(m_groundPath);
 		m_groundPath = NULL;
@@ -205,6 +215,11 @@ void AIGroup::add( Object *obj )
  */
 Bool AIGroup::remove( Object *obj )
 {
+#if !RETAIL_COMPATIBLE_AIGROUP
+	// Defer deletion until the end of this function.
+	AIGroupPtr refThis = AIGroupPtr::Create_AddRef(this);
+#endif
+
 //	DEBUG_LOG(("***AIGROUP %x is removing Object %x (%s).\n", this, obj, obj->getTemplate()->getName().str()));
 	std::list<Object *>::iterator i = std::find( m_memberList.begin(), m_memberList.end(), obj );
 
@@ -223,15 +238,42 @@ Bool AIGroup::remove( Object *obj )
 	// list has changed, properties need recomputation
 	m_dirty = true;
 
-	// if the group is empty, no-one is using it any longer, so destroy it
 	if (isEmpty()) {
+#if RETAIL_COMPATIBLE_AIGROUP
+		// if the group is empty, no-one is using it any longer, so destroy it
 		TheAI->destroyGroup( this );
+#endif
 		return TRUE;
 	}
 
 	return FALSE;
 }
 
+/**
+ * Remove all objects from group
+ */
+void AIGroup::removeAll( void )
+{
+#if !RETAIL_COMPATIBLE_AIGROUP
+	// Defer deletion until the end of this function.
+	AIGroupPtr refThis = AIGroupPtr::Create_AddRef(this);
+#endif
+
+	std::list<Object *> memberList;
+	memberList.swap(m_memberList);
+	m_memberListSize = 0;
+
+	std::list<Object *>::iterator i;
+	for ( i = memberList.begin(); i != memberList.end(); ++i )
+	{
+		Object *member = *i;
+		if (member)
+			member->leaveGroup();
+	}
+
+	m_dirty = true;
+}
+
 /**
  * If the group contains any objects not owned by ownerPlayer, return TRUE.
  */
@@ -2722,6 +2764,11 @@ void AIGroup::groupCheer( CommandSourceType cmdSource )
 	*/
 void AIGroup::groupSell( CommandSourceType cmdSource )
 {
+#if !RETAIL_COMPATIBLE_AIGROUP
+	// Defer deletion until the end of this function.
+	AIGroupPtr refThis = AIGroupPtr::Create_AddRef(this);
+#endif
+
 	std::list<Object *>::iterator i;
 	std::vector<Object *> groupObjectsCopy;
 	groupObjectsCopy.reserve(m_memberListSize);

Generals/Code/GameEngine/Source/GameLogic/AI/AIPlayer.cpp

@@ -893,11 +893,15 @@ void AIPlayer::guardSupplyCenter( Team *team, Int minSupplies )
 	}
 	if (warehouse) {
 
-		AIGroup* theGroup = TheAI->createGroup();
+		AIGroupPtr theGroup = TheAI->createGroup();
 		if (!theGroup) {
 			return;
 		}
+#if RETAIL_COMPATIBLE_AIGROUP
 		team->getTeamAsAIGroup(theGroup);
+#else
+		team->getTeamAsAIGroup(theGroup.Peek());
+#endif
 		Coord3D location = *warehouse->getPosition();
 		// It's probably a defensive move - position towards the enemy.
 		Region2D bounds;
@@ -2467,9 +2471,9 @@ void AIPlayer::checkReadyTeams( void )
 					/*
 					if (team->m_team->getPrototype()->getTemplateInfo()->m_hasHomeLocation && 
 							!team->m_reinforcement) {
- 						AIGroup* theGroup = TheAI->createGroup();
+ 						AIGroupPtr theGroup = TheAI->createGroup();
 						if (theGroup) {
-							team->m_team->getTeamAsAIGroup(theGroup);
+							team->m_team->getTeamAsAIGroup(theGroup.Peek());
 							Coord3D destination = team->m_team->getPrototype()->getTemplateInfo()->m_homeLocation;
 							theGroup->groupTightenToPosition( &destination, false, CMD_FROM_AI );
 							team->m_frameStarted = TheGameLogic->getFrame();

Generals/Code/GameEngine/Source/GameLogic/AI/AIStates.cpp

@@ -4053,8 +4053,12 @@ StateReturnType AIFollowWaypointPathState::update()
 	if (m_moveAsGroup) {
 		if (obj->getControllingPlayer()->isSkirmishAIPlayer()) {
 			Team *team = obj->getTeam();
-			AIGroup *group = TheAI->createGroup();
+			AIGroupPtr group = TheAI->createGroup();
+#if RETAIL_COMPATIBLE_AIGROUP
 			team->getTeamAsAIGroup(group);
+#else
+			team->getTeamAsAIGroup(group.Peek());
+#endif
 
 			Coord3D pos;
 			group->getCenter(&pos);

Generals/Code/GameEngine/Source/GameLogic/Object/Object.cpp

@@ -563,8 +563,7 @@ Object::~Object()
 		ThePartitionManager->unRegisterObject( this );
 
 	// if we are in a group, remove us
-	if (m_group)
-		m_group->remove( this );
+	leaveGroup();
 
 	// note, do NOT free these, there are just a shadow copy!
 	m_ai = NULL;
@@ -3841,7 +3840,7 @@ void Object::xfer( Xfer *xfer )
 	}
 
 	// Doesn't need to be saved.  These are created as needed.  jba.
-	//AIGroup*		m_group;															///< if non-NULL, we are part of this group of agents
+	//m_group;
 
 	// don't need to save m_partitionData.
 	DEBUG_ASSERTCRASH(!(xfer->getXferMode() == XFER_LOAD && m_partitionData == NULL), ("should not be in partitionmgr yet"));
@@ -5446,7 +5445,11 @@ RadarPriorityType Object::getRadarPriority( void ) const
 // ------------------------------------------------------------------------------------------------
 AIGroup *Object::getGroup(void)
 {
+#if RETAIL_COMPATIBLE_AIGROUP
 	return m_group;
+#else
+	return m_group.Peek();
+#endif
 }
 
 //-------------------------------------------------------------------------------------------------
@@ -5456,7 +5459,11 @@ void Object::enterGroup( AIGroup *group )
 	// if we are in another group, remove ourselves from it first
 	leaveGroup();
 
+#if RETAIL_COMPATIBLE_AIGROUP
 	m_group = group;
+#else
+	m_group = AIGroupPtr::Create_AddRef(group);
+#endif
 }
 
 //-------------------------------------------------------------------------------------------------
@@ -5467,7 +5474,7 @@ void Object::leaveGroup( void )
 	if (m_group)
 	{
 		// to avoid recursion, set m_group to NULL before removing
-		AIGroup *group = m_group;
+		AIGroupPtr group = m_group;
 		m_group = NULL;
 		group->remove( this );
 	}