From 1ad77f86977578e851c12287b52f873cb28ab273 Mon Sep 17 00:00:00 2001
From: Jonathan G Rennison <j.g.rennison@gmail.com>
Date: Fri, 24 Mar 2023 21:47:14 +0000
Subject: [PATCH] TBTR: Re-index recursive replacements on group topology
 changes

---
 src/group_cmd.cpp             |  6 ++++--
 src/tbtr_template_vehicle.cpp | 16 +++++++++++-----
 src/tbtr_template_vehicle.h   |  3 ++-
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/src/group_cmd.cpp b/src/group_cmd.cpp
index 7a2787d72b..0caca3b18d 100644
--- a/src/group_cmd.cpp
+++ b/src/group_cmd.cpp
@@ -385,6 +385,7 @@ CommandCost CmdCreateGroup(TileIndex tile, DoCommandFlag flags, uint32 p1, uint3
 			g->livery.colour1 = pg->livery.colour1;
 			g->livery.colour2 = pg->livery.colour2;
 			g->flags = pg->flags;
+			if (vt == VEH_TRAIN) ReindexTemplateReplacementsRecursive();
 		}
 
 		_new_group_id = g->index;
@@ -440,7 +441,7 @@ CommandCost CmdDeleteGroup(TileIndex tile, DoCommandFlag flags, uint32 p1, uint3
 		VehicleType vt = g->vehicle_type;
 
 		/* Delete all template replacements using the just deleted group */
-		DeleteTemplateReplacementsByGroupID(g->index);
+		DeleteTemplateReplacementsByGroupID(g);
 
 		/* notify tracerestrict that group is about to be deleted */
 		TraceRestrictRemoveGroupID(g->index);
@@ -506,6 +507,7 @@ CommandCost CmdAlterGroup(TileIndex tile, DoCommandFlag flags, uint32 p1, uint32
 		if (flags & DC_EXEC) {
 			g->parent = (pg == nullptr) ? INVALID_GROUP : pg->index;
 			GroupStatistics::UpdateAutoreplace(g->owner);
+			if (g->vehicle_type == VEH_TRAIN) ReindexTemplateReplacementsRecursive();
 
 			if (g->livery.in_use == 0) {
 				const Livery *livery = GetParentLivery(g);
@@ -1050,7 +1052,7 @@ void RemoveAllGroupsForCompany(const CompanyID company)
 {
 	for (Group *g : Group::Iterate()) {
 		if (company == g->owner) {
-			DeleteTemplateReplacementsByGroupID(g->index);
+			DeleteTemplateReplacementsByGroupID(g);
 			delete g;
 		}
 	}
diff --git a/src/tbtr_template_vehicle.cpp b/src/tbtr_template_vehicle.cpp
index e13abea6ce..7d7539c30d 100644
--- a/src/tbtr_template_vehicle.cpp
+++ b/src/tbtr_template_vehicle.cpp
@@ -54,7 +54,6 @@ INSTANTIATE_POOL_METHODS(TemplateReplacement)
 robin_hood::unordered_flat_map<GroupID, TemplateID> _template_replacement_index;
 robin_hood::unordered_flat_map<GroupID, TemplateID> _template_replacement_index_recursive;
 
-static void ReindexTemplateReplacementsRecursive();
 static void MarkTrainsInGroupAsPendingTemplateReplacement(GroupID gid, const TemplateVehicle *tv);
 
 void TemplateVehicleImageDimensions::SetFromTrain(const Train *t)
@@ -291,13 +290,20 @@ uint TemplateVehicle::NumGroupsUsingTemplate() const
 	return amount;
 }
 
-uint DeleteTemplateReplacementsByGroupID(GroupID g_id)
+uint DeleteTemplateReplacementsByGroupID(const Group *g)
 {
-	if (GetTemplateIDByGroupID(g_id) == INVALID_TEMPLATE) return 0;
+	if (g->vehicle_type != VEH_TRAIN) return 0;
+
+	if (g->parent != INVALID_GROUP) {
+		/* Erase any inherited replacement */
+		_template_replacement_index_recursive.erase(g->index);
+	}
+
+	if (GetTemplateIDByGroupID(g->index) == INVALID_TEMPLATE) return 0;
 
 	uint del_amount = 0;
 	for (const TemplateReplacement *tr : TemplateReplacement::Iterate()) {
-		if (tr->group == g_id) {
+		if (tr->group == g->index) {
 			delete tr;
 			del_amount++;
 		}
@@ -314,7 +320,7 @@ void ReindexTemplateReplacements()
 	ReindexTemplateReplacementsRecursive();
 }
 
-static void ReindexTemplateReplacementsRecursive()
+void ReindexTemplateReplacementsRecursive()
 {
 	_template_replacement_index_recursive.clear();
 	for (const Group *group : Group::Iterate()) {
diff --git a/src/tbtr_template_vehicle.h b/src/tbtr_template_vehicle.h
index c8630fbb3a..0958cafc92 100644
--- a/src/tbtr_template_vehicle.h
+++ b/src/tbtr_template_vehicle.h
@@ -227,9 +227,10 @@ bool IssueTemplateReplacement(GroupID gid, TemplateID tid);
 bool ShouldServiceTrainForTemplateReplacement(const Train *t, const TemplateVehicle *tv);
 void MarkTrainsUsingTemplateAsPendingTemplateReplacement(const TemplateVehicle *tv);
 
-uint DeleteTemplateReplacementsByGroupID(GroupID g_id);
+uint DeleteTemplateReplacementsByGroupID(const Group *g);
 
 void ReindexTemplateReplacements();
+void ReindexTemplateReplacementsRecursive();
 
 int GetTemplateVehicleEstimatedMaxAchievableSpeed(const TemplateVehicle *tv, int mass, const int speed_cap);