Tracerestrict: Add button to highlight all signals using program

src/lang/english.txt

@@ -3452,6 +3452,7 @@ STR_TRACE_RESTRICT_SHARE_TOOLTIP                                :{BLACK}Share pr
 STR_TRACE_RESTRICT_UNSHARE_TOOLTIP                              :{BLACK}Stop sharing program with other signals, create a copy of the program
 STR_TRACE_RESTRICT_SIGNAL_GUI_TOOLTIP                           :{BLACK}Routefinding restriction
 STR_TRACE_RESTRICT_INSTRUCTION_LIST_TOOLTIP                     :{BLACK}Click an instruction to select it{}Ctrl+Click to scroll to the instruction's target (if any)
+STR_TRACE_RESTRICT_HIGHLIGHT_TOOLTIP                            :{BLACK}Toggle highlighting all signals sharing this program
 STR_TRACE_RESTRICT_ERROR_CAN_T_INSERT_ITEM                      :{WHITE}Can't insert instruction
 STR_TRACE_RESTRICT_ERROR_CAN_T_MODIFY_ITEM                      :{WHITE}Can't modify instruction
 STR_TRACE_RESTRICT_ERROR_CAN_T_REMOVE_ITEM                      :{WHITE}Can't remove instruction

src/tracerestrict.cpp

@@ -871,6 +871,11 @@ void TraceRestrictProgram::DecrementRefCount(TraceRestrictRefId ref_id) {
 		free(ptr);
 	}
 	if (this->refcount == 0) {
+		extern const TraceRestrictProgram *_viewport_highlight_tracerestrict_program;
+		if (_viewport_highlight_tracerestrict_program == this) {
+			_viewport_highlight_tracerestrict_program = nullptr;
+			InvalidateWindowClassesData(WC_TRACE_RESTRICT);
+		}
 		delete this;
 	}
 }

src/tracerestrict_gui.cpp

@@ -51,6 +51,7 @@
 /** Widget IDs */
 enum TraceRestrictWindowWidgets {
 	TR_WIDGET_CAPTION,
+	TR_WIDGET_HIGHLIGHT,
 	TR_WIDGET_INSTRUCTION_LIST,
 	TR_WIDGET_SCROLLBAR,
 
@@ -1632,6 +1633,17 @@ public:
 		this->ReloadProgramme();
 	}
 
+	~TraceRestrictWindow()
+	{
+		extern const TraceRestrictProgram *_viewport_highlight_tracerestrict_program;
+		if (_viewport_highlight_tracerestrict_program != nullptr) {
+			const TraceRestrictProgram *prog = this->GetProgram();
+			if (prog != nullptr && prog == _viewport_highlight_tracerestrict_program) {
+				SetViewportCatchmentTraceRestrictProgram(prog, false);
+			}
+		}
+	}
+
 	virtual void OnClick(Point pt, int widget, int click_count) override
 	{
 		switch (widget) {
@@ -1983,6 +1995,15 @@ public:
 				TraceRestrictProgMgmtDoCommandP(tile, track, TRDCT_PROG_UNSHARE, STR_TRACE_RESTRICT_ERROR_CAN_T_UNSHARE_PROGRAM);
 				break;
 			}
+
+			case TR_WIDGET_HIGHLIGHT: {
+				const TraceRestrictProgram *prog = this->GetProgram();
+				if (prog != nullptr) {
+					extern const TraceRestrictProgram *_viewport_highlight_tracerestrict_program;
+					SetViewportCatchmentTraceRestrictProgram(prog, _viewport_highlight_tracerestrict_program != prog);
+				}
+				break;
+			}
 		}
 	}
 
@@ -2670,6 +2691,10 @@ private:
 		this->GetWidget<NWidgetCore>(TR_WIDGET_CAPTION)->widget_data =
 				(prog && prog->refcount > 1) ? STR_TRACE_RESTRICT_CAPTION_SHARED : STR_TRACE_RESTRICT_CAPTION;
 
+		this->SetWidgetDisabledState(TR_WIDGET_HIGHLIGHT, prog == nullptr);
+		extern const TraceRestrictProgram *_viewport_highlight_tracerestrict_program;
+		this->SetWidgetLoweredState(TR_WIDGET_HIGHLIGHT, prog != nullptr && _viewport_highlight_tracerestrict_program == prog);
+
 		auto left_aux_guard = scope_guard([&]() {
 			if (this->current_left_aux_plane != left_aux_sel->shown_plane) {
 				this->current_left_aux_plane = left_aux_sel->shown_plane;
@@ -3137,6 +3162,7 @@ static const NWidgetPart _nested_program_widgets[] = {
 	NWidget(NWID_HORIZONTAL),
 		NWidget(WWT_CLOSEBOX, COLOUR_GREY),
 		NWidget(WWT_CAPTION, COLOUR_GREY, TR_WIDGET_CAPTION), SetDataTip(STR_TRACE_RESTRICT_CAPTION, STR_TOOLTIP_WINDOW_TITLE_DRAG_THIS),
+		NWidget(WWT_IMGBTN, COLOUR_GREY, TR_WIDGET_HIGHLIGHT), SetMinimalSize(12, 12), SetDataTip(SPR_SHARED_ORDERS_ICON, STR_TRACE_RESTRICT_HIGHLIGHT_TOOLTIP),
 		NWidget(WWT_SHADEBOX, COLOUR_GREY),
 		NWidget(WWT_STICKYBOX, COLOUR_GREY),
 	EndContainer(),

src/viewport.cpp

@@ -112,6 +112,7 @@
 #include "object_map.h"
 #include "newgrf_object.h"
 #include "infrastructure_func.h"
+#include "tracerestrict.h"
 
 #include <map>
 #include <vector>
@@ -1391,6 +1392,7 @@ enum TileHighlightType {
 
 const Station *_viewport_highlight_station; ///< Currently selected station for coverage area highlight
 const Town *_viewport_highlight_town;       ///< Currently selected town for coverage area highlight
+const TraceRestrictProgram *_viewport_highlight_tracerestrict_program; ///< Currently selected tracerestrict program for highlight
 
 /**
  * Get tile highlight type of coverage area for a given tile.
@@ -1422,6 +1424,13 @@ static TileHighlightType GetTileHighlightType(TileIndex t)
 		}
 	}
 
+	if (_viewport_highlight_tracerestrict_program != nullptr) {
+		const TraceRestrictRefId *refs = _viewport_highlight_tracerestrict_program->GetRefIdsPtr();
+		for (uint i = 0; i < _viewport_highlight_tracerestrict_program->refcount; i++) {
+			if (GetTraceRestrictRefIdTileIndex(refs[i]) == t) return THT_LIGHT_BLUE;
+		}
+	}
+
 	return THT_NONE;
 }
 
@@ -6208,6 +6217,8 @@ void SetViewportCatchmentStation(const Station *st, bool sel)
 		MarkCatchmentTilesDirty();
 		_viewport_highlight_station = st;
 		_viewport_highlight_town = nullptr;
+		if (_viewport_highlight_tracerestrict_program != nullptr) InvalidateWindowClassesData(WC_TRACE_RESTRICT);
+		_viewport_highlight_tracerestrict_program = nullptr;
 		MarkCatchmentTilesDirty();
 	} else if (!sel && _viewport_highlight_station == st) {
 		MarkCatchmentTilesDirty();
@@ -6229,6 +6240,8 @@ void SetViewportCatchmentTown(const Town *t, bool sel)
 	if (sel && _viewport_highlight_town != t) {
 		_viewport_highlight_station = nullptr;
 		_viewport_highlight_town = t;
+		if (_viewport_highlight_tracerestrict_program != nullptr) InvalidateWindowClassesData(WC_TRACE_RESTRICT);
+		_viewport_highlight_tracerestrict_program = nullptr;
 		MarkWholeNonMapViewportsDirty();
 	} else if (!sel && _viewport_highlight_town == t) {
 		_viewport_highlight_town = nullptr;
@@ -6237,6 +6250,23 @@ void SetViewportCatchmentTown(const Town *t, bool sel)
 	if (_viewport_highlight_town != nullptr) SetWindowDirty(WC_TOWN_VIEW, _viewport_highlight_town->index);
 }
 
+void SetViewportCatchmentTraceRestrictProgram(const TraceRestrictProgram *prog, bool sel)
+{
+	if (_viewport_highlight_town != nullptr) SetWindowDirty(WC_TOWN_VIEW, _viewport_highlight_town->index);
+	if (_viewport_highlight_station != nullptr) SetWindowDirty(WC_STATION_VIEW, _viewport_highlight_station->index);
+	if (sel && _viewport_highlight_tracerestrict_program != prog) {
+		_viewport_highlight_station = nullptr;
+		_viewport_highlight_town = nullptr;
+		_viewport_highlight_tracerestrict_program = prog;
+		InvalidateWindowClassesData(WC_TRACE_RESTRICT);
+		MarkWholeNonMapViewportsDirty();
+	} else if (!sel && _viewport_highlight_tracerestrict_program == prog) {
+		_viewport_highlight_tracerestrict_program = nullptr;
+		InvalidateWindowClassesData(WC_TRACE_RESTRICT);
+		MarkWholeNonMapViewportsDirty();
+	}
+}
+
 int GetSlopeTreeBrightnessAdjust(Slope slope)
 {
 	switch (slope) {

src/viewport_func.h

@@ -136,8 +136,10 @@ void DrawTileSelectionRect(const TileInfo *ti, PaletteID pal);
 void DrawSelectionSprite(SpriteID image, PaletteID pal, const TileInfo *ti, int z_offset, FoundationPart foundation_part, const SubSprite *sub = nullptr);
 
 struct Town;
+struct TraceRestrictProgram;
 void SetViewportCatchmentStation(const Station *st, bool sel);
 void SetViewportCatchmentTown(const Town *t, bool sel);
+void SetViewportCatchmentTraceRestrictProgram(const TraceRestrictProgram *prog, bool sel);
 
 void MarkBridgeDirty(TileIndex begin, TileIndex end, DiagDirection direction, uint bridge_height, ViewportMarkDirtyFlags flags = VMDF_NONE);
 void MarkBridgeDirty(TileIndex tile, ViewportMarkDirtyFlags flags = VMDF_NONE);