From 9489c6ed1afd33790adc2e0b4e3c0152da8203e0 Mon Sep 17 00:00:00 2001
From: Alexander Gavrilov <angavrilov@gmail.com>
Date: Thu, 26 Apr 2012 11:05:35 +0400
Subject: [PATCH] Add support for a few more viewscreens to
 Units::getSelectedUnit.

---
 library/modules/Gui.cpp | 67 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/library/modules/Gui.cpp b/library/modules/Gui.cpp
index 6841421b6..2d9402fca 100644
--- a/library/modules/Gui.cpp
+++ b/library/modules/Gui.cpp
@@ -50,6 +50,11 @@ using namespace DFHack;
 #include "df/viewscreen_joblistst.h"
 #include "df/viewscreen_unitlistst.h"
 #include "df/viewscreen_itemst.h"
+#include "df/viewscreen_layerst.h"
+#include "df/viewscreen_layer_workshop_profilest.h"
+#include "df/viewscreen_layer_noblelistst.h"
+#include "df/viewscreen_layer_overall_healthst.h"
+#include "df/viewscreen_petst.h"
 #include "df/ui_unit_view_mode.h"
 #include "df/ui_sidebar_menus.h"
 #include "df/ui_look_list.h"
@@ -63,12 +68,18 @@ using namespace DFHack;
 #include "df/popup_message.h"
 #include "df/interfacest.h"
 #include "df/graphic.h"
+#include "df/layer_object_listst.h"
 
 using namespace df::enums;
 using df::global::gview;
 using df::global::init;
 using df::global::gps;
 
+static df::layer_object_listst *getLayerList(df::viewscreen_layerst *layer, int idx)
+{
+    return virtual_cast<df::layer_object_listst>(vector_get(layer->layer_objects,idx));
+}
+
 // Predefined common guard functions
 
 bool Gui::default_hotkey(df::viewscreen *top)
@@ -294,6 +305,62 @@ static df::unit *getAnyUnit(df::viewscreen *top)
         return ref ? ref->getUnit() : NULL;
     }
 
+    if (VIRTUAL_CAST_VAR(screen, df::viewscreen_layer_workshop_profilest, top))
+    {
+        if (auto list1 = getLayerList(screen, 0))
+            return vector_get(screen->workers, list1->cursor);
+        return NULL;
+    }
+
+    if (VIRTUAL_CAST_VAR(screen, df::viewscreen_layer_noblelistst, top))
+    {
+        switch (screen->mode)
+        {
+        case df::viewscreen_layer_noblelistst::List:
+            if (auto list1 = getLayerList(screen, 0))
+            {
+                if (auto info = vector_get(screen->info, list1->cursor))
+                    return info->unit;
+            }
+            return NULL;
+
+        case df::viewscreen_layer_noblelistst::Appoint:
+            if (auto list2 = getLayerList(screen, 1))
+            {
+                if (auto info = vector_get(screen->candidates, list2->cursor))
+                    return info->unit;
+            }
+            return NULL;
+
+        default:
+            return NULL;
+        }
+    }
+
+    if (VIRTUAL_CAST_VAR(screen, df::viewscreen_petst, top))
+    {
+        switch (screen->mode)
+        {
+        case df::viewscreen_petst::List:
+            if (!vector_get(screen->is_vermin, screen->cursor))
+                return (df::unit*)vector_get(screen->animal, screen->cursor);
+            return NULL;
+
+        case df::viewscreen_petst::SelectTrainer:
+            return vector_get(screen->trainer_unit, screen->trainer_cursor);
+
+        default:
+            return NULL;
+        }
+    }
+
+    if (VIRTUAL_CAST_VAR(screen, df::viewscreen_layer_overall_healthst, top))
+    {
+        if (auto list1 = getLayerList(screen, 0))
+            return vector_get(screen->unit, list1->cursor);
+        return NULL;
+    }
+
     if (!Gui::dwarfmode_hotkey(top))
         return NULL;