1 files changed, 79 insertions, 17 deletions

diff --git a/src/pbs.cpp b/src/pbs.cpp
index 71a9689a0..12e54d7f4 100644
--- a/src/pbs.cpp
+++ b/src/pbs.cpp
@@ -14,6 +14,7 @@
 #include "settings_type.h"
 #include "road_func.h"
 #include "vehicle_base.h"
+#include "vehicle_func.h"
 #include "yapf/follow_track.hpp"
 
 /**
@@ -159,23 +160,12 @@ bool TryReserveRailTrack(TileIndex tile, Track t)
 	}
 }
 
-/**
- * Follow a train reservation to the last tile.
- *
- * @param v the vehicle
- * @returns The last tile of the reservation or the current train tile if no reservation present.
- */
-PBSTileInfo FollowTrainReservation(const Vehicle *v)
-{
-	assert(v->type == VEH_TRAIN);
-
-	TileIndex tile = v->tile;
-	Trackdir  trackdir = GetVehicleTrackdir(v);
-
-	if (IsRailDepotTile(tile) && !GetRailDepotReservation(tile)) return PBSTileInfo(tile, trackdir, false);
 
+/** Follow a reservation starting from a specific tile to the end. */
+static PBSTileInfo FollowReservation(Owner o, RailTypes rts, TileIndex tile, Trackdir trackdir, bool ignore_oneway = false)
+{
 	/* Do not disallow 90 deg turns as the setting might have changed between reserving and now. */
-	CFollowTrackRail ft(v, GetRailTypeInfo(v->u.rail.railtype)->compatible_railtypes);
+	CFollowTrackRail ft(o, rts);
 	while (ft.Follow(tile, trackdir)) {
 		TrackdirBits reserved = (TrackdirBits)(ft.m_new_td_bits & (GetReservedTrackbits(ft.m_new_tile) * 0x101));
 
@@ -187,7 +177,7 @@ PBSTileInfo FollowTrainReservation(const Vehicle *v)
 
 		/* One-way signal against us. The reservation can't be ours as it is not
 		 * a safe position from our direction and we can never pass the signal. */
-		if (HasOnewaySignalBlockingTrackdir(ft.m_new_tile, new_trackdir)) break;
+		if (!ignore_oneway && HasOnewaySignalBlockingTrackdir(ft.m_new_tile, new_trackdir)) break;
 
 		tile = ft.m_new_tile;
 		trackdir = new_trackdir;
@@ -198,7 +188,79 @@ PBSTileInfo FollowTrainReservation(const Vehicle *v)
 		if (IsTileType(tile, MP_RAILWAY) && HasSignalOnTrackdir(tile, trackdir) && !IsPbsSignal(GetSignalType(tile, TrackdirToTrack(trackdir)))) break;
 	}
 
-	return PBSTileInfo(tile, trackdir, IsSafeWaitingPosition(v, tile, trackdir, true, _settings_game.pf.forbid_90_deg));
+	return PBSTileInfo(tile, trackdir, false);
+}
+
+/**
+ * Follow a train reservation to the last tile.
+ *
+ * @param v the vehicle
+ * @returns The last tile of the reservation or the current train tile if no reservation present.
+ */
+PBSTileInfo FollowTrainReservation(const Vehicle *v)
+{
+	assert(v->type == VEH_TRAIN);
+
+	TileIndex tile = v->tile;
+	Trackdir  trackdir = GetVehicleTrackdir(v);
+
+	if (IsRailDepotTile(tile) && !GetRailDepotReservation(tile)) return PBSTileInfo(tile, trackdir, false);
+
+	PBSTileInfo res = FollowReservation(v->owner, GetRailTypeInfo(v->u.rail.railtype)->compatible_railtypes, tile, trackdir);
+	res.okay = IsSafeWaitingPosition(v, res.tile, res.trackdir, true, _settings_game.pf.forbid_90_deg);
+	return res;
+}
+
+/** Callback for VehicleFromPos to find a train on a specific track. */
+static Vehicle *FindTrainOnTrackEnum(Vehicle *v, void *data)
+{
+	PBSTileInfo info = *(PBSTileInfo *)data;
+
+	if (v->type == VEH_TRAIN && !(v->vehstatus & VS_CRASHED) && HasBit((TrackBits)v->u.rail.track, TrackdirToTrack(info.trackdir))) return v;
+
+	return NULL;
+}
+
+/**
+ * Find the train which has reserved a specific path.
+ *
+ * @param tile A tile on the path.
+ * @param track A reserved track on the tile.
+ * @return The vehicle holding the reservation or NULL if the path is stray.
+ */
+Vehicle *GetTrainForReservation(TileIndex tile, Track track)
+{
+	assert(HasReservedTracks(tile, TrackToTrackBits(track)));
+	Trackdir  trackdir = TrackToTrackdir(track);
+
+	RailTypes rts = GetRailTypeInfo(GetTileRailType(tile))->compatible_railtypes;
+
+	/* Follow the path from tile to both ends, one of the end tiles should
+	 * have a train on it. We need FollowReservation to ignore one-way signals
+	 * here, as one of the two search directions will be the "wrong" way. */
+	for (int i = 0; i < 2; ++i, trackdir = ReverseTrackdir(trackdir)) {
+		PBSTileInfo dest = FollowReservation(GetTileOwner(tile), rts, tile, trackdir, true);
+
+		Vehicle *v = VehicleFromPos(dest.tile, &dest, FindTrainOnTrackEnum);
+		if (v != NULL) return v->First();
+
+		/* Special case for stations: check the whole platform for a vehicle. */
+		if (IsRailwayStationTile(dest.tile)) {
+			TileIndexDiff diff = TileOffsByDiagDir(TrackdirToExitdir(ReverseTrackdir(dest.trackdir)));
+			for (TileIndex st_tile = dest.tile + diff; IsCompatibleTrainStationTile(st_tile, dest.tile); st_tile += diff) {
+				v = VehicleFromPos(st_tile, &dest, FindTrainOnTrackEnum);
+				if (v != NULL) return v->First();
+			}
+		}
+
+		/* Special case for bridges/tunnels: check the other end as well. */
+		if (IsTileType(dest.tile, MP_TUNNELBRIDGE)) {
+			v = VehicleFromPos(GetOtherTunnelBridgeEnd(dest.tile), &dest, FindTrainOnTrackEnum);
+			if (v != NULL) return v->First();
+		}
+	}
+
+	return NULL;
 }
 
 /**