Allow adding/removing builder/factory build options dynamically. It m…

…ight not take into account some corner cases related to the use case. Should be harmless otherwise.

rts/Sim/Units/CommandAI/CommandAI.cpp

@@ -393,6 +393,11 @@ void CCommandAI::UpdateCommandDescription(unsigned int cmdDescIdx, SCommandDescr
 	if (!curCmdDesc->queueing)
 		nonQueingCommands.erase(curCmdDesc->id);
 
+	const bool boUpdate = (curCmdDesc->id != modCmdDesc.id);
+
+	if (boUpdate)
+		HandleBuildOptionRemoval(curCmdDesc->id);
+
 	// re-insert otherwise (possibly with a different cmdID!)
 	if (!modCmdDesc.queueing)
 		nonQueingCommands.insert(modCmdDesc.id);
@@ -402,6 +407,9 @@ void CCommandAI::UpdateCommandDescription(unsigned int cmdDescIdx, SCommandDescr
 	// update
 	possibleCommands[cmdDescIdx] = commandDescriptionCache.GetPtr(std::move(modCmdDesc));
 
+	if (boUpdate)
+		HandleBuildOptionInsertion(possibleCommands[cmdDescIdx]->id);
+
 	selectedUnitsHandler.PossibleCommandChange(owner);
 }
 
@@ -416,6 +424,8 @@ void CCommandAI::InsertCommandDescription(unsigned int cmdDescIdx, SCommandDescr
 		possibleCommands.insert(possibleCommands.begin() + cmdDescIdx, cmdDescPtr);
 	}
 
+	HandleBuildOptionInsertion(cmdDescPtr->id < 0);
+
 	// NB: cmdDesc is moved into cache, but id remains valid
 	if (!cmdDesc.queueing)
 		nonQueingCommands.insert(cmdDesc.id);
@@ -428,10 +438,14 @@ bool CCommandAI::RemoveCommandDescription(unsigned int cmdDescIdx)
 	if (cmdDescIdx >= possibleCommands.size())
 		return false;
 
-	if (!possibleCommands[cmdDescIdx]->queueing)
-		nonQueingCommands.erase(possibleCommands[cmdDescIdx]->id);
+	const auto* cmdDescPtr = possibleCommands[cmdDescIdx];
 
-	commandDescriptionCache.DecRef(*possibleCommands[cmdDescIdx]);
+	HandleBuildOptionRemoval(cmdDescPtr->id);
+
+	if (!cmdDescPtr->queueing)
+		nonQueingCommands.erase(cmdDescPtr->id);
+
+	commandDescriptionCache.DecRef(*cmdDescPtr);
 	// preserve order
 	possibleCommands.erase(possibleCommands.begin() + cmdDescIdx);
 	selectedUnitsHandler.PossibleCommandChange(owner);
@@ -476,6 +490,50 @@ void CCommandAI::AddCommandDependency(const Command& c) {
 }
 
 
+bool CCommandAI::HandleBuildOptionInsertion(int cmdId)
+{
+	if (cmdId >= 0)
+		return false;
+
+	if (auto* bcai = dynamic_cast<CBuilderCAI*>(this); bcai != nullptr)
+		bcai->buildOptions.insert(cmdId);
+	else if (auto* fcai = dynamic_cast<CFactoryCAI*>(this); fcai != nullptr)
+		fcai->buildOptions.insert(cmdId, 0);
+
+	return true;
+}
+
+bool CCommandAI::HandleBuildOptionRemoval(int cmdId)
+{
+	if (cmdId >= 0)
+		return false;
+
+	if (cmdId < 0) {
+		if (auto* bcai = dynamic_cast<CBuilderCAI*>(this); bcai != nullptr) {
+			// clear the removed unitDef from the construction queue
+			for (size_t i = 0; i < bcai->commandQue.size(); /*NOOP*/) {
+				if (const auto& q = bcai->commandQue[i]; q.GetID() == cmdId)
+					bcai->commandQue.erase(commandQue.begin() + i);
+				else
+					++i;
+			}
+			bcai->buildOptions.erase(cmdId);
+		}
+		else if (auto* fcai = dynamic_cast<CFactoryCAI*>(this); fcai != nullptr) {
+			// clear the removed unitDef from the construction queue
+			for (size_t i = 0; i < fcai->commandQue.size(); /*NOOP*/) {
+				if (const auto& q = fcai->commandQue[i]; q.GetID() == cmdId)
+					fcai->commandQue.erase(commandQue.begin() + i);
+				else
+					++i;
+			}
+			fcai->buildOptions.erase(cmdId);
+		}
+	}
+
+	return true;
+}
+
 bool CCommandAI::IsAttackCapable() const
 {
 	return (owner->unitDef->CanAttack());

rts/Sim/Units/CommandAI/CommandAI.h

@@ -134,6 +134,8 @@ class CCommandAI : public CObject
 	int lastSelectedCommandPage;
 
 protected:
+	bool HandleBuildOptionInsertion(int cmdId);
+	bool HandleBuildOptionRemoval(int cmdId);
 	// return true by default so non-AirCAI's trigger FinishCommand
 	virtual bool SelectNewAreaAttackTargetOrPos(const Command& ac) { return true; }