Codechange: Make a k-d tree index of stations

14 files changed, 127 insertions, 33 deletions

diff --git a/projects/openttd_vs140.vcxproj b/projects/openttd_vs140.vcxproj
index 1d5365216..c1a23497a 100644
--- a/projects/openttd_vs140.vcxproj
+++ b/projects/openttd_vs140.vcxproj
@@ -646,6 +646,7 @@
     <ClInclude Include="..\src\station_base.h" />
     <ClInclude Include="..\src\station_func.h" />
     <ClInclude Include="..\src\station_gui.h" />
+    <ClInclude Include="..\src\station_kdtree.h" />
     <ClInclude Include="..\src\station_type.h" />
     <ClInclude Include="..\src\statusbar_gui.h" />
     <ClInclude Include="..\src\stdafx.h" />
diff --git a/projects/openttd_vs140.vcxproj.filters b/projects/openttd_vs140.vcxproj.filters
index b06f9576d..9bd33721a 100644
--- a/projects/openttd_vs140.vcxproj.filters
+++ b/projects/openttd_vs140.vcxproj.filters
@@ -1026,6 +1026,9 @@
     <ClInclude Include="..\src\station_gui.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\src\station_kdtree.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
     <ClInclude Include="..\src\station_type.h">
       <Filter>Header Files</Filter>
     </ClInclude>
diff --git a/projects/openttd_vs141.vcxproj b/projects/openttd_vs141.vcxproj
index 8d2a9dfc9..2b6c5c178 100644
--- a/projects/openttd_vs141.vcxproj
+++ b/projects/openttd_vs141.vcxproj
@@ -646,6 +646,7 @@
     <ClInclude Include="..\src\station_base.h" />
     <ClInclude Include="..\src\station_func.h" />
     <ClInclude Include="..\src\station_gui.h" />
+    <ClInclude Include="..\src\station_kdtree.h" />
     <ClInclude Include="..\src\station_type.h" />
     <ClInclude Include="..\src\statusbar_gui.h" />
     <ClInclude Include="..\src\stdafx.h" />
diff --git a/projects/openttd_vs141.vcxproj.filters b/projects/openttd_vs141.vcxproj.filters
index b06f9576d..9bd33721a 100644
--- a/projects/openttd_vs141.vcxproj.filters
+++ b/projects/openttd_vs141.vcxproj.filters
@@ -1026,6 +1026,9 @@
     <ClInclude Include="..\src\station_gui.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\src\station_kdtree.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
     <ClInclude Include="..\src\station_type.h">
       <Filter>Header Files</Filter>
     </ClInclude>
diff --git a/projects/openttd_vs142.vcxproj b/projects/openttd_vs142.vcxproj
index 56455c2f8..66b07576a 100644
--- a/projects/openttd_vs142.vcxproj
+++ b/projects/openttd_vs142.vcxproj
@@ -646,6 +646,7 @@
     <ClInclude Include="..\src\station_base.h" />
     <ClInclude Include="..\src\station_func.h" />
     <ClInclude Include="..\src\station_gui.h" />
+    <ClInclude Include="..\src\station_kdtree.h" />
     <ClInclude Include="..\src\station_type.h" />
     <ClInclude Include="..\src\statusbar_gui.h" />
     <ClInclude Include="..\src\stdafx.h" />
diff --git a/projects/openttd_vs142.vcxproj.filters b/projects/openttd_vs142.vcxproj.filters
index b06f9576d..9bd33721a 100644
--- a/projects/openttd_vs142.vcxproj.filters
+++ b/projects/openttd_vs142.vcxproj.filters
@@ -1026,6 +1026,9 @@
     <ClInclude Include="..\src\station_gui.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\src\station_kdtree.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
     <ClInclude Include="..\src\station_type.h">
       <Filter>Header Files</Filter>
     </ClInclude>
diff --git a/source.list b/source.list
index 969d90b9c..74068d48d 100644
--- a/source.list
+++ b/source.list
@@ -333,6 +333,7 @@ spritecache.h
 station_base.h
 station_func.h
 station_gui.h
+station_kdtree.h
 station_type.h
 statusbar_gui.h
 stdafx.h
diff --git a/src/misc.cpp b/src/misc.cpp
index c3130d85b..b1d480d82 100644
--- a/src/misc.cpp
+++ b/src/misc.cpp
@@ -28,6 +28,7 @@
 #include "core/pool_type.hpp"
 #include "game/game.hpp"
 #include "linkgraph/linkgraphschedule.h"
+#include "station_kdtree.h"
 #include "town_kdtree.h"
 
 #include "safeguards.h"
@@ -76,6 +77,7 @@ void InitializeGame(uint size_x, uint size_y, bool reset_date, bool reset_settin
 	LinkGraphSchedule::Clear();
 	PoolBase::Clean(PT_NORMAL);
 
+	RebuildStationKdtree();
 	RebuildTownKdtree();
 
 	ResetPersistentNewGRFData();
diff --git a/src/saveload/afterload.cpp b/src/saveload/afterload.cpp
index 07de26038..a80114f91 100644
--- a/src/saveload/afterload.cpp
+++ b/src/saveload/afterload.cpp
@@ -537,6 +537,7 @@ bool AfterLoadGame()
 	GamelogTestMode();
 
 	RebuildTownKdtree();
+	RebuildStationKdtree();
 
 	if (IsSavegameVersionBefore(SLV_98)) GamelogGRFAddList(_grfconfig);
 
diff --git a/src/station.cpp b/src/station.cpp
index 27063dcd7..5722fcdcf 100644
--- a/src/station.cpp
+++ b/src/station.cpp
@@ -21,6 +21,7 @@
 #include "vehiclelist.h"
 #include "core/pool_func.hpp"
 #include "station_base.h"
+#include "station_kdtree.h"
 #include "roadstop_base.h"
 #include "industry.h"
 #include "town.h"
@@ -36,6 +37,20 @@
 StationPool _station_pool("Station");
 INSTANTIATE_POOL_METHODS(Station)
 
+
+StationKdtree _station_kdtree(Kdtree_StationXYFunc);
+
+void RebuildStationKdtree()
+{
+	std::vector<StationID> stids;
+	BaseStation *st;
+	FOR_ALL_STATIONS(st) {
+		stids.push_back(st->index);
+	}
+	_station_kdtree.Build(stids.begin(), stids.end());
+}
+
+
 BaseStation::~BaseStation()
 {
 	free(this->name);
@@ -146,6 +161,8 @@ Station::~Station()
 	}
 
 	CargoPacket::InvalidateAllFrom(this->index);
+
+	_station_kdtree.Remove(this->index);
 }
 
 
diff --git a/src/station_base.h b/src/station_base.h
index a48daa653..b3a4a76c7 100644
--- a/src/station_base.h
+++ b/src/station_base.h
@@ -556,4 +556,6 @@ public:
 	}
 };
 
+void RebuildStationKdtree();
+
 #endif /* STATION_BASE_H */
diff --git a/src/station_cmd.cpp b/src/station_cmd.cpp
index cb588ef80..67ea4daf5 100644
--- a/src/station_cmd.cpp
+++ b/src/station_cmd.cpp
@@ -37,6 +37,7 @@
 #include "animated_tile_func.h"
 #include "elrail_func.h"
 #include "station_base.h"
+#include "station_kdtree.h"
 #include "roadstop_base.h"
 #include "newgrf_railtype.h"
 #include "waypoint_base.h"
@@ -359,19 +360,21 @@ static StringID GenerateStationName(Station *st, TileIndex tile, StationNaming n
 static Station *GetClosestDeletedStation(TileIndex tile)
 {
 	uint threshold = 8;
-	Station *best_station = NULL;
-	Station *st;
 
-	FOR_ALL_STATIONS(st) {
+	Station *best_station = NULL;
+	ForAllStationsRadius(tile, threshold, [&](Station *st) {
 		if (!st->IsInUse() && st->owner == _current_company) {
 			uint cur_dist = DistanceManhattan(tile, st->xy);
 
 			if (cur_dist < threshold) {
 				threshold = cur_dist;
 				best_station = st;
+			} else if (cur_dist == threshold && best_station != NULL) {
+				/* In case of a tie, lowest station ID wins */
+				if (st->index < best_station->index) best_station = st;
 			}
 		}
-	}
+	});
 
 	return best_station;
 }
@@ -667,8 +670,13 @@ static void UpdateStationSignCoord(BaseStation *st)
 	if (r->IsEmpty()) return; // no tiles belong to this station
 
 	/* clamp sign coord to be inside the station rect */
-	st->xy = TileXY(ClampU(TileX(st->xy), r->left, r->right), ClampU(TileY(st->xy), r->top, r->bottom));
-	st->UpdateVirtCoord();
+	TileIndex new_xy = TileXY(ClampU(TileX(st->xy), r->left, r->right), ClampU(TileY(st->xy), r->top, r->bottom));
+	if (new_xy != st->xy) {
+		_station_kdtree.Remove(st->index);
+		st->xy = new_xy;
+		_station_kdtree.Insert(st->index);
+		st->UpdateVirtCoord();
+	}
 
 	if (!Station::IsExpected(st)) return;
 	Station *full_station = Station::From(st);
@@ -706,6 +714,7 @@ static CommandCost BuildStationPart(Station **st, DoCommandFlag flags, bool reus
 
 		if (flags & DC_EXEC) {
 			*st = new Station(area.tile);
+			_station_kdtree.Insert((*st)->index);
 
 			(*st)->town = ClosestTownFromTile(area.tile, UINT_MAX);
 			(*st)->string_id = GenerateStationName(*st, area.tile, name_class);
@@ -3694,11 +3703,8 @@ void StationMonthlyLoop()
 
 void ModifyStationRatingAround(TileIndex tile, Owner owner, int amount, uint radius)
 {
-	Station *st;
-
-	FOR_ALL_STATIONS(st) {
-		if (st->owner == owner &&
-				DistanceManhattan(tile, st->xy) <= radius) {
+	ForAllStationsRadius(tile, radius, [&](Station *st) {
+		if (st->owner == owner) {
 			for (CargoID i = 0; i < NUM_CARGO; i++) {
 				GoodsEntry *ge = &st->goods[i];
 
@@ -3707,7 +3713,7 @@ void ModifyStationRatingAround(TileIndex tile, Owner owner, int amount, uint rad
 				}
 			}
 		}
-	}
+	});
 }
 
 static uint UpdateStationWaiting(Station *st, CargoID type, uint amount, SourceType source_type, SourceID source_id)
@@ -3947,6 +3953,7 @@ void BuildOilRig(TileIndex tile)
 	}
 
 	Station *st = new Station(tile);
+	_station_kdtree.Insert(st->index);
 	st->town = ClosestTownFromTile(tile, UINT_MAX);
 
 	st->string_id = GenerateStationName(st, tile, STATIONNAMING_OILRIG);
diff --git a/src/station_kdtree.h b/src/station_kdtree.h
new file mode 100644
index 000000000..321bbacc6
--- /dev/null
+++ b/src/station_kdtree.h
@@ -0,0 +1,42 @@
+/*
+ * This file is part of OpenTTD.
+ * OpenTTD is free software; you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, version 2.
+ * OpenTTD is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+ * See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with OpenTTD. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+/** @file station_kdtree.h Declarations for accessing the k-d tree of stations */
+
+#ifndef STATION_KDTREE_H
+#define STATION_KDTREE_H
+
+#include "core/kdtree.hpp"
+#include "core/math_func.hpp"
+#include "station_base.h"
+#include "map_func.h"
+
+inline uint16 Kdtree_StationXYFunc(StationID stid, int dim) { return (dim == 0) ? TileX(BaseStation::Get(stid)->xy) : TileY(BaseStation::Get(stid)->xy); }
+typedef Kdtree<StationID, decltype(&Kdtree_StationXYFunc), uint16, int> StationKdtree;
+extern StationKdtree _station_kdtree;
+
+/**
+ * Call a function on all stations whose sign is within a radius of a center tile.
+ * @param center  Central tile to search around.
+ * @param radius  Distance in both X and Y to search within.
+ * @param func    The function to call, must take a single parameter which is Station*.
+ */
+template <typename Func>
+void ForAllStationsRadius(TileIndex center, uint radius, Func func)
+{
+	uint16 x1, y1, x2, y2;
+	x1 = (uint16)max<int>(0, TileX(center) - radius);
+	x2 = (uint16)min<int>(TileX(center) + radius + 1, MapSizeX());
+	y1 = (uint16)max<int>(0, TileY(center) - radius);
+	y2 = (uint16)min<int>(TileY(center) + radius + 1, MapSizeY());
+
+	_station_kdtree.FindContained(x1, y1, x2, y2, [&](StationID id) {
+		func(Station::Get(id));
+	});
+}
+
+#endif
diff --git a/src/town_cmd.cpp b/src/town_cmd.cpp
index e56cfd417..0f1ccbba3 100644
--- a/src/town_cmd.cpp
+++ b/src/town_cmd.cpp
@@ -18,6 +18,7 @@
 #include "command_func.h"
 #include "industry.h"
 #include "station_base.h"
+#include "station_kdtree.h"
 #include "company_base.h"
 #include "news_func.h"
 #include "error.h"
@@ -3197,6 +3198,21 @@ CommandCost CmdDoTownAction(TileIndex tile, DoCommandFlag flags, uint32 p1, uint
 	return cost;
 }
 
+template <typename Func>
+static void ForAllStationsNearTown(Town *t, Func func)
+{
+	/* Ideally the search radius should be close to the actual town zone 0 radius.
+	 * The true radius is not stored or calculated anywhere, only the squared radius. */
+	/* The efficiency of this search might be improved for large towns and many stations on the map,
+	 * by using an integer square root approximation giving a value not less than the true square root. */
+	uint search_radius = t->cache.squared_town_zone_radius[0] / 2;
+	ForAllStationsRadius(t->xy, search_radius, [&](const Station * st) {
+		if (DistanceSquare(st->xy, t->xy) <= t->cache.squared_town_zone_radius[0]) {
+			func(st);
+		}
+	});
+}
+
 static void UpdateTownRating(Town *t)
 {
 	/* Increase company ratings if they're low */
@@ -3207,22 +3223,19 @@ static void UpdateTownRating(Town *t)
 		}
 	}
 
-	const Station *st;
-	FOR_ALL_STATIONS(st) {
-		if (DistanceSquare(st->xy, t->xy) <= t->cache.squared_town_zone_radius[0]) {
-			if (st->time_since_load <= 20 || st->time_since_unload <= 20) {
-				if (Company::IsValidID(st->owner)) {
-					int new_rating = t->ratings[st->owner] + RATING_STATION_UP_STEP;
-					t->ratings[st->owner] = min(new_rating, INT16_MAX); // do not let it overflow
-				}
-			} else {
-				if (Company::IsValidID(st->owner)) {
-					int new_rating = t->ratings[st->owner] + RATING_STATION_DOWN_STEP;
-					t->ratings[st->owner] = max(new_rating, INT16_MIN);
-				}
+	ForAllStationsNearTown(t, [&](const Station *st) {
+		if (st->time_since_load <= 20 || st->time_since_unload <= 20) {
+			if (Company::IsValidID(st->owner)) {
+				int new_rating = t->ratings[st->owner] + RATING_STATION_UP_STEP;
+				t->ratings[st->owner] = min(new_rating, INT16_MAX); // do not let it overflow
+			}
+		} else {
+			if (Company::IsValidID(st->owner)) {
+				int new_rating = t->ratings[st->owner] + RATING_STATION_DOWN_STEP;
+				t->ratings[st->owner] = max(new_rating, INT16_MIN);
 			}
 		}
-	}
+	});
 
 	/* clamp all ratings to valid values */
 	for (uint i = 0; i < MAX_COMPANIES; i++) {
@@ -3257,14 +3270,11 @@ static void UpdateTownGrowCounter(Town *t, uint16 prev_growth_rate)
 static int CountActiveStations(Town *t)
 {
 	int n = 0;
-	const Station *st;
-	FOR_ALL_STATIONS(st) {
-		if (DistanceSquare(st->xy, t->xy) <= t->cache.squared_town_zone_radius[0]) {
-			if (st->time_since_load <= 20 || st->time_since_unload <= 20) {
-				n++;
-			}
+	ForAllStationsNearTown(t, [&](const Station * st) {
+		if (st->time_since_load <= 20 || st->time_since_unload <= 20) {
+			n++;
 		}
-	}
+	});
 	return n;
 }