Tracerestrict: Handle slot acquire and wait at PBS on intermediary reservation signals

Wait at PBS is applied as if on the starting signal

src/pathfinder/yapf/yapf_node_rail.hpp

@@ -194,15 +194,15 @@ struct CYapfRailNodeT
 		m_segment->m_last_td = td;
 	}
 
-	template <class Tbase, class Tfunc, class Tpf>
-	bool IterateTiles(const Train *v, Tpf &yapf, Tbase &obj, bool (Tfunc::*func)(TileIndex, Trackdir)) const
+	template <class Tbase, class Tpf, class Tfunc>
+	bool IterateTiles(const Train *v, Tpf &yapf, Tfunc func) const
 	{
 		typename Tbase::TrackFollower ft(v, yapf.GetCompatibleRailTypes());
 		TileIndex cur = base::GetTile();
 		Trackdir  cur_td = base::GetTrackdir();
 
 		while (cur != GetLastTile() || cur_td != GetLastTrackdir()) {
-			if (!((obj.*func)(cur, cur_td))) return false;
+			if (!(func(cur, cur_td))) return false;
 
 			if (!ft.Follow(cur, cur_td)) break;
 			cur = ft.m_new_tile;
@@ -210,7 +210,15 @@ struct CYapfRailNodeT
 			cur_td = FindFirstTrackdir(ft.m_new_td_bits);
 		}
 
-		return (obj.*func)(cur, cur_td);
+		return func(cur, cur_td);
+	}
+
+	template <class Tbase, class Tpf, class Tfunc>
+	bool IterateTiles(const Train *v, Tpf &yapf, Tbase &obj, bool (Tfunc::*func)(TileIndex, Trackdir)) const
+	{
+		return this->template IterateTiles<Tbase>(v, yapf, [&](TileIndex tile, Trackdir td) -> bool {
+			return (obj.*func)(tile, td);
+		});
 	}
 
 	template <class Tbase, class Tpf>

src/pathfinder/yapf/yapf_rail.cpp

@@ -186,11 +186,38 @@ public:
 		}
 
 		/* Don't bother if the target is reserved. */
-		PBSWaitingPositionRestrictedSignalInfo restricted_signal_info;
-		if (!IsWaitingPositionFree(Yapf().GetVehicle(), m_res_dest, m_res_dest_td, false, &restricted_signal_info)) return false;
+		PBSWaitingPositionRestrictedSignalState restricted_signal_state;
+		restricted_signal_state.defer_test_if_slot_conditional = true;
+		if (!IsWaitingPositionFree(Yapf().GetVehicle(), m_res_dest, m_res_dest_td, false, &restricted_signal_state)) return false;
+
+		static TraceRestrictSlotTemporaryState temporary_slot_state;
+		assert(temporary_slot_state.IsEmpty());
+
+		struct IntermediaryTraceRestrictSignalInfo {
+			const TraceRestrictProgram *prog;
+			TileIndex tile;
+			Trackdir  trackdir;
+		};
+		/* Nodes are iterated in reverse order (from the target), but tiles within the node are iterated in forward order (towards the target).
+		 * intermediary_restricted_signals is in reverse order, (the first signal to evaluate at the end).
+		 */
+		static std::vector<IntermediaryTraceRestrictSignalInfo> intermediary_restricted_signals;
+		intermediary_restricted_signals.clear();
 
 		for (Node *node = m_res_node; node->m_parent != nullptr; node = node->m_parent) {
-			node->IterateTiles(Yapf().GetVehicle(), Yapf(), *this, &CYapfReserveTrack<Types>::ReserveSingleTrack);
+			const size_t intermediary_restricted_signals_current_size = intermediary_restricted_signals.size();
+			node->template IterateTiles<CYapfReserveTrack>(Yapf().GetVehicle(), Yapf(), [&](TileIndex tile, Trackdir td) -> bool {
+				/* Cheapest tests first */
+				if (IsTileType(tile, MP_RAILWAY) && HasSignals(tile) && IsRestrictedSignal(tile) && HasSignalOnTrack(tile, TrackdirToTrack(td))) {
+					const TraceRestrictProgram *prog = GetExistingTraceRestrictProgram(tile, TrackdirToTrack(td));
+					if (prog != nullptr && prog->actions_used_flags & (TRPAUF_WAIT_AT_PBS | TRPAUF_SLOT_ACQUIRE)) {
+						/* Insert at intermediary_restricted_signals_current_size, such that if there are multiple signals for this node, they end up in reverse order */
+						intermediary_restricted_signals.insert(intermediary_restricted_signals.begin() + intermediary_restricted_signals_current_size, { prog, tile, td });
+					}
+				}
+
+				return this->ReserveSingleTrack(tile, td);
+			});
 			if (m_res_fail_tile != INVALID_TILE) {
 				/* Reservation failed, undo. */
 				Node *fail_node = m_res_node;
@@ -201,17 +228,52 @@ public:
 					fail_node->IterateTiles(Yapf().GetVehicle(), Yapf(), *this, &CYapfReserveTrack<Types>::UnreserveSingleTrack);
 				} while (fail_node != node && (fail_node = fail_node->m_parent) != nullptr);
 
+				temporary_slot_state.RevertTemporaryChanges(Yapf().GetVehicle()->index);
 				return false;
 			}
 		}
 
-		if (restricted_signal_info.tile != INVALID_TILE) {
-			const TraceRestrictProgram *prog = GetExistingTraceRestrictProgram(restricted_signal_info.tile, TrackdirToTrack(restricted_signal_info.trackdir));
-			if (prog && prog->actions_used_flags & TRPAUF_PBS_RES_END_SLOT) {
+		auto undo_reservation = [&]() {
+			for (Node *node = m_res_node; node->m_parent != nullptr; node = node->m_parent) {
+				node->IterateTiles(Yapf().GetVehicle(), Yapf(), *this, &CYapfReserveTrack<Types>::UnreserveSingleTrack);
+			}
+			temporary_slot_state.RevertTemporaryChanges(Yapf().GetVehicle()->index);
+		};
+
+		/* Iterate in reverse order */
+		for (auto iter = intermediary_restricted_signals.rbegin(); iter != intermediary_restricted_signals.rend(); ++iter) {
+			extern TileIndex VehiclePosTraceRestrictPreviousSignalCallback(const Train *v, const void *, TraceRestrictPBSEntrySignalAuxField mode);
+
+			TraceRestrictProgramInput input(iter->tile, iter->trackdir, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
+			input.permitted_slot_operations = TRPISP_ACQUIRE_TEMP_STATE;
+			input.slot_temporary_state = &temporary_slot_state;
+			TraceRestrictProgramResult out;
+			iter->prog->Execute(Yapf().GetVehicle(), input, out);
+			if (out.flags & TRPRF_WAIT_AT_PBS) {
+				/* Wait at PBS is set, take this as waiting at the start signal */
+				undo_reservation();
+				return false;
+			}
+		}
+
+		if (restricted_signal_state.deferred_test) {
+			/* The IsWaitingPositionFree restricted signal test was deferred due to possible slot changes during reservation, test it now */
+			if (!IsWaitingPositionFreeTraceRestrictExecute(restricted_signal_state.prog, Yapf().GetVehicle(), restricted_signal_state.tile, restricted_signal_state.trackdir)) {
+				/* Target is reserved, undo reservation */
+				undo_reservation();
+				return false;
+			}
+		}
+
+		temporary_slot_state.ApplyTemporaryChanges(Yapf().GetVehicle()->index);
+
+		if (restricted_signal_state.prog != nullptr) {
+			const TraceRestrictProgram *prog = restricted_signal_state.prog;
+			if (prog != nullptr && prog->actions_used_flags & TRPAUF_PBS_RES_END_SLOT) {
 				extern TileIndex VehiclePosTraceRestrictPreviousSignalCallback(const Train *v, const void *, TraceRestrictPBSEntrySignalAuxField mode);
 
 				TraceRestrictProgramResult out;
-				TraceRestrictProgramInput input(restricted_signal_info.tile, restricted_signal_info.trackdir, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
+				TraceRestrictProgramInput input(restricted_signal_state.tile, restricted_signal_state.trackdir, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
 				input.permitted_slot_operations = TRPISP_PBS_RES_END_ACQUIRE | TRPISP_PBS_RES_END_RELEASE;
 				prog->Execute(Yapf().GetVehicle(), input, out);
 			}

src/pbs.cpp

@@ -1531,6 +1531,20 @@ bool IsSafeWaitingPosition(const Train *v, TileIndex tile, Trackdir trackdir, bo
 	return false;
 }
 
+bool IsWaitingPositionFreeTraceRestrictExecute(const TraceRestrictProgram *prog, const Train *v, TileIndex tile, Trackdir trackdir)
+{
+	if (prog != nullptr && prog->actions_used_flags & TRPAUF_PBS_RES_END_WAIT) {
+		TraceRestrictProgramInput input(tile, trackdir, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
+		input.permitted_slot_operations = TRPISP_PBS_RES_END_ACQ_DRY;
+		TraceRestrictProgramResult out;
+		prog->Execute(v, input, out);
+		if (out.flags & TRPRF_PBS_RES_END_WAIT) {
+			return false;
+		}
+	}
+	return true;
+}
+
 /**
  * Check if a safe position is free.
  *
@@ -1538,9 +1552,10 @@ bool IsSafeWaitingPosition(const Train *v, TileIndex tile, Trackdir trackdir, bo
  * @param tile The tile
  * @param trackdir The trackdir to test
  * @param forbid_90deg Don't allow trains to make 90 degree turns
+ * @param restricted_signal_state Restricted signal state in/out
  * @return True if the position is free
  */
-bool IsWaitingPositionFree(const Train *v, TileIndex tile, Trackdir trackdir, bool forbid_90deg, PBSWaitingPositionRestrictedSignalInfo *restricted_signal_info)
+bool IsWaitingPositionFree(const Train *v, TileIndex tile, Trackdir trackdir, bool forbid_90deg, PBSWaitingPositionRestrictedSignalState *restricted_signal_state)
 {
 	Track     track = TrackdirToTrack(trackdir);
 	TrackBits reserved = GetReservedTrackbits(tile);
@@ -1551,22 +1566,19 @@ bool IsWaitingPositionFree(const Train *v, TileIndex tile, Trackdir trackdir, bo
 	/* Not reserved and depot or not a pbs signal -> free. */
 	if (IsRailDepotTile(tile)) return true;
 
-	auto pbs_res_end_wait_test = [v, restricted_signal_info](TileIndex t, Trackdir td, bool tunnel_bridge) -> bool {
+	auto pbs_res_end_wait_test = [v, restricted_signal_state](TileIndex t, Trackdir td, bool tunnel_bridge) -> bool {
 		if (tunnel_bridge ? IsTunnelBridgeRestrictedSignal(t) : IsRestrictedSignal(t)) {
 			const TraceRestrictProgram *prog = GetExistingTraceRestrictProgram(t, TrackdirToTrack(td));
-			if (restricted_signal_info && prog) {
-				restricted_signal_info->tile = t;
-				restricted_signal_info->trackdir = td;
-			}
-			if (prog && prog->actions_used_flags & TRPAUF_PBS_RES_END_WAIT) {
-				TraceRestrictProgramInput input(t, td, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
-				input.permitted_slot_operations = TRPISP_PBS_RES_END_ACQ_DRY;
-				TraceRestrictProgramResult out;
-				prog->Execute(v, input, out);
-				if (out.flags & TRPRF_PBS_RES_END_WAIT) {
-					return false;
+			if (restricted_signal_state != nullptr && prog != nullptr) {
+				restricted_signal_state->prog = prog;
+				restricted_signal_state->tile = t;
+				restricted_signal_state->trackdir = td;
+				if (restricted_signal_state->defer_test_if_slot_conditional && (prog->actions_used_flags & TRPAUF_SLOT_CONDITIONALS) && (prog->actions_used_flags & TRPAUF_PBS_RES_END_WAIT)) {
+					restricted_signal_state->deferred_test = true;
+					return true;
 				}
 			}
+			return IsWaitingPositionFreeTraceRestrictExecute(prog, v, t, td);
 		}
 		return true;
 	};

src/pbs.h

@@ -45,11 +45,6 @@ struct PBSTileInfo {
 	PBSTileInfo(TileIndex _t, Trackdir _td, bool _okay) : tile(_t), trackdir(_td), okay(_okay) {}
 };
 
-struct PBSWaitingPositionRestrictedSignalInfo {
-	TileIndex tile = INVALID_TILE;
-	Trackdir  trackdir = INVALID_TRACKDIR;
-};
-
 enum TrainReservationLookAheadItemType : byte {
 	TRLIT_STATION                = 0,     ///< Station/waypoint
 	TRLIT_REVERSE                = 1,     ///< Reverse behind signal
@@ -179,7 +174,19 @@ void SetTrainReservationLookaheadEnd(Train *v);
 void FillTrainReservationLookAhead(Train *v);
 bool TrainReservationPassesThroughTile(const Train *v, TileIndex search_tile);
 bool IsSafeWaitingPosition(const Train *v, TileIndex tile, Trackdir trackdir, bool include_line_end, bool forbid_90deg = false);
-bool IsWaitingPositionFree(const Train *v, TileIndex tile, Trackdir trackdir, bool forbid_90deg = false, PBSWaitingPositionRestrictedSignalInfo *restricted_signal_info = nullptr);
+
+struct TraceRestrictProgram;
+
+struct PBSWaitingPositionRestrictedSignalState {
+	const TraceRestrictProgram *prog = nullptr;
+	TileIndex tile = INVALID_TILE;
+	Trackdir  trackdir = INVALID_TRACKDIR;
+	bool defer_test_if_slot_conditional = false;
+	bool deferred_test = false;
+};
+
+bool IsWaitingPositionFree(const Train *v, TileIndex tile, Trackdir trackdir, bool forbid_90deg = false, PBSWaitingPositionRestrictedSignalState *restricted_signal_state = nullptr);
+bool IsWaitingPositionFreeTraceRestrictExecute(const TraceRestrictProgram *prog, const Train *v, TileIndex tile, Trackdir trackdir);
 
 Train *GetTrainForReservation(TileIndex tile, Track track);
 CommandCost CheckTrainReservationPreventsTrackModification(TileIndex tile, Track track);

src/train_cmd.cpp

@@ -3844,6 +3844,7 @@ static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &ori
 		/* Station, depot or waypoint are a possible target. */
 		bool target_seen = ft.m_is_station || (IsTileType(ft.m_new_tile, MP_RAILWAY) && !IsPlainRail(ft.m_new_tile));
 		if (target_seen || KillFirstBit(ft.m_new_td_bits) != TRACKDIR_BIT_NONE) {
+			target_seen_path:
 			/* Choice found or possible target encountered.
 			 * On finding a possible target, we need to stop and let the pathfinder handle the
 			 * remaining path. This is because we don't know if this target is in one of our
@@ -3866,15 +3867,15 @@ static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &ori
 		cur_td = FindFirstTrackdir(ft.m_new_td_bits);
 
 		if (IsSafeWaitingPosition(v, tile, cur_td, true, _settings_game.pf.forbid_90_deg)) {
-			PBSWaitingPositionRestrictedSignalInfo restricted_signal_info;
-			bool wp_free = IsWaitingPositionFree(v, tile, cur_td, _settings_game.pf.forbid_90_deg, &restricted_signal_info);
+			PBSWaitingPositionRestrictedSignalState restricted_signal_state;
+			bool wp_free = IsWaitingPositionFree(v, tile, cur_td, _settings_game.pf.forbid_90_deg, &restricted_signal_state);
 			if (!(wp_free && TryReserveRailTrackdir(v, tile, cur_td))) break;
 			/* Safe position is all good, path valid and okay. */
-			if (restricted_signal_info.tile != INVALID_TILE) {
-				const TraceRestrictProgram *prog = GetExistingTraceRestrictProgram(restricted_signal_info.tile, TrackdirToTrack(restricted_signal_info.trackdir));
-				if (prog && prog->actions_used_flags & TRPAUF_PBS_RES_END_SLOT) {
+			if (restricted_signal_state.prog != nullptr) {
+				const TraceRestrictProgram *prog = restricted_signal_state.prog;
+				if (prog != nullptr && (prog->actions_used_flags & TRPAUF_PBS_RES_END_SLOT)) {
 					TraceRestrictProgramResult out;
-					TraceRestrictProgramInput input(restricted_signal_info.tile, restricted_signal_info.trackdir, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
+					TraceRestrictProgramInput input(restricted_signal_state.tile, restricted_signal_state.trackdir, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
 					input.permitted_slot_operations = TRPISP_PBS_RES_END_ACQUIRE | TRPISP_PBS_RES_END_RELEASE;
 					prog->Execute(v, input, out);
 				}
@@ -3882,6 +3883,16 @@ static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &ori
 			return PBSTileInfo(tile, cur_td, true);
 		}
 
+		if (IsTileType(tile, MP_RAILWAY) && HasSignals(tile) && IsRestrictedSignal(tile) && HasSignalOnTrack(tile, TrackdirToTrack(cur_td))) {
+			const TraceRestrictProgram *prog = GetExistingTraceRestrictProgram(tile, TrackdirToTrack(cur_td));
+			if (prog != nullptr && prog->actions_used_flags & (TRPAUF_WAIT_AT_PBS | TRPAUF_SLOT_ACQUIRE)) {
+				/* The pathfinder must deal with this, because temporary slot states can't be nested.
+				 * See target_seen path above.
+				 */
+				goto target_seen_path;
+			}
+		}
+
 		if (!TryReserveRailTrackdir(v, tile, cur_td)) break;
 	}