Tracerestrict: Allow nesting temporary slot states

Use a temporary state stack instead of passing the state directly
in TraceRestrictProgramInput

No longer exit early from ExtendTrainReservation to handle slots, this
fixes reservation failures when starting from a signal with a slot acquire

More thorough fix for 9e8a4314

src/pathfinder/yapf/yapf_rail.cpp

@@ -190,8 +190,9 @@ public:
 		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;
 
+		/* The temporary slot state only needs to be pushed to the stack (i.e. activated) on first use */
 		static TraceRestrictSlotTemporaryState temporary_slot_state;
-		assert(temporary_slot_state.IsEmpty());
+		assert(temporary_slot_state.IsEmpty() && !temporary_slot_state.IsActive());
 
 		struct IntermediaryTraceRestrictSignalInfo {
 			const TraceRestrictProgram *prog;
@@ -228,7 +229,7 @@ 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);
+				if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackRevertTemporaryChanges(Yapf().GetVehicle()->index);
 				return false;
 			}
 		}
@@ -237,16 +238,20 @@ public:
 			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);
+			if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackRevertTemporaryChanges(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);
 
+			if (!temporary_slot_state.IsActive()) {
+				/* The temporary slot state needs to be be pushed because permission to use it is granted by TRPISP_ACQUIRE_TEMP_STATE */
+				temporary_slot_state.PushToChangeStack();
+			}
+
 			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) {
@@ -266,7 +271,7 @@ public:
 		}
 
 		/* This must be done before calling TraceRestrictExecuteResEndSlot */
-		temporary_slot_state.ApplyTemporaryChanges(Yapf().GetVehicle());
+		TraceRestrictSlotTemporaryState::ClearChangeStackApplyAllTemporaryChanges(Yapf().GetVehicle());
 
 		restricted_signal_state.TraceRestrictExecuteResEndSlot(Yapf().GetVehicle());
 

src/tracerestrict.cpp

@@ -732,7 +732,7 @@ void TraceRestrictProgram::Execute(const Train* v, const TraceRestrictProgramInp
 								if (input.permitted_slot_operations & TRPISP_ACQUIRE) {
 									if (!slot->Occupy(v)) out.flags |= TRPRF_WAIT_AT_PBS;
 								} else if (input.permitted_slot_operations & TRPISP_ACQUIRE_TEMP_STATE) {
-									if (!slot->OccupyUsingTemporaryState(v->index, input.slot_temporary_state)) out.flags |= TRPRF_WAIT_AT_PBS;
+									if (!slot->OccupyUsingTemporaryState(v->index, TraceRestrictSlotTemporaryState::GetCurrent())) out.flags |= TRPRF_WAIT_AT_PBS;
 								}
 								break;
 
@@ -740,7 +740,7 @@ void TraceRestrictProgram::Execute(const Train* v, const TraceRestrictProgramInp
 								if (input.permitted_slot_operations & TRPISP_ACQUIRE) {
 									slot->Occupy(v);
 								} else if (input.permitted_slot_operations & TRPISP_ACQUIRE_TEMP_STATE) {
-									slot->OccupyUsingTemporaryState(v->index, input.slot_temporary_state);
+									slot->OccupyUsingTemporaryState(v->index, TraceRestrictSlotTemporaryState::GetCurrent());
 								}
 								break;
 
@@ -748,7 +748,7 @@ void TraceRestrictProgram::Execute(const Train* v, const TraceRestrictProgramInp
 								if (input.permitted_slot_operations & TRPISP_ACQUIRE) {
 									slot->Vacate(v);
 								} else if (input.permitted_slot_operations & TRPISP_ACQUIRE_TEMP_STATE) {
-									slot->VacateUsingTemporaryState(v->index, input.slot_temporary_state);
+									slot->VacateUsingTemporaryState(v->index, TraceRestrictSlotTemporaryState::GetCurrent());
 								}
 								break;
 
@@ -2716,6 +2716,8 @@ void TraceRestrictSlot::PreCleanPool()
 	slot_vehicle_index.clear();
 }
 
+std::vector<TraceRestrictSlotTemporaryState *> TraceRestrictSlotTemporaryState::change_stack;
+
 /** Revert any temporary changes */
 void TraceRestrictSlotTemporaryState::RevertTemporaryChanges(VehicleID veh)
 {
@@ -2754,6 +2756,38 @@ void TraceRestrictSlotTemporaryState::ApplyTemporaryChanges(const Vehicle *v)
 	this->veh_temporarily_removed.clear();
 }
 
+/** Apply any temporary changes to a parent temporary state */
+void TraceRestrictSlotTemporaryState::ApplyTemporaryChangesToParent(VehicleID veh, TraceRestrictSlotTemporaryState *parent)
+{
+	for (TraceRestrictSlotID id : this->veh_temporarily_added) {
+		if (find_index(parent->veh_temporarily_removed, id) < 0) {
+			include(parent->veh_temporarily_added, id);
+		}
+	}
+	for (TraceRestrictSlotID id : this->veh_temporarily_removed) {
+		if (find_index(parent->veh_temporarily_added, id) < 0) {
+			include(parent->veh_temporarily_removed, id);
+		}
+	}
+
+	this->veh_temporarily_added.clear();
+	this->veh_temporarily_removed.clear();
+}
+
+/** Pop from change stack and apply any temporary changes (to the parent temporary state if present) */
+void TraceRestrictSlotTemporaryState::PopFromChangeStackApplyTemporaryChanges(const Vehicle *v)
+{
+	assert(this->change_stack.back() == this);
+	this->change_stack.pop_back();
+	this->is_active = false;
+
+	if (this->change_stack.empty()) {
+		this->ApplyTemporaryChanges(v);
+	} else {
+		this->ApplyTemporaryChangesToParent(v->index, this->change_stack.back());
+	}
+}
+
 /** Remove vehicle ID from all slot occupants */
 void TraceRestrictRemoveVehicleFromAllSlots(VehicleID vehicle_id)
 {

src/tracerestrict.h

@@ -489,7 +489,7 @@ enum TraceRestrictProgramInputSlotPermissions : uint8_t {
 	TRPISP_RELEASE_FRONT          = 1 << 2,  ///< Slot release (front) is permitted
 	TRPISP_PBS_RES_END_ACQUIRE    = 1 << 3,  ///< Slot acquire/release (PBS reservations ending at this signal) is permitted
 	TRPISP_PBS_RES_END_ACQ_DRY    = 1 << 4,  ///< Dry-run slot acquire/release (PBS reservations ending at this signal) is permitted
-	TRPISP_ACQUIRE_TEMP_STATE     = 1 << 5,  ///< Slot acquire/release is permitted, using temporary state, TraceRestrictProgramInput::slot_temporary_state must be set
+	TRPISP_ACQUIRE_TEMP_STATE     = 1 << 5,  ///< Slot acquire/release is permitted, using temporary state, TraceRestrictSlotTemporaryState::change_stack must be non-empty
 	TRPISP_CHANGE_COUNTER         = 1 << 6,  ///< Change counter value is permitted
 };
 DECLARE_ENUM_AS_BIT_SET(TraceRestrictProgramInputSlotPermissions)
@@ -507,8 +507,44 @@ struct TraceRestrictSlotTemporaryState {
 	std::vector<TraceRestrictSlotID> veh_temporarily_added;
 	std::vector<TraceRestrictSlotID> veh_temporarily_removed;
 
-	void RevertTemporaryChanges(VehicleID veh);
+private:
+	bool is_active = false;
+
+	static std::vector<TraceRestrictSlotTemporaryState *> change_stack;
+
 	void ApplyTemporaryChanges(const Vehicle *v);
+	void ApplyTemporaryChangesToParent(VehicleID veh, TraceRestrictSlotTemporaryState *parent);
+
+public:
+	static TraceRestrictSlotTemporaryState *GetCurrent() { return change_stack.back(); }
+
+	static void ClearChangeStackApplyAllTemporaryChanges(const Vehicle *v)
+	{
+		while (!change_stack.empty()) {
+			change_stack.back()->PopFromChangeStackApplyTemporaryChanges(v);
+		}
+	}
+
+	bool IsActive() const { return this->is_active; }
+
+	void RevertTemporaryChanges(VehicleID veh);
+
+	void PushToChangeStack()
+	{
+		assert(!this->is_active);
+		this->is_active = true;
+		this->change_stack.push_back(this);
+	}
+
+	void PopFromChangeStackRevertTemporaryChanges(VehicleID veh)
+	{
+		assert(this->change_stack.back() == this);
+		this->change_stack.pop_back();
+		this->RevertTemporaryChanges(veh);
+		this->is_active = false;
+	}
+
+	void PopFromChangeStackApplyTemporaryChanges(const Vehicle *v);
 
 	bool IsEmpty() const
 	{
@@ -528,7 +564,6 @@ struct TraceRestrictProgramInput {
 	TraceRestrictProgramInputSlotPermissions permitted_slot_operations; ///< Permitted slot operations
 	PreviousSignalProc *previous_signal_callback; ///< Callback to retrieve tile and direction of previous signal, may be nullptr
 	const void *previous_signal_ptr;              ///< Opaque pointer suitable to be passed to previous_signal_callback
-	TraceRestrictSlotTemporaryState *slot_temporary_state = nullptr; ///< Slot temporary state, must be set when permitted_slot_operations includes TRPISP_ACQUIRE_TEMP_STATE
 
 	TraceRestrictProgramInput(TileIndex tile_, Trackdir trackdir_, PreviousSignalProc *previous_signal_callback_, const void *previous_signal_ptr_)
 			: tile(tile_), trackdir(trackdir_), input_flags(TRPIF_NONE), permitted_slot_operations(TRPISP_NONE),

src/train_cmd.cpp

@@ -3822,9 +3822,10 @@ static Track DoTrainPathfind(const Train *v, TileIndex tile, DiagDirection enter
  * @param origin The tile from which the reservation have to be extended
  * @param new_tracks [out] Tracks to choose from when encountering a choice
  * @param enterdir [out] The direction from which the choice tile is to be entered
+ * @param temporary_slot_state The temporary slot to use (will be activated/deactivated as necessary if it isn't already)
  * @return INVALID_TILE indicates that the reservation failed.
  */
-static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &origin, TrackBits *new_tracks, DiagDirection *enterdir)
+static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &origin, TrackBits *new_tracks, DiagDirection *enterdir, TraceRestrictSlotTemporaryState &temporary_slot_state)
 {
 	CFollowTrackRail ft(v);
 
@@ -3844,7 +3845,6 @@ 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
@@ -3878,10 +3878,20 @@ static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &ori
 		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 (!temporary_slot_state.IsActive()) {
+					/* The temporary slot state needs to be be pushed because permission to use it is granted by TRPISP_ACQUIRE_TEMP_STATE */
+					temporary_slot_state.PushToChangeStack();
+				}
+
+				TraceRestrictProgramInput input(tile, cur_td, &VehiclePosTraceRestrictPreviousSignalCallback, nullptr);
+				input.permitted_slot_operations = TRPISP_ACQUIRE_TEMP_STATE;
+				TraceRestrictProgramResult out;
+				prog->Execute(v, input, out);
+				if (out.flags & TRPRF_WAIT_AT_PBS) {
+					/* Wait at PBS is set, take this as waiting at the start signal, handle as a reservation failure */
+					break;
+				}
 			}
 		}
 
@@ -3913,6 +3923,8 @@ static PBSTileInfo ExtendTrainReservation(const Train *v, const PBSTileInfo &ori
 		UnreserveRailTrackdir(tile, cur_td);
 	}
 
+	if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackRevertTemporaryChanges(v->index);
+
 	/* Path invalid. */
 	return PBSTileInfo();
 }
@@ -4415,11 +4427,19 @@ static Track ChooseTrainTrack(Train *v, TileIndex tile, DiagDirection enterdir,
 		ClearLookAheadIfInvalid(v);
 	}
 
+	/* The temporary slot state only needs to be pushed to the stack (i.e. activated) on first use */
+	TraceRestrictSlotTemporaryState temporary_slot_state;
+
+	/* All exit paths except success should revert the temporary slot state if required */
+	auto slot_state_guard = scope_guard([&]() {
+		if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackRevertTemporaryChanges(v->index);
+	});
+
 	PBSTileInfo   origin = FollowTrainReservation(v, nullptr, FTRF_OKAY_UNUSED);
 	PBSTileInfo   res_dest(tile, INVALID_TRACKDIR, false);
 	DiagDirection dest_enterdir = enterdir;
 	if (do_track_reservation) {
-		res_dest = ExtendTrainReservation(v, origin, &tracks, &dest_enterdir);
+		res_dest = ExtendTrainReservation(v, origin, &tracks, &dest_enterdir, temporary_slot_state);
 		if (res_dest.tile == INVALID_TILE) {
 			/* Reservation failed? */
 			if (mark_stuck) MarkTrainAsStuck(v);
@@ -4427,6 +4447,7 @@ static Track ChooseTrainTrack(Train *v, TileIndex tile, DiagDirection enterdir,
 			return FindFirstTrack(tracks);
 		}
 		if (res_dest.okay) {
+			if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackApplyTemporaryChanges(v);
 			bool long_reserve = (CheckLongReservePbsTunnelBridgeOnTrackdir(v, res_dest.tile, res_dest.trackdir) != INVALID_TILE);
 			if (!long_reserve) {
 				CFollowTrackRail ft(v);
@@ -4444,8 +4465,6 @@ static Track ChooseTrainTrack(Train *v, TileIndex tile, DiagDirection enterdir,
 				if (_settings_game.vehicle.train_braking_model == TBM_REALISTIC) FillTrainReservationLookAhead(v);
 				return best_track;
 			}
-		} else if (res_dest.tile == tile) {
-			if (changed_signal != INVALID_TRACKDIR) SetSignalStateByTrackdir(tile, changed_signal, SIGNAL_STATE_RED);
 		}
 
 		/* Check if the train needs service here, so it has a chance to always find a depot.
@@ -4493,6 +4512,7 @@ static Track ChooseTrainTrack(Train *v, TileIndex tile, DiagDirection enterdir,
 		/* Try to find any safe destination. */
 		PBSTileInfo path_end = FollowTrainReservation(v, nullptr, FTRF_OKAY_UNUSED);
 		if (TryReserveSafeTrack(v, path_end.tile, path_end.trackdir, false)) {
+			if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackApplyTemporaryChanges(v);
 			TrackBits res = GetReservedTrackbits(tile) & DiagdirReachesTracks(enterdir);
 			best_track = FindFirstTrack(res);
 			if (!HasBit(lookahead_state.flags, CTTLASF_NO_RES_VEH_TILE)) TryReserveRailTrack(v->tile, TrackdirToTrack(v->GetVehicleTrackdir()));
@@ -4550,6 +4570,7 @@ static Track ChooseTrainTrack(Train *v, TileIndex tile, DiagDirection enterdir,
 				if (mark_stuck) MarkTrainAsStuck(v);
 				got_reservation = false;
 				changed_signal = INVALID_TRACKDIR;
+				if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackRevertTemporaryChanges(v->index);
 				break;
 			}
 		}
@@ -4559,11 +4580,13 @@ static Track ChooseTrainTrack(Train *v, TileIndex tile, DiagDirection enterdir,
 			if (mark_stuck) MarkTrainAsStuck(v);
 			got_reservation = false;
 			changed_signal = INVALID_TRACKDIR;
+			if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackRevertTemporaryChanges(v->index);
 		}
 		break;
 	}
 
 	if (got_reservation) {
+		if (temporary_slot_state.IsActive()) temporary_slot_state.PopFromChangeStackApplyTemporaryChanges(v);
 		if (v->current_order.IsBaseStationOrder() && HasStationTileRail(res_dest.tile) && v->current_order.GetDestination() == GetStationIndex(res_dest.tile)) {
 			if (v->current_order.ShouldStopAtStation(v, v->current_order.GetDestination(), v->current_order.IsType(OT_GOTO_WAYPOINT))) {
 				v->last_station_visited = v->current_order.GetDestination();