From 5bb1f00af323860446671381bee85b2b0ec6d023 Mon Sep 17 00:00:00 2001
From: LHoG <1476261+lhog@users.noreply.github.com>
Date: Sat, 10 Jun 2023 22:29:07 +0400
Subject: [PATCH] Allow adding/removing builder/factory build options
 dynamically. It might not take into account some corner cases related to the
 use case. Should be harmless otherwise.

---
 rts/Sim/Units/CommandAI/CommandAI.cpp | 64 +++++++++++++++++++++++++--
 rts/Sim/Units/CommandAI/CommandAI.h   |  2 +
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/rts/Sim/Units/CommandAI/CommandAI.cpp b/rts/Sim/Units/CommandAI/CommandAI.cpp
index 08a15bc95b..4196927fde 100644
--- a/rts/Sim/Units/CommandAI/CommandAI.cpp
+++ b/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());
diff --git a/rts/Sim/Units/CommandAI/CommandAI.h b/rts/Sim/Units/CommandAI/CommandAI.h
index 13407c2ed9..d587966ce3 100644
--- a/rts/Sim/Units/CommandAI/CommandAI.h
+++ b/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; }