From 047a66030b3e51aa27c60e6d2f29a2799c1e282d Mon Sep 17 00:00:00 2001
From: lethosor <lethosor@gmail.com>
Date: Wed, 7 Jun 2017 20:51:40 -0400
Subject: [PATCH] getAnyUnit: add support for viewscreen_unitst,
 viewscreen_textviewerst

---
 library/modules/Gui.cpp | 87 ++++++++++++++++++++++++-----------------
 1 file changed, 51 insertions(+), 36 deletions(-)

diff --git a/library/modules/Gui.cpp b/library/modules/Gui.cpp
index 941f652de..ccbad8f00 100644
--- a/library/modules/Gui.cpp
+++ b/library/modules/Gui.cpp
@@ -46,55 +46,58 @@ using namespace DFHack;
 #include "modules/Maps.h"
 
 #include "DataDefs.h"
-#include "df/world.h"
+
+#include "df/announcement_flags.h"
+#include "df/announcements.h"
+#include "df/assign_trade_status.h"
+#include "df/building_civzonest.h"
+#include "df/building_furnacest.h"
+#include "df/building_trapst.h"
+#include "df/building_workshopst.h"
+#include "df/game_mode.h"
+#include "df/general_ref.h"
 #include "df/global_objects.h"
-#include "df/viewscreen_dwarfmodest.h"
-#include "df/viewscreen_dungeonmodest.h"
+#include "df/graphic.h"
+#include "df/interfacest.h"
+#include "df/job.h"
+#include "df/layer_object_listst.h"
+#include "df/occupation.h"
+#include "df/plant.h"
+#include "df/popup_message.h"
+#include "df/report.h"
+#include "df/route_stockpile_link.h"
+#include "df/stop_depart_condition.h"
+#include "df/ui_advmode.h"
+#include "df/ui_build_selector.h"
+#include "df/ui_look_list.h"
+#include "df/ui_sidebar_menus.h"
+#include "df/ui_unit_view_mode.h"
+#include "df/unit.h"
+#include "df/unit_inventory_item.h"
+#include "df/viewscreen_buildinglistst.h"
 #include "df/viewscreen_dungeon_monsterstatusst.h"
-#include "df/viewscreen_jobst.h"
+#include "df/viewscreen_dungeonmodest.h"
+#include "df/viewscreen_dwarfmodest.h"
+#include "df/viewscreen_itemst.h"
 #include "df/viewscreen_joblistst.h"
 #include "df/viewscreen_jobmanagementst.h"
-#include "df/viewscreen_unitlistst.h"
-#include "df/viewscreen_buildinglistst.h"
-#include "df/viewscreen_itemst.h"
+#include "df/viewscreen_jobst.h"
 #include "df/viewscreen_layer.h"
-#include "df/viewscreen_layer_noblelistst.h"
-#include "df/viewscreen_layer_overall_healthst.h"
 #include "df/viewscreen_layer_assigntradest.h"
 #include "df/viewscreen_layer_militaryst.h"
+#include "df/viewscreen_layer_noblelistst.h"
+#include "df/viewscreen_layer_overall_healthst.h"
 #include "df/viewscreen_layer_stockpilest.h"
 #include "df/viewscreen_locationsst.h"
 #include "df/viewscreen_petst.h"
-#include "df/viewscreen_tradegoodsst.h"
 #include "df/viewscreen_storesst.h"
+#include "df/viewscreen_textviewerst.h"
+#include "df/viewscreen_tradegoodsst.h"
+#include "df/viewscreen_unitlistst.h"
+#include "df/viewscreen_unitst.h"
 #include "df/viewscreen_workquota_conditionst.h"
 #include "df/viewscreen_workshop_profilest.h"
-#include "df/ui_unit_view_mode.h"
-#include "df/ui_sidebar_menus.h"
-#include "df/ui_look_list.h"
-#include "df/ui_advmode.h"
-#include "df/job.h"
-#include "df/ui_build_selector.h"
-#include "df/building_workshopst.h"
-#include "df/building_furnacest.h"
-#include "df/building_trapst.h"
-#include "df/building_civzonest.h"
-#include "df/general_ref.h"
-#include "df/unit_inventory_item.h"
-#include "df/report.h"
-#include "df/popup_message.h"
-#include "df/interfacest.h"
-#include "df/graphic.h"
-#include "df/layer_object_listst.h"
-#include "df/assign_trade_status.h"
-#include "df/announcement_flags.h"
-#include "df/announcements.h"
-#include "df/stop_depart_condition.h"
-#include "df/route_stockpile_link.h"
-#include "df/game_mode.h"
-#include "df/unit.h"
-#include "df/occupation.h"
-#include "df/plant.h"
+#include "df/world.h"
 
 using namespace df::enums;
 using df::global::gview;
@@ -815,6 +818,11 @@ df::unit *Gui::getAnyUnit(df::viewscreen *top)
     using df::global::ui_look_list;
     using df::global::ui_selected_unit;
 
+    if (VIRTUAL_CAST_VAR(screen, df::viewscreen_unitst, top))
+    {
+        return screen->unit;
+    }
+
     if (VIRTUAL_CAST_VAR(screen, df::viewscreen_joblistst, top))
     {
         if (auto unit = vector_get(screen->units, screen->cursor_pos))
@@ -910,6 +918,13 @@ df::unit *Gui::getAnyUnit(df::viewscreen *top)
         return NULL;
     }
 
+    if (VIRTUAL_CAST_VAR(screen, df::viewscreen_textviewerst, top))
+    {
+        if (screen->parent)
+            return getAnyUnit(screen->parent);
+        return NULL;
+    }
+
     if (auto dfscreen = dfhack_viewscreen::try_cast(top))
         return dfscreen->getSelectedUnit();