From 3282ac3db216ef43cab5744631b57ed05bcf8a12 Mon Sep 17 00:00:00 2001 From: Alexander Gavrilov Date: Sat, 21 Apr 2012 15:43:52 +0400 Subject: [PATCH] Add a hotkey command that sorts units in lists using lua comparators. --- library/DataStaticsFields.cpp | 1 + library/LuaApi.cpp | 1 + library/include/DataIdentity.h | 1 + library/include/modules/Units.h | 2 + library/lua/utils.lua | 106 ++++++++++++++ library/modules/Units.cpp | 16 +++ library/xml | 2 +- plugins/CMakeLists.txt | 1 + plugins/lua/sort.lua | 32 +++++ plugins/lua/sort/units.lua | 20 +++ plugins/sort.cpp | 235 ++++++++++++++++++++++++++++++++ 11 files changed, 416 insertions(+), 1 deletion(-) create mode 100644 library/lua/utils.lua create mode 100644 plugins/lua/sort.lua create mode 100644 plugins/lua/sort/units.lua create mode 100644 plugins/sort.cpp diff --git a/library/DataStaticsFields.cpp b/library/DataStaticsFields.cpp index a34938f95..79aa3bcf1 100644 --- a/library/DataStaticsFields.cpp +++ b/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::identity; stl_string_identity identity_traits::identity; diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp index ccd057c9c..00192c058 100644 --- a/library/LuaApi.cpp +++ b/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 } }; diff --git a/library/include/DataIdentity.h b/library/include/DataIdentity.h index e4197a506..dcd0ae979 100644 --- a/library/include/DataIdentity.h +++ b/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 { static bool_identity identity; diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h index e093ed1ef..775cc5181 100644 --- a/library/include/modules/Units.h +++ b/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 diff --git a/library/lua/utils.lua b/library/lua/utils.lua new file mode 100644 index 000000000..d884f2f62 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp index 55bd98dec..e9699edbe 100644 --- a/library/modules/Units.cpp +++ b/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; +} diff --git a/library/xml b/library/xml index ee172f69f..e217d28c4 160000 --- a/library/xml +++ b/library/xml @@ -1 +1 @@ -Subproject commit ee172f69f613716c8d740bbd22054f48b1a22d5f +Subproject commit e217d28c4800fadd3b37e153a363656dc7beb3e3 diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index 4bfb392a2..70b40cedb 100644 --- a/plugins/CMakeLists.txt +++ b/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() diff --git a/plugins/lua/sort.lua b/plugins/lua/sort.lua new file mode 100644 index 000000000..2318d7e07 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/plugins/lua/sort/units.lua b/plugins/lua/sort/units.lua new file mode 100644 index 000000000..872e49429 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/plugins/sort.cpp b/plugins/sort.cpp new file mode 100644 index 000000000..a373d54c3 --- /dev/null +++ b/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 + +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 & parameters); + +DFHACK_PLUGIN("sort"); + +DFhackCExport command_result plugin_init (color_ostream &out, std::vector &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 +void reorder_vector(std::vector *pvec, const std::vector &order) +{ + assert(pvec->size() == order.size()); + + std::vector 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 ¶ms) +{ + 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 *order, size_t size) +{ + std::vector 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 +bool compute_order(color_ostream &out, lua_State *L, int base, std::vector *order, const std::vector &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(screen)) + return true; + if (strict_virtual_cast(screen)) + return true; + + if (strict_virtual_cast(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 ¶meters) +{ + 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 order; + + if (auto units = strict_virtual_cast(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(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(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; +}