Merge https://github.com/angavrilov/dfhack

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

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>

library/Core.cpp

@@ -75,7 +75,7 @@ using df::global::world;
 
 // FIXME: A lot of code in one file, all doing different things... there's something fishy about it.
 
-static bool parseKeySpec(std::string keyspec, int *psym, int *pmod);
+static bool parseKeySpec(std::string keyspec, int *psym, int *pmod, std::string *pfocus = NULL);
 
 struct Core::Cond
 {
@@ -558,10 +558,15 @@ command_result Core::runCommand(color_ostream &con, const std::string &first, ve
             {
                 con << "Usage:" << endl
                     << "  keybinding list <key>" << endl
-                    << "  keybinding clear <key> <key>..." << endl
-                    << "  keybinding set <key> \"cmdline\" \"cmdline\"..." << endl
-                    << "  keybinding add <key> \"cmdline\" \"cmdline\"..." << endl
-                    << "Later adds, and earlier items within one command have priority." << endl;
+                    << "  keybinding clear <key>[@context]..." << endl
+                    << "  keybinding set <key>[@context] \"cmdline\" \"cmdline\"..." << endl
+                    << "  keybinding add <key>[@context] \"cmdline\" \"cmdline\"..." << endl
+                    << "Later adds, and earlier items within one command have priority." << endl
+                    << "Supported keys: [Ctrl-][Alt-][Shift-](A-Z, or F1-F9, or Enter)." << endl
+                    << "Context may be used to limit the scope of the binding, by" << endl
+                    << "requiring the current context to have a certain prefix." << endl
+                    << "Current UI context is: "
+                    << Gui::getFocusString(Core::getTopViewscreen()) << endl;
             }
         }
         else if(first == "fpause")
@@ -1254,7 +1259,7 @@ int Core::SDL_Event(SDL::Event* ev)
     if(!started) return true;
     if(!ev)
         return true;
-    if(ev && ev->type == SDL::ET_KEYDOWN || ev->type == SDL::ET_KEYUP)
+    if(ev && (ev->type == SDL::ET_KEYDOWN || ev->type == SDL::ET_KEYUP))
     {
         SDL::KeyboardEvent * ke = (SDL::KeyboardEvent *)ev;
 
@@ -1299,6 +1304,9 @@ bool Core::SelectHotkey(int sym, int modifiers)
     while (screen->child)
         screen = screen->child;
 
+    if (sym == SDL::K_KP_ENTER)
+        sym = SDL::K_RETURN;
+
     std::string cmd;
 
     {
@@ -1309,6 +1317,9 @@ bool Core::SelectHotkey(int sym, int modifiers)
         for (int i = bindings.size()-1; i >= 0; --i) {
             if (bindings[i].modifiers != modifiers)
                 continue;
+            if (!bindings[i].focus.empty() &&
+                !prefix_matches(bindings[i].focus, Gui::getFocusString(screen)))
+                continue;
             if (!plug_mgr->CanInvokeHotkey(bindings[i].command[0], screen))
                 continue;
             cmd = bindings[i].cmdline;
@@ -1340,10 +1351,22 @@ bool Core::SelectHotkey(int sym, int modifiers)
         return false;
 }
 
-static bool parseKeySpec(std::string keyspec, int *psym, int *pmod)
+static bool parseKeySpec(std::string keyspec, int *psym, int *pmod, std::string *pfocus)
 {
     *pmod = 0;
 
+    if (pfocus)
+    {
+        *pfocus = "";
+
+        size_t idx = keyspec.find('@');
+        if (idx != std::string::npos)
+        {
+            *pfocus = keyspec.substr(idx+1);
+            keyspec = keyspec.substr(0, idx);
+        }
+    }
+
     // ugh, ugly
     for (;;) {
         if (keyspec.size() > 6 && keyspec.substr(0, 6) == "Shift-") {
@@ -1365,6 +1388,9 @@ static bool parseKeySpec(std::string keyspec, int *psym, int *pmod)
     } else if (keyspec.size() == 2 && keyspec[0] == 'F' && keyspec[1] >= '1' && keyspec[1] <= '9') {
         *psym = SDL::K_F1 + (keyspec[1]-'1');
         return true;
+    } else if (keyspec == "Enter") {
+        *psym = SDL::K_RETURN;
+        return true;
     } else
         return false;
 }
@@ -1372,14 +1398,15 @@ static bool parseKeySpec(std::string keyspec, int *psym, int *pmod)
 bool Core::ClearKeyBindings(std::string keyspec)
 {
     int sym, mod;
-    if (!parseKeySpec(keyspec, &sym, &mod))
+    std::string focus;
+    if (!parseKeySpec(keyspec, &sym, &mod, &focus))
         return false;
 
     tthread::lock_guard<tthread::mutex> lock(*HotkeyMutex);
 
     std::vector<KeyBinding> &bindings = key_bindings[sym];
     for (int i = bindings.size()-1; i >= 0; --i) {
-        if (bindings[i].modifiers == mod)
+        if (bindings[i].modifiers == mod && prefix_matches(focus, bindings[i].focus))
             bindings.erase(bindings.begin()+i);
     }
 
@@ -1390,7 +1417,7 @@ bool Core::AddKeyBinding(std::string keyspec, std::string cmdline)
 {
     int sym;
     KeyBinding binding;
-    if (!parseKeySpec(keyspec, &sym, &binding.modifiers))
+    if (!parseKeySpec(keyspec, &sym, &binding.modifiers, &binding.focus))
         return false;
 
     cheap_tokenise(cmdline, binding.command);
@@ -1403,7 +1430,8 @@ bool Core::AddKeyBinding(std::string keyspec, std::string cmdline)
     std::vector<KeyBinding> &bindings = key_bindings[sym];
     for (int i = bindings.size()-1; i >= 0; --i) {
         if (bindings[i].modifiers == binding.modifiers &&
-            bindings[i].cmdline == cmdline)
+            bindings[i].cmdline == cmdline &&
+            bindings[i].focus == binding.focus)
             return true;
     }
 
@@ -1424,7 +1452,12 @@ std::vector<std::string> Core::ListKeyBindings(std::string keyspec)
     std::vector<KeyBinding> &bindings = key_bindings[sym];
     for (int i = bindings.size()-1; i >= 0; --i) {
         if (bindings[i].modifiers == mod)
-            rv.push_back(bindings[i].cmdline);
+        {
+            std::string cmd = bindings[i].cmdline;
+            if (!bindings[i].focus.empty())
+                cmd = "@" + bindings[i].focus + ": " + cmd;
+            rv.push_back(cmd);
+        }
     }
 
     return rv;

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),

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()
 {

library/include/Core.h

@@ -205,6 +205,7 @@ namespace DFHack
             int modifiers;
             std::vector<std::string> command;
             std::string cmdline;
+            std::string focus;
         };
 
         std::map<int, std::vector<KeyBinding> > key_bindings;

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)
     {

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
  */

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

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,340 @@ 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)
+{
+    if (!top)
+        return "";
+
+    if (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;
+    }
+    else
+    {
+        Core &core = Core::getInstance();
+        std::string name = core.p->readClassName(*(void**)top);
+        return name.substr(11, name.size()-11-2);
+    }
+}
+
 // Predefined common guard functions
 
 bool Gui::default_hotkey(df::viewscreen *top)

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,24 +231,28 @@ 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);
-
-    if (!prepare_sort(pout, L))
-        return false;
+typedef void (*SortHandler)(color_ostream *pout, lua_State *L, int top,
+                            df::viewscreen *screen, vector<string> &parameters);
 
-    std::vector<unsigned> order;
+#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)
 
-    if (auto units = strict_virtual_cast<df::viewscreen_unitlistst>(screen))
-    {
-        if (!L) return true;
+static std::map<std::string, SortHandler> unit_sorters;
 
-        /*
+/*
  * Sort units in the 'u'nit list screen.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, unitlist, "", units)
+{
     PARSE_SPEC("units", parameters);
 
     int page = units->page;
@@ -258,17 +263,14 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_vector(&units->units[page], order);
         reorder_vector(&units->jobs[page], order);
     }
+}
 
-        return true;
-    }
-    else if (auto jobs = strict_virtual_cast<df::viewscreen_joblistst>(screen))
-    {
-        if (!L) return true;
-
-        /*
+/*
  * Sort units in the 'j'ob list screen.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, joblist, "", jobs)
+{
     PARSE_SPEC("units", parameters);
 
     if (compute_order(*pout, L, top, &order, jobs->units))
@@ -277,25 +279,16 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_vector(&jobs->units, order);
         reorder_vector(&jobs->jobs, order);
     }
+}
 
-        return true;
-    }
-    else if (auto military = strict_virtual_cast<df::viewscreen_layer_militaryst>(screen))
-    {
-        switch (military->page)
-        {
-        case df::viewscreen_layer_militaryst::Positions:
-            {
-                auto &candidates = military->positions.candidates;
-                auto list3 = getLayerList(military, 2);
-
-                /*
+/*
  * Sort candidate units in the 'p'osition page of the 'm'ilitary screen.
  */
 
-                if (list3 && !candidates.empty() && list3->bright)
-                {
-                    if (!L) return true;
+DEFINE_SORT_HANDLER(unit_sorters, layer_military, "/Positions/Candidates", military)
+{
+    auto &candidates = military->positions.candidates;
+    auto list3 = getLayerList(military, 2);
 
     PARSE_SPEC("units", parameters);
 
@@ -304,28 +297,17 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_cursor(&list3->cursor, order);
         reorder_vector(&candidates, order);
     }
+}
 
-                    return true;
-                }
-
-                return false;
-            }
-
-        default:
-            return false;
-        }
-    }
-    else if (auto profile = strict_virtual_cast<df::viewscreen_layer_workshop_profilest>(screen))
-    {
-        auto list1 = getLayerList(profile, 0);
-
-        if (!list1) return false;
-        if (!L) return true;
 
-        /*
+/*
  * Sort units in the workshop 'q'uery 'P'rofile modification screen.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, layer_workshop_profile, "/Unit", profile)
+{
+    auto list1 = getLayerList(profile, 0);
+
     PARSE_SPEC("units", parameters);
 
     if (compute_order(*pout, L, top, &order, profile->workers))
@@ -333,25 +315,12 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_cursor(&list1->cursor, order);
         reorder_vector(&profile->workers, order);
     }
+}
 
-        return true;
-    }
-    else if (auto nobles = strict_virtual_cast<df::viewscreen_layer_noblelistst>(screen))
-    {
-        switch (nobles->mode)
-        {
-        case df::viewscreen_layer_noblelistst::Appoint:
-            {
+DEFINE_SORT_HANDLER(unit_sorters, layer_noblelist, "/Appoint", nobles)
+{
     auto list2 = getLayerList(nobles, 1);
 
-                /*
-                 * Sort units in the appointment candidate list of the 'n'obles screen.
-                 */
-
-                if (list2)
-                {
-                    if (!L) return true;
-
     sort_null_first(parameters);
     PARSE_SPEC("units", parameters);
 
@@ -364,29 +333,14 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_cursor(&list2->cursor, order);
         reorder_vector(&nobles->candidates, order);
     }
+}
 
-                    return true;
-                }
-
-                return false;
-            }
-
-        default:
-            return false;
-        }
-    }
-    else if (auto animals = strict_virtual_cast<df::viewscreen_petst>(screen))
-    {
-        switch (animals->mode)
-        {
-        case df::viewscreen_petst::List:
-            {
-                if (!L) return true;
-
-                /*
+/*
  * Sort animal units in the Animal page of the 'z' status screen.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, pet, "/List", animals)
+{
     PARSE_SPEC("units", parameters);
 
     std::vector<df::unit*> units;
@@ -402,18 +356,14 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         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.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, pet, "/SelectTrainer", animals)
+{
     sort_null_first(parameters);
     PARSE_SPEC("units", parameters);
 
@@ -423,25 +373,16 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_vector(&animals->trainer_unit, order);
         reorder_vector(&animals->trainer_mode, order);
     }
+}
 
-                return true;
-            }
-
-        default:
-            return false;
-        }
-    }
-    else if (auto health = strict_virtual_cast<df::viewscreen_layer_overall_healthst>(screen))
-    {
-        auto list1 = getLayerList(health, 0);
-
-        if (!list1) return false;
-        if (!L) return true;
-
-        /*
+/*
  * Sort units in the Health page of the 'z' status screen.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, layer_overall_health, "/Units", health)
+{
+    auto list1 = getLayerList(health, 0);
+
     PARSE_SPEC("units", parameters);
 
     if (compute_order(*pout, L, top, &order, health->unit))
@@ -452,20 +393,14 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         reorder_vector(&health->bits2, order);
         reorder_vector(&health->bits3, order);
     }
+}
 
-        return true;
-    }
-    else if (strict_virtual_cast<df::viewscreen_dwarfmodest>(screen))
-    {
-        switch (ui->main.mode)
-        {
-        case ui_sidebar_mode::Burrows:
-            if (!L) return true;
-
-            /*
+/*
  * Sort burrow member candidate units in the 'w' sidebar mode.
  */
 
+DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/Burrows/AddUnits", screen)
+{
     PARSE_SPEC("units", parameters);
 
     if (compute_order(*pout, L, top, &order, ui->burrows.list_units))
@@ -474,34 +409,14 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         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.
+/*
+ * Sort building owner candidate units in the 'q' sidebar mode, or cage assignment.
  */
 
-                // TODO: better way
-                bool is_assign_owner = ((*ui_building_assign_type)[0] == -1);
-
-                if (is_assign_owner)
+DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/QueryBuilding/Some/Assign", screen)
+{
     sort_null_first(parameters);
 
     PARSE_SPEC("units", parameters);
@@ -511,29 +426,31 @@ static bool maybe_sort_units(color_ostream *pout, lua_State *L,
         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;
+/*
+ * Sort pen assignment candidate units in 'z'->'N'.
+ */
 
-        default:
-            return false;
-        }
+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,30 +461,28 @@ 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))
-        return CR_WRONG_USAGE;
+    Lua::StackUnwinder top(L);
 
-    return CR_OK;
-}
+    if (!prepare_sort(&out, L))
+        return CR_WRONG_USAGE;
 
-static bool maybe_sort_items(color_ostream *pout, lua_State *L,
-                             df::viewscreen *screen, vector<string> &parameters)
-{
-    Lua::StackUnwinder top(L);
+    auto focus = Gui::getFocusString(screen);
+    auto handler = findPrefixInMap(unit_sorters, focus);
 
-    if (!prepare_sort(pout, L))
-        return false;
+    if (!handler)
+        return CR_WRONG_USAGE;
+    else
+        handler(&out, L, top, screen, parameters);
 
-    std::vector<unsigned> order;
+    return CR_OK;
+}
 
-    if (auto trade = strict_virtual_cast<df::viewscreen_tradegoodsst>(screen))
-    {
-        if (!L) return true;
+static std::map<std::string, SortHandler> item_sorters;
 
+DEFINE_SORT_HANDLER(item_sorters, tradegoods, "/Items/Broker", trade)
+{
     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);
@@ -575,9 +490,12 @@ static bool maybe_sort_items(color_ostream *pout, lua_State *L,
         reorder_vector(&trade->broker_selected, order);
         reorder_vector(&trade->broker_count, order);
     }
-        }
-        else
-        {
+}
+
+DEFINE_SORT_HANDLER(item_sorters, tradegoods, "/Items/Trader", trade)
+{
+    PARSE_SPEC("items", parameters);
+
     if (compute_order(*pout, L, top, &order, trade->trader_items))
     {
         reorder_cursor(&trade->trader_cursor, order);
@@ -585,23 +503,13 @@ static bool maybe_sort_items(color_ostream *pout, lua_State *L,
         reorder_vector(&trade->trader_selected, order);
         reorder_vector(&trade->trader_count, order);
     }
-        }
+}
 
-        return true;
-    }
-    else if (auto bring = strict_virtual_cast<df::viewscreen_layer_assigntradest>(screen))
-    {
+DEFINE_SORT_HANDLER(item_sorters, layer_assigntrade, "/Items", bring)
+{
     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;
-
-        if (!L) return true;
 
     PARSE_SPEC("items", parameters);
 
@@ -616,17 +524,12 @@ static bool maybe_sort_items(color_ostream *pout, lua_State *L,
         reorder_cursor(&list2->cursor, order);
         reorder_vector(&vec, order);
     }
-
-        return true;
-    }
-    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;
 }