Add a hotkey command that sorts units in lists using lua comparators.

library/DataStaticsFields.cpp

@@ -28,6 +28,7 @@ namespace df {
     NUMBER_IDENTITY_TRAITS(int64_t);
     NUMBER_IDENTITY_TRAITS(uint64_t);
     NUMBER_IDENTITY_TRAITS(float);
+    NUMBER_IDENTITY_TRAITS(double);
 
     bool_identity identity_traits<bool>::identity;
     stl_string_identity identity_traits<std::string>::identity;

library/LuaApi.cpp

@@ -615,6 +615,7 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, clearBurrowMembers),
     WRAPM(Units, isInBurrow),
     WRAPM(Units, setInBurrow),
+    WRAPM(Units, getAge),
     { NULL, NULL }
 };
 

library/include/DataIdentity.h

@@ -447,6 +447,7 @@ namespace df
     NUMBER_IDENTITY_TRAITS(int64_t);
     NUMBER_IDENTITY_TRAITS(uint64_t);
     NUMBER_IDENTITY_TRAITS(float);
+    NUMBER_IDENTITY_TRAITS(double);
 
     template<> struct DFHACK_EXPORT identity_traits<bool> {
         static bool_identity identity;

library/include/modules/Units.h

@@ -214,6 +214,8 @@ DFHACK_EXPORT void clearBurrowMembers(df::burrow *burrow);
 DFHACK_EXPORT bool isInBurrow(df::unit *unit, df::burrow *burrow);
 DFHACK_EXPORT void setInBurrow(df::unit *unit, df::burrow *burrow, bool enable);
 
+DFHACK_EXPORT double getAge(df::unit *unit);
+
 }
 }
 #endif

library/lua/utils.lua

@@ -0,0 +1,106 @@
+local _ENV = mkmodule('utils')
+
+-- Comparator function
+function compare(a,b)
+    if a < b then
+        return -1
+    elseif a > b then
+        return 1
+    else
+        return 0
+    end
+end
+
+-- Sort strings; compare empty last
+function compare_name(a,b)
+    if a == '' then
+        if b == '' then
+            return 0
+        else
+            return 1
+        end
+    elseif b == '' then
+        return -1
+    else
+        return compare(a,b)
+    end
+end
+
+--[[
+    Sort items in data according to ordering.
+
+    Each ordering spec is a table with possible fields:
+
+    * key = function(value)
+        Computes comparison key from a data value. Not called on nil.
+    * key_table = function(data)
+        Computes a key table from the data table in one go.
+    * compare = function(a,b)
+        Comparison function. Defaults to compare above.
+        Called on non-nil keys; nil sorts last.
+
+    Returns a table of integer indices into data.
+--]]
+function make_sort_order(data,ordering)
+    -- Compute sort keys and comparators
+    local keys = {}
+    local cmps = {}
+
+    for i=1,#ordering do
+        local order = ordering[i]
+
+        if order.key_table then
+            keys[i] = order.key_table(data)
+        elseif order.key then
+            local kt = {}
+            local kf = order.key
+            for j=1,#data do
+                if data[j] == nil then
+                    kt[j] = nil
+                else
+                    kt[j] = kf(data[j])
+                end
+            end
+            keys[i] = kt
+        else
+            keys[i] = data
+        end
+
+        cmps[i] = order.compare or compare
+    end
+
+    -- Make an order table
+    local index = {}
+    for i=1,#data do
+        index[i] = i
+    end
+
+    -- Sort the ordering table
+    table.sort(index, function(ia,ib)
+        for i=1,#keys do
+            local ka = keys[i][ia]
+            local kb = keys[i][ib]
+
+            -- Sort nil keys to the end
+            if ka == nil then
+                if kb ~= nil then
+                    return false
+                end
+            elseif kb == nil then
+                return true
+            else
+                local cmpv = cmps[i](ka,kb)
+                if cmpv < 0 then
+                    return true
+                elseif cmpv > 0 then
+                    return false
+                end
+            end
+        end
+        return ia < ib -- this should ensure stable sort
+    end)
+
+    return index
+end
+
+return _ENV

library/modules/Units.cpp

@@ -734,3 +734,19 @@ void DFHack::Units::setInBurrow(df::unit *unit, df::burrow *burrow, bool enable)
     }
 }
 
+double DFHack::Units::getAge(df::unit *unit)
+{
+    using df::global::cur_year;
+    using df::global::cur_year_tick;
+
+    CHECK_NULL_POINTER(unit);
+
+    if (!cur_year || !cur_year_tick)
+        return -1;
+
+    double year_ticks = 403200.0;
+    double birth_time = unit->relations.birth_year + unit->relations.birth_time/year_ticks;
+    double cur_time = *cur_year + *cur_year_tick / year_ticks;
+
+    return cur_time - birth_time;
+}

library/xml

@@ -1 +1 @@
-Subproject commit ee172f69f613716c8d740bbd22054f48b1a22d5f
+Subproject commit e217d28c4800fadd3b37e153a363656dc7beb3e3

plugins/CMakeLists.txt

@@ -106,6 +106,7 @@ if (BUILD_SUPPORTED)
     DFHACK_PLUGIN(regrass regrass.cpp)
     # this one exports functions to lua
     DFHACK_PLUGIN(burrows burrows.cpp LINK_LIBRARIES lua)
+    DFHACK_PLUGIN(sort sort.cpp LINK_LIBRARIES lua)
     # not yet. busy with other crud again...
     #DFHACK_PLUGIN(versionosd versionosd.cpp)
 endif()

plugins/lua/sort.lua

@@ -0,0 +1,32 @@
+local _ENV = mkmodule('plugins.sort')
+
+local utils = require('utils')
+local units = require('plugins.sort.units')
+
+orders = orders or {}
+orders.units = units.orders
+
+function parse_ordering_spec(type,...)
+    local group = orders[type]
+    if group == nil then
+        dfhack.printerr('Invalid ordering class: '..tostring(type))
+        return nil
+    end
+
+    local specs = table.pack(...)
+    local rv = { }
+    for _,spec in ipairs(specs) do
+        local cm = group[spec]
+        if cm == nil then
+            dfhack.printerr('Unknown order for '..type..': '..tostring(spec))
+            return nil
+        end
+        rv[#rv+1] = cm
+    end
+
+    return rv
+end
+
+make_sort_order = utils.make_sort_order
+
+return _ENV

plugins/lua/sort/units.lua

@@ -0,0 +1,20 @@
+local _ENV = mkmodule('plugins.sort.units')
+
+local utils = require('utils')
+
+orders = orders or {}
+
+orders.name = {
+    key = function(unit)
+        return dfhack.TranslateName(unit.name)
+    end,
+    compare = utils.compare_name
+}
+
+orders.age = {
+    key = function(unit)
+        return dfhack.units.getAge(unit)
+    end
+}
+
+return _ENV

plugins/sort.cpp

@@ -0,0 +1,235 @@
+#include "Core.h"
+#include "Console.h"
+#include "Export.h"
+#include "PluginManager.h"
+
+#include "modules/Gui.h"
+#include "modules/Translation.h"
+#include "modules/Units.h"
+
+#include "LuaTools.h"
+
+#include "DataDefs.h"
+#include "df/ui.h"
+#include "df/world.h"
+#include "df/viewscreen_joblistst.h"
+#include "df/viewscreen_unitlistst.h"
+#include "df/viewscreen_dwarfmodest.h"
+
+#include "MiscUtils.h"
+
+#include <stdlib.h>
+
+using std::vector;
+using std::string;
+using std::endl;
+using namespace DFHack;
+using namespace df::enums;
+
+using df::global::ui;
+using df::global::world;
+
+static bool unit_list_hotkey(df::viewscreen *top);
+
+static command_result sort_units(color_ostream &out, vector <string> & parameters);
+
+DFHACK_PLUGIN("sort");
+
+DFhackCExport command_result plugin_init (color_ostream &out, std::vector <PluginCommand> &commands)
+{
+    commands.push_back(PluginCommand(
+        "sort-units", "Sort the visible unit list.", sort_units, unit_list_hotkey,
+        "  sort-units filter...\n"
+        "    Sort the unit list using the given sequence of comparisons.\n"
+    ));
+    return CR_OK;
+}
+
+DFhackCExport command_result plugin_shutdown ( color_ostream &out )
+{
+    return CR_OK;
+}
+
+template<class T>
+void reorder_vector(std::vector<T> *pvec, const std::vector<unsigned> &order)
+{
+    assert(pvec->size() == order.size());
+
+    std::vector<T> tmp(*pvec);
+    for (size_t i = 0; i < order.size(); i++)
+        (*pvec)[i] = tmp[order[i]];
+}
+
+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))
+        return false;
+
+    if (!Lua::PushModulePublic(out, L, "plugins.sort", "parse_ordering_spec"))
+        return false;
+
+    Lua::Push(L, type);
+    for (size_t i = 0; i < params.size(); i++)
+        Lua::Push(L, params[i]);
+
+    if (!Lua::SafeCall(out, L, params.size()+1, 1))
+        return false;
+
+    if (!lua_istable(L, -1))
+    {
+        lua_pop(L, 1);
+        return false;
+    }
+
+    return true;
+}
+
+bool read_order(color_ostream &out, lua_State *L, std::vector<unsigned> *order, size_t size)
+{
+    std::vector<char> found;
+
+    if (!lua_istable(L, -1))
+    {
+        out.printerr("Not a table returned as ordering.\n");
+        goto fail;
+    }
+
+    if (lua_rawlen(L, -1) != size)
+    {
+        out.printerr("Invalid ordering size: expected %d, actual %d\n", size, lua_rawlen(L, -1));
+        goto fail;
+    }
+
+    order->clear();
+    order->resize(size);
+    found.resize(size);
+
+    for (size_t i = 1; i <= size; i++)
+    {
+        lua_rawgeti(L, -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;
+        }
+
+        if (found[v-1])
+        {
+            out.printerr("Duplicate order value: %d\n", v);
+            goto fail;
+        }
+
+        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_pushvalue(L, base+2);
+
+    if (!Lua::SafeCall(out, L, 2, 1))
+        return false;
+
+    return read_order(out, L, order, key.size());
+}
+
+
+static bool unit_list_hotkey(df::viewscreen *screen)
+{
+    if (strict_virtual_cast<df::viewscreen_unitlistst>(screen))
+        return true;
+    if (strict_virtual_cast<df::viewscreen_joblistst>(screen))
+        return true;
+
+    if (strict_virtual_cast<df::viewscreen_dwarfmodest>(screen))
+    {
+        using namespace df::enums::ui_sidebar_mode;
+
+        switch (ui->main.mode)
+        {
+        case Burrows:
+            return ui->burrows.in_add_units_mode;
+        default:
+            return false;
+        }
+    }
+
+    return false;
+}
+
+static command_result sort_units(color_ostream &out, vector <string> &parameters)
+{
+    if (parameters.empty())
+        return CR_WRONG_USAGE;
+
+    auto L = Lua::Core::State;
+    int top = lua_gettop(L);
+
+    if (!Lua::Core::PushModulePublic(out, "plugins.sort", "make_sort_order"))
+    {
+        out.printerr("Cannot access the sorter function.\n");
+        return CR_WRONG_USAGE;
+    }
+
+    if (!parse_ordering_spec(out, L, "units", parameters))
+    {
+        out.printerr("Invalid unit ordering specification.\n");
+        lua_settop(L, top);
+        return CR_WRONG_USAGE;
+    }
+
+    auto screen = Core::getInstance().getTopViewscreen();
+    std::vector<unsigned> order;
+
+    if (auto units = strict_virtual_cast<df::viewscreen_unitlistst>(screen))
+    {
+        for (int i = 0; i < 4; i++)
+        {
+            if (compute_order(out, L, top, &order, units->units[i]))
+            {
+                reorder_vector(&units->units[i], order);
+                reorder_vector(&units->jobs[i], order);
+            }
+        }
+    }
+    else if (auto jobs = strict_virtual_cast<df::viewscreen_joblistst>(screen))
+    {
+        if (compute_order(out, L, top, &order, jobs->units))
+        {
+            reorder_vector(&jobs->units, order);
+            reorder_vector(&jobs->jobs, order);
+        }
+    }
+    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))
+            {
+                reorder_vector(&ui->burrows.list_units, order);
+                reorder_vector(&ui->burrows.sel_units, order);
+            }
+            break;
+
+        default:
+            break;;
+        }
+    }
+
+    lua_settop(L, top);
+    return CR_OK;
+}