From b076369469fd74e1d6e1a9858f6567fabf9cff02 Mon Sep 17 00:00:00 2001
From: Tyler Trahan <tyler@tylertrahan.com>
Date: Tue, 22 Feb 2022 14:00:37 -0700
Subject: [PATCH] Fix #9020: Update station coverage highlight when
 adding/removing tiles

---
 src/station_cmd.cpp | 13 +++++++++++--
 src/viewport.cpp    |  2 +-
 src/viewport_func.h |  1 +
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/station_cmd.cpp b/src/station_cmd.cpp
index 3a7114d9f4..e71f2de16a 100644
--- a/src/station_cmd.cpp
+++ b/src/station_cmd.cpp
@@ -741,9 +741,16 @@ static void DeleteStationIfEmpty(BaseStation *st)
 void Station::AfterStationTileSetChange(bool adding, StationType type)
 {
 	this->UpdateVirtCoord();
-	this->RecomputeCatchment();
 	DirtyCompanyInfrastructureWindows(this->owner);
-	if (adding) InvalidateWindowData(WC_STATION_LIST, this->owner, 0);
+
+	if (adding) {
+		this->RecomputeCatchment();
+		MarkCatchmentTilesDirty();
+		InvalidateWindowData(WC_STATION_LIST, this->owner, 0);
+	} else {
+		MarkCatchmentTilesDirty();
+		this->RecomputeCatchment();
+	}
 
 	switch (type) {
 		case STATION_RAIL:
@@ -1639,6 +1646,7 @@ CommandCost RemoveFromRailBaseStation(TileArea ta, std::vector<T *> &affected_st
 		if (st->train_station.tile == INVALID_TILE) {
 			st->facilities &= ~FACIL_TRAIN;
 			SetWindowWidgetDirty(WC_STATION_VIEW, st->index, WID_SV_TRAINS);
+			MarkCatchmentTilesDirty();
 			st->UpdateVirtCoord();
 			DeleteStationIfEmpty(st);
 		}
@@ -1675,6 +1683,7 @@ CommandCost CmdRemoveFromRailStation(TileIndex start, DoCommandFlag flags, uint3
 
 		if (st->train_station.tile == INVALID_TILE) SetWindowWidgetDirty(WC_STATION_VIEW, st->index, WID_SV_TRAINS);
 		st->MarkTilesDirty(false);
+		MarkCatchmentTilesDirty();
 		st->RecomputeCatchment();
 	}
 
diff --git a/src/viewport.cpp b/src/viewport.cpp
index b3ccba2b5e..11dfa15606 100644
--- a/src/viewport.cpp
+++ b/src/viewport.cpp
@@ -3496,7 +3496,7 @@ CommandCost CmdScrollViewport(TileIndex tile, DoCommandFlag flags, uint32 p1, ui
 	return CommandCost();
 }
 
-static void MarkCatchmentTilesDirty()
+void MarkCatchmentTilesDirty()
 {
 	if (_viewport_highlight_town != nullptr) {
 		MarkWholeScreenDirty();
diff --git a/src/viewport_func.h b/src/viewport_func.h
index 58e1706eb6..e60df4d31e 100644
--- a/src/viewport_func.h
+++ b/src/viewport_func.h
@@ -98,5 +98,6 @@ struct Town;
 
 void SetViewportCatchmentStation(const Station *st, bool sel);
 void SetViewportCatchmentTown(const Town *t, bool sel);
+void MarkCatchmentTilesDirty();
 
 #endif /* VIEWPORT_FUNC_H */