Support sorting units in many more ui contexts.

library/include/LuaTools.h

@@ -248,9 +248,15 @@ namespace DFHack {namespace Lua {
     }
 
     template<class T>
-    void PushVector(lua_State *state, const T &pvec)
+    void PushVector(lua_State *state, const T &pvec, bool addn = false)
     {
-        lua_createtable(state,pvec.size(),0);
+        lua_createtable(state,pvec.size(), addn?1:0);
+
+        if (addn)
+        {
+            lua_pushinteger(state, pvec.size());
+            lua_setfield(state, -2, "n");
+        }
 
         for (size_t i = 0; i < pvec.size(); i++)
         {
@@ -267,6 +273,26 @@ namespace DFHack {namespace Lua {
     DFHACK_EXPORT void MakeEvent(lua_State *state, void *key);
     DFHACK_EXPORT void InvokeEvent(color_ostream &out, lua_State *state, void *key, int num_args);
 
+    class StackUnwinder {
+        lua_State *state;
+        int top;
+    public:
+        StackUnwinder(lua_State *state, int bias = 0) : state(state), top(0) {
+            if (state) top = lua_gettop(state) - bias;
+        }
+        ~StackUnwinder() {
+            if (state) lua_settop(state, top);
+        }
+        operator int () { return top; }
+        int operator+ (int v) { return top + v; }
+        int operator- (int v) { return top + v; }
+        int operator[] (int v) { return top + v; }
+        StackUnwinder &operator += (int v) { top += v; return *this; }
+        StackUnwinder &operator -= (int v) { top += v; return *this; }
+        StackUnwinder &operator ++ () { top++; return *this; }
+        StackUnwinder &operator -- () { top--; return *this; }
+    };
+
     /**
      * Namespace for the common lua interpreter state.
      * All accesses must be done under CoreSuspender.

library/lua/utils.lua

@@ -49,6 +49,7 @@ function make_sort_order(data,ordering)
     -- Compute sort keys and comparators
     local keys = {}
     local cmps = {}
+    local size = data.n or #data
 
     for i=1,#ordering do
         local order = ordering[i]
@@ -58,7 +59,7 @@ function make_sort_order(data,ordering)
         elseif order.key then
             local kt = {}
             local kf = order.key
-            for j=1,#data do
+            for j=1,size do
                 if data[j] == nil then
                     kt[j] = nil
                 else
@@ -75,7 +76,7 @@ function make_sort_order(data,ordering)
 
     -- Make an order table
     local index = {}
-    for i=1,#data do
+    for i=1,size do
         index[i] = i
     end
 

library/xml

@@ -1 +1 @@
-Subproject commit e217d28c4800fadd3b37e153a363656dc7beb3e3
+Subproject commit 0ec109ce5643be2c3ebcd4d652aeafaf78f33da7

plugins/lua/sort/units.lua

@@ -4,6 +4,13 @@ local utils = require('utils')
 
 orders = orders or {}
 
+-- Relies on NULL being auto-translated to NULL, and then sorted
+orders.exists = {
+    key = function(unit)
+        return 1
+    end
+}
+
 orders.name = {
     key = function(unit)
         if unit.name.has_name then

plugins/sort.cpp

@@ -14,7 +14,13 @@
 #include "df/world.h"
 #include "df/viewscreen_joblistst.h"
 #include "df/viewscreen_unitlistst.h"
+#include "df/viewscreen_layer_militaryst.h"
+#include "df/viewscreen_layer_workshop_profilest.h"
+#include "df/viewscreen_layer_noblelistst.h"
+#include "df/viewscreen_layer_overall_healthst.h"
 #include "df/viewscreen_dwarfmodest.h"
+#include "df/viewscreen_petst.h"
+#include "df/layer_object_listst.h"
 
 #include "MiscUtils.h"
 
@@ -28,6 +34,9 @@ using namespace df::enums;
 
 using df::global::ui;
 using df::global::world;
+using df::global::ui_building_in_owner;
+using df::global::ui_building_item_cursor;
+using df::global::ui_owner_candidates;
 
 static bool unit_list_hotkey(df::viewscreen *top);
 
@@ -65,6 +74,22 @@ void reorder_vector(std::vector<T> *pvec, const std::vector<unsigned> &order)
         (*pvec)[i] = tmp[order[i]];
 }
 
+template<class T>
+void reorder_cursor(T *cursor, const std::vector<unsigned> &order)
+{
+    if (*cursor == 0)
+        return;
+
+    for (size_t i = 0; i < order.size(); i++)
+    {
+        if (unsigned(*cursor) == order[i])
+        {
+            *cursor = T(i);
+            break;
+        }
+    }
+}
+
 bool parse_ordering_spec(color_ostream &out, lua_State *L, std::string type, const std::vector<std::string> &params)
 {
     if (!lua_checkstack(L, params.size() + 2))
@@ -93,16 +118,18 @@ bool read_order(color_ostream &out, lua_State *L, std::vector<unsigned> *order,
 {
     std::vector<char> found;
 
+    Lua::StackUnwinder frame(L, 1);
+
     if (!lua_istable(L, -1))
     {
         out.printerr("Not a table returned as ordering.\n");
-        goto fail;
+        return false;
     }
 
     if (lua_rawlen(L, -1) != size)
     {
         out.printerr("Invalid ordering size: expected %d, actual %d\n", size, lua_rawlen(L, -1));
-        goto fail;
+        return false;
     }
 
     order->clear();
@@ -111,38 +138,34 @@ bool read_order(color_ostream &out, lua_State *L, std::vector<unsigned> *order,
 
     for (size_t i = 1; i <= size; i++)
     {
-        lua_rawgeti(L, -1, i);
+        lua_rawgeti(L, frame[1], i);
         int v = lua_tointeger(L, -1);
         lua_pop(L, 1);
 
         if (v < 1 || size_t(v) > size)
         {
             out.printerr("Order value out of range: %d\n", v);
-            goto fail;
+            return false;
         }
 
         if (found[v-1])
         {
             out.printerr("Duplicate order value: %d\n", v);
-            goto fail;
+            return false;
         }
 
         found[v-1] = 1;
         (*order)[i-1] = v-1;
     }
 
-    lua_pop(L, 1);
     return true;
-fail:
-    lua_pop(L, 1);
-    return false;
 }
 
 template<class T>
 bool compute_order(color_ostream &out, lua_State *L, int base, std::vector<unsigned> *order, const std::vector<T> &key)
 {
     lua_pushvalue(L, base+1);
-    Lua::PushVector(L, key);
+    Lua::PushVector(L, key, true);
     lua_pushvalue(L, base+2);
 
     if (!Lua::SafeCall(out, L, 2, 1))
@@ -151,90 +174,330 @@ bool compute_order(color_ostream &out, lua_State *L, int base, std::vector<unsig
     return read_order(out, L, order, key.size());
 }
 
+static bool ParseSpec(color_ostream &out, lua_State *L, const char *type, vector<string> &params)
+{
+    if (!parse_ordering_spec(out, L, "units", params))
+    {
+        out.printerr("Invalid ordering specification for %s.\n", type);
+        return false;
+    }
+
+    return true;
+}
+
+#define PARSE_SPEC(type, params) \
+    if (!ParseSpec(*pout, L, type, params)) return false;
 
-static bool unit_list_hotkey(df::viewscreen *screen)
+static void sort_null_first(vector<string> &parameters)
+{
+    vector_insert_at(parameters, 0, std::string("<exists"));
+}
+
+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 bool maybe_sort_units(color_ostream *pout, lua_State *L,
+                             df::viewscreen *screen, vector<string> &parameters)
 {
-    if (strict_virtual_cast<df::viewscreen_unitlistst>(screen))
+    Lua::StackUnwinder top(L);
+
+    if (L)
+    {
+        if (!Lua::PushModulePublic(*pout, L, "plugins.sort", "make_sort_order"))
+        {
+            pout->printerr("Cannot access the sorter function.\n");
+            return false;
+        }
+    }
+
+    std::vector<unsigned> order;
+
+    if (auto units = strict_virtual_cast<df::viewscreen_unitlistst>(screen))
+    {
+        if (!L) return true;
+
+        /*
+         * Sort units in the 'u'nit list screen.
+         */
+
+        PARSE_SPEC("units", parameters);
+
+        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);
+        }
+
         return true;
-    if (strict_virtual_cast<df::viewscreen_joblistst>(screen))
+    }
+    else if (auto jobs = strict_virtual_cast<df::viewscreen_joblistst>(screen))
+    {
+        if (!L) return true;
+
+        /*
+         * Sort units in the 'j'ob list screen.
+         */
+
+        PARSE_SPEC("units", parameters);
+
+        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);
+        }
+
         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 (strict_virtual_cast<df::viewscreen_dwarfmodest>(screen))
+                if (list3 && !candidates.empty() && list3->bright)
                 {
-        using namespace df::enums::ui_sidebar_mode;
+                    if (!L) return true;
 
-        switch (ui->main.mode)
+                    PARSE_SPEC("units", parameters);
+
+                    if (compute_order(*pout, L, top, &order, candidates))
                     {
-        case Burrows:
-            return ui->burrows.in_add_units_mode;
+                        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);
 
-    return false;
-}
+        if (!list1) return false;
+        if (!L) return true;
 
-static command_result sort_units(color_ostream &out, vector <string> &parameters)
-{
-    if (parameters.empty())
-        return CR_WRONG_USAGE;
+        /*
+         * Sort units in the workshop 'q'uery 'P'rofile modification screen.
+         */
 
-    auto L = Lua::Core::State;
-    int top = lua_gettop(L);
+        PARSE_SPEC("units", parameters);
 
-    if (!Lua::Core::PushModulePublic(out, "plugins.sort", "make_sort_order"))
+        if (compute_order(*pout, L, top, &order, profile->workers))
         {
-        out.printerr("Cannot access the sorter function.\n");
-        return CR_WRONG_USAGE;
+            reorder_cursor(&list1->cursor, order);
+            reorder_vector(&profile->workers, order);
         }
 
-    if (!parse_ordering_spec(out, L, "units", parameters))
+        return true;
+    }
+    else if (auto nobles = strict_virtual_cast<df::viewscreen_layer_noblelistst>(screen))
     {
-        out.printerr("Invalid unit ordering specification.\n");
-        lua_settop(L, top);
-        return CR_WRONG_USAGE;
+        switch (nobles->mode)
+        {
+        case df::viewscreen_layer_noblelistst::Appoint:
+            {
+                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);
+
+                    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, units))
+                    {
+                        reorder_cursor(&list2->cursor, order);
+                        reorder_vector(&nobles->candidates, order);
                     }
 
-    auto screen = Core::getInstance().getTopViewscreen();
-    std::vector<unsigned> order;
+                    return true;
+                }
 
-    if (auto units = strict_virtual_cast<df::viewscreen_unitlistst>(screen))
+                return false;
+            }
+
+        default:
+            return false;
+        }
+    }
+    else if (auto animals = strict_virtual_cast<df::viewscreen_petst>(screen))
+    {
+        switch (animals->mode)
+        {
+        case df::viewscreen_petst::List:
             {
-        for (int i = 0; i < 4; i++)
+                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 (compute_order(out, L, top, &order, units->units[i]))
+                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_vector(&units->units[i], order);
-                reorder_vector(&units->jobs[i], order);
+                    reorder_cursor(&animals->trainer_cursor, order);
+                    reorder_vector(&animals->trainer_unit, order);
+                    reorder_vector(&animals->trainer_mode, order);
                 }
+
+                return true;
             }
+
+        default:
+            return false;
         }
-    else if (auto jobs = strict_virtual_cast<df::viewscreen_joblistst>(screen))
+    }
+    else if (auto health = strict_virtual_cast<df::viewscreen_layer_overall_healthst>(screen))
     {
-        if (compute_order(out, L, top, &order, jobs->units))
+        auto list1 = getLayerList(health, 0);
+
+        if (!list1) return false;
+        if (!L) return true;
+
+        /*
+         * Sort units in the Health page of the 'z' status screen.
+         */
+
+        PARSE_SPEC("units", parameters);
+
+        if (compute_order(*pout, L, top, &order, health->unit))
         {
-            reorder_vector(&jobs->units, order);
-            reorder_vector(&jobs->jobs, order);
+            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);
         }
+
+        return true;
     }
     else if (strict_virtual_cast<df::viewscreen_dwarfmodest>(screen))
     {
         switch (ui->main.mode)
         {
         case ui_sidebar_mode::Burrows:
-            if (compute_order(out, L, top, &order, ui->burrows.list_units))
+            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);
             }
-            break;
+
+            return true;
+
+        case ui_sidebar_mode::QueryBuilding:
+            if (ui_building_in_owner && *ui_building_in_owner &&
+                ui_owner_candidates && ui_building_item_cursor)
+            {
+                if (!L) return true;
+
+                /*
+                 * Sort building owner candidate units in the 'q' sidebar mode.
+                 */
+
+                sort_null_first(parameters);
+                PARSE_SPEC("units", parameters);
+
+                if (compute_order(*pout, L, top, &order, *ui_owner_candidates))
+                {
+                    reorder_cursor(ui_building_item_cursor, order);
+                    reorder_vector(ui_owner_candidates, order);
+                }
+
+                return true;
+            }
+            return false;
 
         default:
-            break;;
+            return false;
         }
     }
+    else
+        return false;
+}
+
+static bool unit_list_hotkey(df::viewscreen *screen)
+{
+    vector<string> dummy;
+    return maybe_sort_units(NULL, NULL, screen, dummy);
+}
+
+static command_result sort_units(color_ostream &out, vector <string> &parameters)
+{
+    if (parameters.empty())
+        return CR_WRONG_USAGE;
+
+    auto L = Lua::Core::State;
+    auto screen = Core::getInstance().getTopViewscreen();
+
+    if (!maybe_sort_units(&out, L, screen, parameters))
+        return CR_WRONG_USAGE;
 
-    lua_settop(L, top);
     return CR_OK;
 }