From 7774f5f2c1b1e16b9bd3eefe0fd8b4287d03fef7 Mon Sep 17 00:00:00 2001
From: Alexander Gavrilov <angavrilov@gmail.com>
Date: Sat, 19 May 2012 19:50:36 +0400
Subject: [PATCH] Add a mechanism converting ui focus to a string
 representation.

The idea is to make ui handling more modular, dispensing with
huge functions that switch or if/else on lots of variables.
For now, used to split up functions in the sort plugin.
---
 LUA_API.rst                   |   9 +
 Lua API.html                  |   7 +
 library/LuaApi.cpp            |   3 +
 library/MiscUtils.cpp         |  24 ++
 library/include/DataDefs.h    |   3 +
 library/include/MiscUtils.h   |  45 +++
 library/include/modules/Gui.h |   4 +
 library/modules/Gui.cpp       | 320 ++++++++++++++++++++
 plugins/sort.cpp              | 553 ++++++++++++++--------------------
 9 files changed, 648 insertions(+), 320 deletions(-)

diff --git a/LUA_API.rst b/LUA_API.rst
index 9ecd2e55e..59a52e596 100644
--- a/LUA_API.rst
+++ b/LUA_API.rst
@@ -613,6 +613,15 @@ can be omitted.
 Gui module
 ----------
 
+* ``dfhack.gui.getCurViewscreen()``
+
+  Returns the viewscreen that is current in the core.
+
+* ``dfhack.gui.getFocusString(viewscreen)``
+
+  Returns a string representation of the current focus position
+  in the ui. The string has a "screen/foo/bar/baz..." format.
+
 * ``dfhack.gui.getSelectedWorkshopJob([silent])``
 
   When a job is selected in *'q'* mode, returns the job, else
diff --git a/Lua API.html b/Lua API.html
index 7ce777c42..e2d4ec23e 100644
--- a/Lua API.html	
+++ b/Lua API.html	
@@ -868,6 +868,13 @@ can be omitted.</p>
 <div class="section" id="gui-module">
 <h3><a class="toc-backref" href="#id14">Gui module</a></h3>
 <ul>
+<li><p class="first"><tt class="docutils literal">dfhack.gui.getCurViewscreen()</tt></p>
+<p>Returns the viewscreen that is current in the core.</p>
+</li>
+<li><p class="first"><tt class="docutils literal">dfhack.gui.getFocusString(viewscreen)</tt></p>
+<p>Returns a string representation of the current focus position
+in the ui. The string has a &quot;screen/foo/bar/baz...&quot; format.</p>
+</li>
 <li><p class="first"><tt class="docutils literal"><span class="pre">dfhack.gui.getSelectedWorkshopJob([silent])</span></tt></p>
 <p>When a job is selected in <em>'q'</em> mode, returns the job, else
 prints error unless silent and returns <em>nil</em>.</p>
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index d4c3b157d..9dfb39754 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -61,6 +61,7 @@ distribution.
 #include "df/unit.h"
 #include "df/item.h"
 #include "df/material.h"
+#include "df/viewscreen.h"
 #include "df/assumed_identity.h"
 #include "df/nemesis_record.h"
 #include "df/historical_figure.h"
@@ -652,6 +653,8 @@ static const LuaWrapper::FunctionReg dfhack_module[] = {
 /***** Gui module *****/
 
 static const LuaWrapper::FunctionReg dfhack_gui_module[] = {
+    WRAPM(Gui, getCurViewscreen),
+    WRAPM(Gui, getFocusString),
     WRAPM(Gui, getSelectedWorkshopJob),
     WRAPM(Gui, getSelectedJob),
     WRAPM(Gui, getSelectedUnit),
diff --git a/library/MiscUtils.cpp b/library/MiscUtils.cpp
index 8658788bb..b73d730c1 100644
--- a/library/MiscUtils.cpp
+++ b/library/MiscUtils.cpp
@@ -36,6 +36,7 @@ distribution.
 
 #include <ctype.h>
 #include <stdarg.h>
+#include <string.h>
 
 #include <sstream>
 #include <map>
@@ -124,6 +125,29 @@ std::string toLower(const std::string &str)
     return rv;
 }
 
+bool prefix_matches(const std::string &prefix, const std::string &key, std::string *tail)
+{
+    size_t ksize = key.size();
+    size_t psize = prefix.size();
+    if (ksize < psize || memcmp(prefix.data(), key.data(), psize) != 0)
+        return false;
+    if (tail)
+        tail->clear();
+    if (ksize == psize)
+        return true;
+    if (psize == 0 || prefix[psize-1] == '/')
+    {
+        if (tail) *tail = key.substr(psize);
+        return true;
+    }
+    if (key[psize] == '/')
+    {
+        if (tail) *tail = key.substr(psize+1);
+        return true;
+    }
+    return false;
+}
+
 #ifdef LINUX_BUILD // Linux
 uint64_t GetTimeMs64()
 {
diff --git a/library/include/DataDefs.h b/library/include/DataDefs.h
index 57773e898..1d485156f 100644
--- a/library/include/DataDefs.h
+++ b/library/include/DataDefs.h
@@ -452,6 +452,9 @@ namespace df
         }
     };
 
+    template<class ET, class IT>
+    struct enum_traits<enum_field<ET, IT> > : public enum_traits<ET> {};
+
     template<class EnumType, class IntType1, class IntType2>
     inline bool operator== (enum_field<EnumType,IntType1> a, enum_field<EnumType,IntType2> b)
     {
diff --git a/library/include/MiscUtils.h b/library/include/MiscUtils.h
index 0cf34c489..bbc88a2bb 100644
--- a/library/include/MiscUtils.h
+++ b/library/include/MiscUtils.h
@@ -262,6 +262,51 @@ Link *linked_list_insert_after(Link *pos, Link *link)
     return link;
 }
 
+template<typename T>
+inline typename T::mapped_type map_find(
+    const T &map, const typename T::key_type &key,
+    const typename T::mapped_type &defval = typename T::mapped_type()
+) {
+    auto it = map.find(key);
+    return (it == map.end()) ? defval : it->second;
+}
+
+DFHACK_EXPORT bool prefix_matches(const std::string &prefix, const std::string &key, std::string *tail = NULL);
+
+template<typename T>
+typename T::mapped_type findPrefixInMap(
+    const T &table, const std::string &key,
+    const typename T::mapped_type& defval = typename T::mapped_type()
+) {
+    auto it = table.lower_bound(key);
+    if (it != table.end() && it->first == key)
+        return it->second;
+    if (it != table.begin()) {
+        --it;
+        if (prefix_matches(it->first, key))
+            return it->second;
+    }
+    return defval;
+}
+
+#ifdef __GNUC__
+#define VARIABLE_IS_NOT_USED __attribute__ ((unused))
+#else
+#define VARIABLE_IS_NOT_USED
+#endif
+
+template<class CT>
+inline bool static_add_to_map(CT *pmap, typename CT::key_type key, typename CT::mapped_type value) {
+    (*pmap)[key] = value;
+    return true;
+}
+
+#define CONCAT_TOKENS2(a,b) a##b
+#define CONCAT_TOKENS(a,b) CONCAT_TOKENS2(a,b)
+#define DFHACK_STATIC_ADD_TO_MAP(pmap,key,value) \
+    static bool VARIABLE_IS_NOT_USED CONCAT_TOKENS(static_add_to_map_,__LINE__)\
+        = static_add_to_map(pmap,key,value)
+
 /*
  * MISC
  */
diff --git a/library/include/modules/Gui.h b/library/include/modules/Gui.h
index f1e0335ae..e7155c436 100644
--- a/library/include/modules/Gui.h
+++ b/library/include/modules/Gui.h
@@ -55,6 +55,10 @@ namespace DFHack
      */
     namespace Gui
     {
+        inline df::viewscreen *getCurViewscreen() { return Core::getTopViewscreen(); }
+
+        DFHACK_EXPORT std::string getFocusString(df::viewscreen *top);
+
         // Full-screen item details view
         DFHACK_EXPORT bool item_details_hotkey(df::viewscreen *top);
         // 'u'nits or 'j'obs full-screen view
diff --git a/library/modules/Gui.cpp b/library/modules/Gui.cpp
index 461921391..074eb8370 100644
--- a/library/modules/Gui.cpp
+++ b/library/modules/Gui.cpp
@@ -55,6 +55,7 @@ using namespace DFHack;
 #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_petst.h"
 #include "df/viewscreen_tradegoodsst.h"
 #include "df/ui_unit_view_mode.h"
@@ -64,6 +65,8 @@ using namespace DFHack;
 #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"
@@ -77,12 +80,329 @@ using namespace df::enums;
 using df::global::gview;
 using df::global::init;
 using df::global::gps;
+using df::global::ui;
+using df::global::world;
 
 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));
 }
 
+static std::string getNameChunk(virtual_identity *id, int start, int end)
+{
+    if (!id)
+        return "UNKNOWN";
+    const char *name = id->getName();
+    int len = strlen(name);
+    if (len > start + end)
+        return std::string(name+start, len-start-end);
+    else
+        return name;
+}
+
+/*
+ * Classifying focus context by means of a string path.
+ */
+
+typedef void (*getFocusStringHandler)(std::string &str, df::viewscreen *screen);
+static std::map<virtual_identity*, getFocusStringHandler> getFocusStringHandlers;
+
+#define VIEWSCREEN(name) df::viewscreen_##name##st
+#define DEFINE_GET_FOCUS_STRING_HANDLER(screen_type) \
+    static void getFocusString_##screen_type(std::string &focus, VIEWSCREEN(screen_type) *screen);\
+    DFHACK_STATIC_ADD_TO_MAP(\
+        &getFocusStringHandlers, &VIEWSCREEN(screen_type)::_identity, \
+        (getFocusStringHandler)getFocusString_##screen_type \
+    ); \
+    static void getFocusString_##screen_type(std::string &focus, VIEWSCREEN(screen_type) *screen)
+
+DEFINE_GET_FOCUS_STRING_HANDLER(dwarfmode)
+{
+    using namespace df::enums::ui_sidebar_mode;
+
+    using df::global::ui_workshop_in_add;
+    using df::global::ui_build_selector;
+    using df::global::ui_selected_unit;
+    using df::global::ui_look_list;
+    using df::global::ui_look_cursor;
+    using df::global::ui_building_item_cursor;
+    using df::global::ui_building_assign_type;
+    using df::global::ui_building_assign_is_marked;
+    using df::global::ui_building_assign_units;
+    using df::global::ui_building_assign_items;
+    using df::global::ui_building_in_assign;
+
+    focus += "/" + enum_item_key(ui->main.mode);
+
+    switch (ui->main.mode)
+    {
+    case QueryBuilding:
+        if (df::building *selected = world->selected_building)
+        {
+            if (!selected->jobs.empty() &&
+                selected->jobs[0]->job_type == job_type::DestroyBuilding)
+            {
+                focus += "/Destroying";
+                break;
+            }
+
+            focus += "/Some";
+
+            virtual_identity *id = virtual_identity::get(selected);
+
+            bool jobs = false;
+
+            if (id == &df::building_workshopst::_identity ||
+                id == &df::building_furnacest::_identity)
+            {
+                focus += "/Workshop";
+                jobs = true;
+            }
+            else if (id == &df::building_trapst::_identity)
+            {
+                auto trap = (df::building_trapst*)selected;
+                if (trap->trap_type == trap_type::Lever) {
+                    focus += "/Lever";
+                    jobs = true;
+                }
+            }
+            else if (ui_building_in_assign && *ui_building_in_assign &&
+                     ui_building_assign_type && ui_building_assign_units &&
+                     ui_building_assign_type->size() == ui_building_assign_units->size())
+            {
+                focus += "/Assign";
+                if (ui_building_item_cursor)
+                {
+                    auto unit = vector_get(*ui_building_assign_units, *ui_building_item_cursor);
+                    focus += unit ? "/Unit" : "/None";
+                }
+            }
+
+            if (jobs)
+            {
+                if (ui_workshop_in_add && *ui_workshop_in_add)
+                    focus += "/AddJob";
+                else if (!selected->jobs.empty())
+                    focus += "/Job";
+                else
+                    focus += "/Empty";
+            }
+        }
+        else
+            focus += "/None";
+        break;
+
+    case Build:
+        if (ui_build_selector)
+        {
+            // Not selecting, or no choices?
+            if (ui_build_selector->building_type < 0)
+                focus += "/Type";
+            else if (ui_build_selector->stage != 2)
+                focus += "/Position";
+            else
+            {
+                focus += "/Material";
+                if (ui_build_selector->is_grouped)
+                    focus += "/Groups";
+                else
+                    focus += "/Items";
+            }
+        }
+        break;
+
+    case ViewUnits:
+        if (ui_selected_unit)
+        {
+            if (auto unit = vector_get(world->units.active, *ui_selected_unit))
+            {
+                focus += "/Some";
+
+                using df::global::ui_unit_view_mode;
+
+                if (ui_unit_view_mode)
+                    focus += "/" + enum_item_key(ui_unit_view_mode->value);
+            }
+            else
+                focus += "/None";
+        }
+        break;
+
+    case LookAround:
+        if (ui_look_list && ui_look_cursor)
+        {
+            auto item = vector_get(ui_look_list->items, *ui_look_cursor);
+            if (item)
+                focus += "/" + enum_item_key(item->type);
+            else
+                focus += "/None";
+        }
+        break;
+
+    case BuildingItems:
+        if (VIRTUAL_CAST_VAR(selected, df::building_actual, world->selected_building))
+        {
+            if (selected->contained_items.empty())
+                focus += "/Some/Empty";
+            else
+                focus += "/Some/Item";
+        }
+        else
+            focus += "/None";
+        break;
+
+    case ZonesPenInfo:
+        if (ui_building_assign_type && ui_building_assign_units &&
+            ui_building_assign_is_marked && ui_building_assign_items &&
+            ui_building_assign_type->size() == ui_building_assign_units->size())
+        {
+            focus += "/Assign";
+            if (ui_building_item_cursor)
+            {
+                if (vector_get(*ui_building_assign_units, *ui_building_item_cursor))
+                    focus += "/Unit";
+                else if (vector_get(*ui_building_assign_items, *ui_building_item_cursor))
+                    focus += "/Vermin";
+                else
+                    focus += "/None";
+            }
+        }
+        break;
+
+    case Burrows:
+        if (ui->burrows.in_add_units_mode)
+            focus += "/AddUnits";
+        else if (ui->burrows.in_edit_name_mode)
+            focus += "/EditName";
+        else if (ui->burrows.in_define_mode)
+            focus += "/Define";
+        else
+            focus += "/List";
+        break;
+
+    default:
+        break;
+    }
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(unitlist)
+{
+    focus += "/" + enum_item_key(screen->page);
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(layer_military)
+{
+    auto list1 = getLayerList(screen, 0);
+    auto list2 = getLayerList(screen, 1);
+    auto list3 = getLayerList(screen, 2);
+    if (!list1 || !list2 || !list3) return;
+
+    focus += "/" + enum_item_key(screen->page);
+
+    int cur_list;
+    if (list1->bright) cur_list = 0;
+    else if (list2->bright) cur_list = 1;
+    else if (list3->bright) cur_list = 2;
+    else return;
+
+    switch (screen->page)
+    {
+    case df::viewscreen_layer_militaryst::Positions:
+        {
+            static const char *lists[] = { "/Squads", "/Positions", "/Candidates" };
+            focus += lists[cur_list];
+            break;
+        }
+
+    default:
+        break;
+    }
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(layer_workshop_profile)
+{
+    auto list1 = getLayerList(screen, 0);
+    if (!list1) return;
+
+    if (vector_get(screen->workers, list1->cursor))
+        focus += "/Unit";
+    else
+        focus += "/None";
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(layer_noblelist)
+{
+    auto list1 = getLayerList(screen, 0);
+    auto list2 = getLayerList(screen, 1);
+    if (!list1 || !list2) return;
+
+    focus += "/" + enum_item_key(screen->mode);
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(pet)
+{
+    focus += "/" + enum_item_key(screen->mode);
+
+    switch (screen->mode)
+    {
+    case df::viewscreen_petst::List:
+        focus += vector_get(screen->is_vermin, screen->cursor) ? "/Vermin" : "/Unit";
+        break;
+
+    case df::viewscreen_petst::SelectTrainer:
+        if (vector_get(screen->trainer_unit, screen->trainer_cursor))
+            focus += "/Unit";
+        break;
+    }
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(layer_overall_health)
+{
+    auto list1 = getLayerList(screen, 0);
+    if (!list1) return;
+
+    focus += "/Units";
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(tradegoods)
+{
+    if (!screen->has_traders || screen->is_unloading)
+        focus += "/NoTraders";
+    else if (screen->in_edit_count)
+        focus += "/EditCount";
+    else
+        focus += (screen->in_right_pane ? "/Items/Broker" : "/Items/Trader");
+}
+
+DEFINE_GET_FOCUS_STRING_HANDLER(layer_assigntrade)
+{
+    auto list1 = getLayerList(screen, 0);
+    auto list2 = getLayerList(screen, 1);
+    if (!list1 || !list2) return;
+
+    int list_idx = vector_get(screen->visible_lists, list1->cursor, (int16_t)-1);
+    unsigned num_lists = sizeof(screen->lists)/sizeof(screen->lists[0]);
+    if (unsigned(list_idx) >= num_lists)
+        return;
+
+    if (list1->bright)
+        focus += "/Groups";
+    else
+        focus += "/Items";
+}
+
+std::string Gui::getFocusString(df::viewscreen *top)
+{
+    virtual_identity *id = virtual_identity::get(top);
+    std::string name = getNameChunk(id, 11, 2);
+
+    auto handler = map_find(getFocusStringHandlers, id);
+    if (handler)
+        handler(name, top);
+
+    return name;
+}
+
 // Predefined common guard functions
 
 bool Gui::default_hotkey(df::viewscreen *top)
diff --git a/plugins/sort.cpp b/plugins/sort.cpp
index 430fb7192..9274a4cfb 100644
--- a/plugins/sort.cpp
+++ b/plugins/sort.cpp
@@ -204,7 +204,8 @@ static bool ParseSpec(color_ostream &out, lua_State *L, const char *type, vector
 }
 
 #define PARSE_SPEC(type, params) \
-    if (!ParseSpec(*pout, L, type, params)) return false;
+    std::vector<unsigned> order; \
+    if (!ParseSpec(*pout, L, type, params)) return;
 
 static bool prepare_sort(color_ostream *pout, lua_State *L)
 {
@@ -230,310 +231,226 @@ static df::layer_object_listst *getLayerList(df::viewscreen_layerst *layer, int
     return virtual_cast<df::layer_object_listst>(vector_get(layer->layer_objects,idx));
 }
 
-static bool maybe_sort_units(color_ostream *pout, lua_State *L,
-                             df::viewscreen *screen, vector<string> &parameters)
-{
-    Lua::StackUnwinder top(L);
+typedef void (*SortHandler)(color_ostream *pout, lua_State *L, int top,
+                            df::viewscreen *screen, vector<string> &parameters);
 
-    if (!prepare_sort(pout, L))
-        return false;
+#define VIEWSCREEN(name) df::viewscreen_##name##st
+#define DEFINE_SORT_HANDLER(map, screen_type, tail, screen) \
+    static void CONCAT_TOKENS(SortHandler_##screen_type,__LINE__)\
+        (color_ostream *pout, lua_State *L, int top, \
+         VIEWSCREEN(screen_type) *screen, vector<string> &parameters); \
+    DFHACK_STATIC_ADD_TO_MAP(&map, #screen_type tail, \
+        (SortHandler)CONCAT_TOKENS(SortHandler_##screen_type,__LINE__) ); \
+    static void CONCAT_TOKENS(SortHandler_##screen_type,__LINE__)\
+        (color_ostream *pout, lua_State *L, int top, \
+         VIEWSCREEN(screen_type) *screen, vector<string> &parameters)
 
-    std::vector<unsigned> order;
+static std::map<std::string, SortHandler> unit_sorters;
 
-    if (auto units = strict_virtual_cast<df::viewscreen_unitlistst>(screen))
-    {
-        if (!L) return true;
+/*
+ * Sort units in the 'u'nit list screen.
+ */
 
-        /*
-         * Sort units in the 'u'nit list screen.
-         */
+DEFINE_SORT_HANDLER(unit_sorters, unitlist, "", units)
+{
+    PARSE_SPEC("units", parameters);
 
-        PARSE_SPEC("units", parameters);
+    int page = units->page;
 
-        int page = units->page;
+    if (compute_order(*pout, L, top, &order, units->units[page]))
+    {
+        reorder_cursor(&units->cursor_pos[page], order);
+        reorder_vector(&units->units[page], order);
+        reorder_vector(&units->jobs[page], order);
+    }
+}
 
-        if (compute_order(*pout, L, top, &order, units->units[page]))
-        {
-            reorder_cursor(&units->cursor_pos[page], order);
-            reorder_vector(&units->units[page], order);
-            reorder_vector(&units->jobs[page], order);
-        }
+/*
+ * Sort units in the 'j'ob list screen.
+ */
 
-        return true;
-    }
-    else if (auto jobs = strict_virtual_cast<df::viewscreen_joblistst>(screen))
+DEFINE_SORT_HANDLER(unit_sorters, joblist, "", jobs)
+{
+    PARSE_SPEC("units", parameters);
+
+    if (compute_order(*pout, L, top, &order, jobs->units))
     {
-        if (!L) return true;
+        reorder_cursor(&jobs->cursor_pos, order);
+        reorder_vector(&jobs->units, order);
+        reorder_vector(&jobs->jobs, order);
+    }
+}
 
-        /*
-         * Sort units in the 'j'ob list screen.
-         */
+/*
+ * Sort candidate units in the 'p'osition page of the 'm'ilitary screen.
+ */
 
-        PARSE_SPEC("units", parameters);
+DEFINE_SORT_HANDLER(unit_sorters, layer_military, "/Positions/Candidates", military)
+{
+    auto &candidates = military->positions.candidates;
+    auto list3 = getLayerList(military, 2);
 
-        if (compute_order(*pout, L, top, &order, jobs->units))
-        {
-            reorder_cursor(&jobs->cursor_pos, order);
-            reorder_vector(&jobs->units, order);
-            reorder_vector(&jobs->jobs, order);
-        }
+    PARSE_SPEC("units", parameters);
 
-        return true;
-    }
-    else if (auto military = strict_virtual_cast<df::viewscreen_layer_militaryst>(screen))
+    if (compute_order(*pout, L, top, &order, candidates))
     {
-        switch (military->page)
-        {
-        case df::viewscreen_layer_militaryst::Positions:
-            {
-                auto &candidates = military->positions.candidates;
-                auto list3 = getLayerList(military, 2);
+        reorder_cursor(&list3->cursor, order);
+        reorder_vector(&candidates, order);
+    }
+}
 
-                /*
-                 * Sort candidate units in the 'p'osition page of the 'm'ilitary screen.
-                 */
 
-                if (list3 && !candidates.empty() && list3->bright)
-                {
-                    if (!L) return true;
+/*
+ * Sort units in the workshop 'q'uery 'P'rofile modification screen.
+ */
 
-                    PARSE_SPEC("units", parameters);
+DEFINE_SORT_HANDLER(unit_sorters, layer_workshop_profile, "/Unit", profile)
+{
+    auto list1 = getLayerList(profile, 0);
 
-                    if (compute_order(*pout, L, top, &order, candidates))
-                    {
-                        reorder_cursor(&list3->cursor, order);
-                        reorder_vector(&candidates, order);
-                    }
+    PARSE_SPEC("units", parameters);
 
-                    return true;
-                }
+    if (compute_order(*pout, L, top, &order, profile->workers))
+    {
+        reorder_cursor(&list1->cursor, order);
+        reorder_vector(&profile->workers, order);
+    }
+}
 
-                return false;
-            }
+DEFINE_SORT_HANDLER(unit_sorters, layer_noblelist, "/Appoint", nobles)
+{
+    auto list2 = getLayerList(nobles, 1);
 
-        default:
-            return false;
-        }
-    }
-    else if (auto profile = strict_virtual_cast<df::viewscreen_layer_workshop_profilest>(screen))
-    {
-        auto list1 = getLayerList(profile, 0);
+    sort_null_first(parameters);
+    PARSE_SPEC("units", parameters);
 
-        if (!list1) return false;
-        if (!L) return true;
+    std::vector<df::unit*> units;
+    for (size_t i = 0; i < nobles->candidates.size(); i++)
+        units.push_back(nobles->candidates[i]->unit);
 
-        /*
-         * Sort units in the workshop 'q'uery 'P'rofile modification screen.
-         */
+    if (compute_order(*pout, L, top, &order, units))
+    {
+        reorder_cursor(&list2->cursor, order);
+        reorder_vector(&nobles->candidates, order);
+    }
+}
 
-        PARSE_SPEC("units", parameters);
+/*
+ * Sort animal units in the Animal page of the 'z' status screen.
+ */
 
-        if (compute_order(*pout, L, top, &order, profile->workers))
-        {
-            reorder_cursor(&list1->cursor, order);
-            reorder_vector(&profile->workers, order);
-        }
+DEFINE_SORT_HANDLER(unit_sorters, pet, "/List", animals)
+{
+    PARSE_SPEC("units", parameters);
 
-        return true;
-    }
-    else if (auto nobles = strict_virtual_cast<df::viewscreen_layer_noblelistst>(screen))
-    {
-        switch (nobles->mode)
-        {
-        case df::viewscreen_layer_noblelistst::Appoint:
-            {
-                auto list2 = getLayerList(nobles, 1);
+    std::vector<df::unit*> units;
+    for (size_t i = 0; i < animals->animal.size(); i++)
+        units.push_back(animals->is_vermin[i] ? NULL : (df::unit*)animals->animal[i]);
 
-                /*
-                 * Sort units in the appointment candidate list of the 'n'obles screen.
-                 */
+    if (compute_order(*pout, L, top, &order, units))
+    {
+        reorder_cursor(&animals->cursor, order);
+        reorder_vector(&animals->animal, order);
+        reorder_vector(&animals->is_vermin, order);
+        reorder_vector(&animals->pet_info, order);
+        reorder_vector(&animals->is_tame, order);
+        reorder_vector(&animals->is_adopting, order);
+    }
+}
 
-                if (list2)
-                {
-                    if (!L) return true;
+/*
+ * Sort candidate trainers in the Animal page of the 'z' status screen.
+ */
 
-                    sort_null_first(parameters);
-                    PARSE_SPEC("units", parameters);
+DEFINE_SORT_HANDLER(unit_sorters, pet, "/SelectTrainer", animals)
+{
+    sort_null_first(parameters);
+    PARSE_SPEC("units", parameters);
 
-                    std::vector<df::unit*> units;
-                    for (size_t i = 0; i < nobles->candidates.size(); i++)
-                        units.push_back(nobles->candidates[i]->unit);
+    if (compute_order(*pout, L, top, &order, animals->trainer_unit))
+    {
+        reorder_cursor(&animals->trainer_cursor, order);
+        reorder_vector(&animals->trainer_unit, order);
+        reorder_vector(&animals->trainer_mode, order);
+    }
+}
 
-                    if (compute_order(*pout, L, top, &order, units))
-                    {
-                        reorder_cursor(&list2->cursor, order);
-                        reorder_vector(&nobles->candidates, order);
-                    }
+/*
+ * Sort units in the Health page of the 'z' status screen.
+ */
 
-                    return true;
-                }
+DEFINE_SORT_HANDLER(unit_sorters, layer_overall_health, "/Units", health)
+{
+    auto list1 = getLayerList(health, 0);
 
-                return false;
-            }
+    PARSE_SPEC("units", parameters);
 
-        default:
-            return false;
-        }
-    }
-    else if (auto animals = strict_virtual_cast<df::viewscreen_petst>(screen))
+    if (compute_order(*pout, L, top, &order, health->unit))
     {
-        switch (animals->mode)
-        {
-        case df::viewscreen_petst::List:
-            {
-                if (!L) return true;
-
-                /*
-                 * Sort animal units in the Animal page of the 'z' status screen.
-                 */
-
-                PARSE_SPEC("units", parameters);
-
-                std::vector<df::unit*> units;
-                for (size_t i = 0; i < animals->animal.size(); i++)
-                    units.push_back(animals->is_vermin[i] ? NULL : (df::unit*)animals->animal[i]);
-
-                if (compute_order(*pout, L, top, &order, units))
-                {
-                    reorder_cursor(&animals->cursor, order);
-                    reorder_vector(&animals->animal, order);
-                    reorder_vector(&animals->is_vermin, order);
-                    reorder_vector(&animals->pet_info, order);
-                    reorder_vector(&animals->is_tame, order);
-                    reorder_vector(&animals->is_adopting, order);
-                }
-
-                return true;
-            }
-
-        case df::viewscreen_petst::SelectTrainer:
-            {
-                if (!L) return true;
-
-                /*
-                 * Sort candidate trainers in the Animal page of the 'z' status screen.
-                 */
-
-                sort_null_first(parameters);
-                PARSE_SPEC("units", parameters);
-
-                if (compute_order(*pout, L, top, &order, animals->trainer_unit))
-                {
-                    reorder_cursor(&animals->trainer_cursor, order);
-                    reorder_vector(&animals->trainer_unit, order);
-                    reorder_vector(&animals->trainer_mode, order);
-                }
-
-                return true;
-            }
-
-        default:
-            return false;
-        }
+        reorder_cursor(&list1->cursor, order);
+        reorder_vector(&health->unit, order);
+        reorder_vector(&health->bits1, order);
+        reorder_vector(&health->bits2, order);
+        reorder_vector(&health->bits3, order);
     }
-    else if (auto health = strict_virtual_cast<df::viewscreen_layer_overall_healthst>(screen))
-    {
-        auto list1 = getLayerList(health, 0);
+}
+
+/*
+ * Sort burrow member candidate units in the 'w' sidebar mode.
+ */
+
+DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/Burrows/AddUnits", screen)
+{
+    PARSE_SPEC("units", parameters);
 
-        if (!list1) return false;
-        if (!L) return true;
+    if (compute_order(*pout, L, top, &order, ui->burrows.list_units))
+    {
+        reorder_cursor(&ui->burrows.unit_cursor_pos, order);
+        reorder_vector(&ui->burrows.list_units, order);
+        reorder_vector(&ui->burrows.sel_units, order);
+    }
+}
 
-        /*
-         * Sort units in the Health page of the 'z' status screen.
-         */
+/*
+ * Sort building owner candidate units in the 'q' sidebar mode, or cage assignment.
+ */
 
-        PARSE_SPEC("units", parameters);
+DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/QueryBuilding/Some/Assign", screen)
+{
+    sort_null_first(parameters);
 
-        if (compute_order(*pout, L, top, &order, health->unit))
-        {
-            reorder_cursor(&list1->cursor, order);
-            reorder_vector(&health->unit, order);
-            reorder_vector(&health->bits1, order);
-            reorder_vector(&health->bits2, order);
-            reorder_vector(&health->bits3, order);
-        }
+    PARSE_SPEC("units", parameters);
 
-        return true;
-    }
-    else if (strict_virtual_cast<df::viewscreen_dwarfmodest>(screen))
+    if (compute_order(*pout, L, top, &order, *ui_building_assign_units))
     {
-        switch (ui->main.mode)
-        {
-        case ui_sidebar_mode::Burrows:
-            if (!L) return true;
-
-            /*
-             * Sort burrow member candidate units in the 'w' sidebar mode.
-             */
-
-            PARSE_SPEC("units", parameters);
-
-            if (compute_order(*pout, L, top, &order, ui->burrows.list_units))
-            {
-                reorder_cursor(&ui->burrows.unit_cursor_pos, order);
-                reorder_vector(&ui->burrows.list_units, order);
-                reorder_vector(&ui->burrows.sel_units, order);
-            }
-
-            return true;
-
-        case ui_sidebar_mode::QueryBuilding:
-            if (!ui_building_in_assign || !*ui_building_in_assign)
-                return false;
-            // fall through for building owner / chain assign animal
-
-        case ui_sidebar_mode::ZonesPenInfo:
-            if (ui_building_item_cursor &&
-                ui_building_assign_type &&
-                ui_building_assign_is_marked &&
-                ui_building_assign_units &&
-                ui_building_assign_items &&
-                ui_building_assign_type->size() == ui_building_assign_units->size() &&
-                !ui_building_assign_type->empty())
-            {
-                if (!L) return true;
-
-                /*
-                 * Sort building owner candidate units in the 'q' sidebar mode,
-                 * or pen assignment candidate units in 'z'->'N', or cage assignment.
-                 */
-
-                // TODO: better way
-                bool is_assign_owner = ((*ui_building_assign_type)[0] == -1);
-
-                if (is_assign_owner)
-                    sort_null_first(parameters);
-
-                PARSE_SPEC("units", parameters);
-
-                if (compute_order(*pout, L, top, &order, *ui_building_assign_units))
-                {
-                    reorder_cursor(ui_building_item_cursor, order);
-                    reorder_vector(ui_building_assign_type, order);
-                    reorder_vector(ui_building_assign_units, order);
-
-                    if (ui_building_assign_units->size() == ui_building_assign_items->size())
-                        reorder_vector(ui_building_assign_items, order);
-                    if (ui_building_assign_units->size() == ui_building_assign_is_marked->size())
-                        reorder_vector(ui_building_assign_is_marked, order);
-                }
-
-                return true;
-            }
-            return false;
+        reorder_cursor(ui_building_item_cursor, order);
+        reorder_vector(ui_building_assign_type, order);
+        reorder_vector(ui_building_assign_units, order);
+    }
+}
 
-        default:
-            return false;
-        }
+/*
+ * Sort pen assignment candidate units in 'z'->'N'.
+ */
+
+DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/ZonesPenInfo/Assign", screen)
+{
+    PARSE_SPEC("units", parameters);
+
+    if (compute_order(*pout, L, top, &order, *ui_building_assign_units))
+    {
+        reorder_cursor(ui_building_item_cursor, order);
+        reorder_vector(ui_building_assign_type, order);
+        reorder_vector(ui_building_assign_units, order);
+        reorder_vector(ui_building_assign_items, order);
+        reorder_vector(ui_building_assign_is_marked, order);
     }
-    else
-        return false;
 }
 
 static bool unit_list_hotkey(df::viewscreen *screen)
 {
-    vector<string> dummy;
-    return maybe_sort_units(NULL, NULL, screen, dummy);
+    auto focus = Gui::getFocusString(screen);
+    return findPrefixInMap(unit_sorters, focus) != NULL;
 }
 
 static command_result sort_units(color_ostream &out, vector <string> &parameters)
@@ -544,89 +461,75 @@ static command_result sort_units(color_ostream &out, vector <string> &parameters
     auto L = Lua::Core::State;
     auto screen = Core::getInstance().getTopViewscreen();
 
-    if (!maybe_sort_units(&out, L, screen, parameters))
+    Lua::StackUnwinder top(L);
+
+    if (!prepare_sort(&out, L))
         return CR_WRONG_USAGE;
 
+    auto focus = Gui::getFocusString(screen);
+    auto handler = findPrefixInMap(unit_sorters, focus);
+
+    if (!handler)
+        return CR_WRONG_USAGE;
+    else
+        handler(&out, L, top, screen, parameters);
+
     return CR_OK;
 }
 
-static bool maybe_sort_items(color_ostream *pout, lua_State *L,
-                             df::viewscreen *screen, vector<string> &parameters)
-{
-    Lua::StackUnwinder top(L);
-
-    if (!prepare_sort(pout, L))
-        return false;
+static std::map<std::string, SortHandler> item_sorters;
 
-    std::vector<unsigned> order;
+DEFINE_SORT_HANDLER(item_sorters, tradegoods, "/Items/Broker", trade)
+{
+    PARSE_SPEC("items", parameters);
 
-    if (auto trade = strict_virtual_cast<df::viewscreen_tradegoodsst>(screen))
+    if (compute_order(*pout, L, top, &order, trade->broker_items))
     {
-        if (!L) return true;
-
-        PARSE_SPEC("items", parameters);
-
-        if (trade->in_right_pane)
-        {
-            if (compute_order(*pout, L, top, &order, trade->broker_items))
-            {
-                reorder_cursor(&trade->broker_cursor, order);
-                reorder_vector(&trade->broker_items, order);
-                reorder_vector(&trade->broker_selected, order);
-                reorder_vector(&trade->broker_count, order);
-            }
-        }
-        else
-        {
-            if (compute_order(*pout, L, top, &order, trade->trader_items))
-            {
-                reorder_cursor(&trade->trader_cursor, order);
-                reorder_vector(&trade->trader_items, order);
-                reorder_vector(&trade->trader_selected, order);
-                reorder_vector(&trade->trader_count, order);
-            }
-        }
-
-        return true;
+        reorder_cursor(&trade->broker_cursor, order);
+        reorder_vector(&trade->broker_items, order);
+        reorder_vector(&trade->broker_selected, order);
+        reorder_vector(&trade->broker_count, order);
     }
-    else if (auto bring = strict_virtual_cast<df::viewscreen_layer_assigntradest>(screen))
-    {
-        auto list1 = getLayerList(bring, 0);
-        auto list2 = getLayerList(bring, 1);
-        if (!list1 || !list2 || !list2->bright)
-            return false;
+}
 
-        int list_idx = vector_get(bring->visible_lists, list1->cursor, (int16_t)-1);
-        unsigned num_lists = sizeof(bring->lists)/sizeof(std::vector<int32_t>);
-        if (unsigned(list_idx) >= num_lists)
-            return false;
+DEFINE_SORT_HANDLER(item_sorters, tradegoods, "/Items/Trader", trade)
+{
+    PARSE_SPEC("items", parameters);
 
-        if (!L) return true;
+    if (compute_order(*pout, L, top, &order, trade->trader_items))
+    {
+        reorder_cursor(&trade->trader_cursor, order);
+        reorder_vector(&trade->trader_items, order);
+        reorder_vector(&trade->trader_selected, order);
+        reorder_vector(&trade->trader_count, order);
+    }
+}
 
-        PARSE_SPEC("items", parameters);
+DEFINE_SORT_HANDLER(item_sorters, layer_assigntrade, "/Items", bring)
+{
+    auto list1 = getLayerList(bring, 0);
+    auto list2 = getLayerList(bring, 1);
+    int list_idx = vector_get(bring->visible_lists, list1->cursor, (int16_t)-1);
 
-        auto &vec = bring->lists[list_idx];
+    PARSE_SPEC("items", parameters);
 
-        std::vector<df::item*> items;
-        for (size_t i = 0; i < vec.size(); i++)
-            items.push_back(bring->info[vec[i]]->item);
+    auto &vec = bring->lists[list_idx];
 
-        if (compute_order(*pout, L, top, &order, items))
-        {
-            reorder_cursor(&list2->cursor, order);
-            reorder_vector(&vec, order);
-        }
+    std::vector<df::item*> items;
+    for (size_t i = 0; i < vec.size(); i++)
+        items.push_back(bring->info[vec[i]]->item);
 
-        return true;
+    if (compute_order(*pout, L, top, &order, items))
+    {
+        reorder_cursor(&list2->cursor, order);
+        reorder_vector(&vec, order);
     }
-    else
-        return false;
 }
 
 static bool item_list_hotkey(df::viewscreen *screen)
 {
-    vector<string> dummy;
-    return maybe_sort_items(NULL, NULL, screen, dummy);
+    auto focus = Gui::getFocusString(screen);
+    return findPrefixInMap(item_sorters, focus) != NULL;
 }
 
 static command_result sort_items(color_ostream &out, vector <string> &parameters)
@@ -637,8 +540,18 @@ static command_result sort_items(color_ostream &out, vector <string> &parameters
     auto L = Lua::Core::State;
     auto screen = Core::getInstance().getTopViewscreen();
 
-    if (!maybe_sort_items(&out, L, screen, parameters))
+    Lua::StackUnwinder top(L);
+
+    if (!prepare_sort(&out, L))
         return CR_WRONG_USAGE;
 
+    auto focus = Gui::getFocusString(screen);
+    auto handler = findPrefixInMap(item_sorters, focus);
+
+    if (!handler)
+        return CR_WRONG_USAGE;
+    else
+        handler(&out, L, top, screen, parameters);
+
     return CR_OK;
 }