#include
#include
#include "Debug.h"
#include "LuaTools.h"
#include "PluginManager.h"
#include "PluginLua.h"
#include "modules/Gui.h"
#include "modules/Units.h"
#include "df/gamest.h"
#include "df/unit.h"
#include "df/widget_unit_list.h"
#include "df/world.h"
using std::vector;
using std::string;
using namespace DFHack;
DFHACK_PLUGIN("sort");
REQUIRE_GLOBAL(game);
REQUIRE_GLOBAL(world);
namespace DFHack {
DBG_DECLARE(sort, log, DebugCategory::LINFO);
}
using item_or_unit = std::variant<:unit df::item>;
using filter_func = bool(item_or_unit);
using filter_vec_type = std::vector<:function>>;
// recreated here since our autogenerated df::sort_entry lacks template params
struct sort_entry {
std::function fn;
string ident;
};
static const string DFHACK_SORT_IDENT = "dfhack_sort";
//
// filter logic
//
static bool do_filter(const char *module_name, const char *fn_name, const item_or_unit &elem) {
if (std::holds_alternative<:item>(elem)) return true;
auto unit = std::get<:unit>(elem);
bool ret = true;
color_ostream &out = Core::getInstance().getConsole();
Lua::CallLuaModuleFunction(out, module_name, fn_name, std::make_tuple(unit),
1, [&](lua_State *L){
ret = lua_toboolean(L, 1);
}
);
TRACE(log).print("filter result for {}: {}\n", Units::getReadableName(unit), ret);
return !ret;
}
static bool do_squad_filter(item_or_unit elem) {
return do_filter("plugins.sort", "do_squad_filter", elem);
}
static bool do_justice_filter(item_or_unit elem) {
return do_filter("plugins.sort.info", "do_justice_filter", elem);
}
static bool do_work_animal_assignment_filter(item_or_unit elem) {
return do_filter("plugins.sort.info", "do_work_animal_assignment_filter", elem);
}
static int32_t our_filter_idx(filter_func* filter, df::widget_unit_list* unitlist)
{
int32_t idx = 0;
filter_vec_type* filter_vec = reinterpret_cast(&unitlist->filter_func);
TRACE(log).print("probing for our filter function\n");
for (auto& fn : *filter_vec)
{
auto t = fn.target();
if (t && *t == filter)
{
TRACE(log).print("found our filter function at idx {}\n", idx);
return idx;
}
++idx;
}
return -1;
}
static df::widget_unit_list* get_squad_unit_list() {
return virtual_cast<:widget_unit_list>(
Gui::getWidget(&game->main_interface.unit_selector, "Unit selector"));
}
static df::widget_container * get_justice_panel(const char *which) {
auto tabs = virtual_cast<:widget_container>(
Gui::getWidget(&game->main_interface.info.justice, "Tabs"));
if (!tabs) return NULL;
auto open_cases = virtual_cast<:widget_container>(Gui::getWidget(tabs, which));
if (!open_cases) return NULL;
return virtual_cast<:widget_container>(Gui::getWidget(open_cases, "Right panel"));
}
static df::widget_unit_list * get_interrogate_unit_list(const char *which) {
auto right_panel = get_justice_panel(which);
if (!right_panel) return NULL;
return virtual_cast<:widget_unit_list>(Gui::getWidget(right_panel, "Interrogate"));
}
static df::widget_unit_list * get_convict_unit_list(const char *which) {
auto right_panel = get_justice_panel(which);
if (!right_panel) return NULL;
return virtual_cast<:widget_unit_list>(Gui::getWidget(right_panel, "Convict"));
}
static df::widget_unit_list * get_work_animal_assignment_unit_list() {
auto tabs = virtual_cast<:widget_container>(
Gui::getWidget(&game->main_interface.info.creatures, "Tabs"));
if (!tabs) return NULL;
auto pets = virtual_cast<:widget_container>(Gui::getWidget(tabs, "Pets/Livestock"));
if (!pets) return NULL;
return virtual_cast<:widget_unit_list>(Gui::getWidget(pets, "Hunting assignment"));
}
//
// sorting logic
//
static bool sort_proxy(const item_or_unit &a, const item_or_unit &b) {
if (std::holds_alternative<:item>(a) || std::holds_alternative<:item>(b))
return true;
bool ret = true;
color_ostream &out = Core::getInstance().getConsole();
Lua::CallLuaModuleFunction(out, "plugins.sort", "do_sort",
std::make_tuple(std::get<:unit>(a), std::get<:unit>(b)),
1, [&](lua_State *L){
ret = lua_toboolean(L, 1);
}
);
return ret;
}
static sort_entry do_sort{
sort_proxy,
DFHACK_SORT_IDENT
};
int32_t our_sort_idx(const std::vector &sorting_by) {
for (size_t i = 0; i < sorting_by.size(); ++i) {
if (sorting_by[i].ident == DFHACK_SORT_IDENT) {
return (int32_t)i;
}
}
return -1;
}
//
// plugin logic
//
DFhackCExport command_result plugin_init(color_ostream &out, vector &commands) {
return CR_OK;
}
static void remove_filter_function(color_ostream &out, filter_func* filter, const char *which, df::widget_unit_list *unitlist) {
int32_t idx = our_filter_idx(filter, unitlist);
if (idx >= 0) {
DEBUG(log,out).print("removing {} filter function\n", which);
filter_vec_type *filter_vec = reinterpret_cast(&unitlist->filter_func);
vector_erase_at(*filter_vec, idx);
}
}
static void remove_sort_function(color_ostream &out, const char *which, df::widget_unit_list *unitlist) {
std::vector *sorting_by = reinterpret_cast<:vector> *>(&unitlist->sorting_by);
int32_t idx = our_sort_idx(*sorting_by);
if (idx >= 0) {
DEBUG(log).print("removing {} sort function\n", which);
vector_erase_at(*sorting_by, idx);
}
}
DFhackCExport command_result plugin_shutdown(color_ostream &out) {
if (auto unitlist = get_squad_unit_list()) {
remove_filter_function(out, do_squad_filter, "squad", unitlist);
remove_sort_function(out, "squad", unitlist);
}
if (auto unitlist = get_interrogate_unit_list("Open cases"))
remove_filter_function(out, do_justice_filter, "open cases interrogate", unitlist);
if (auto unitlist = get_interrogate_unit_list("Cold cases"))
remove_filter_function(out, do_justice_filter, "cold cases interrogate", unitlist);
if (auto unitlist = get_convict_unit_list("Open cases"))
remove_filter_function(out, do_justice_filter, "open cases convict", unitlist);
if (auto unitlist = get_convict_unit_list("Cold cases"))
remove_filter_function(out, do_justice_filter, "cold cases convict", unitlist);
if (auto unitlist = get_work_animal_assignment_unit_list())
remove_filter_function(out, do_work_animal_assignment_filter, "work animal assignment", unitlist);
return CR_OK;
}
//
// Lua API
//
static void sort_set_squad_filter_fn(color_ostream &out) {
auto unitlist = get_squad_unit_list();
if (unitlist && our_filter_idx(do_squad_filter, unitlist) == -1) {
DEBUG(log).print("adding squad filter function\n");
auto filter_vec = reinterpret_cast(&unitlist->filter_func);
filter_vec->emplace_back(do_squad_filter);
DEBUG(log).print("clearing partitions\n"); // removes sorting other squads to end
auto partitions_vec = reinterpret_cast(&unitlist->partitions);
partitions_vec->clear();
unitlist->sort_flags.bits.NEEDS_RESORTED = true;
}
}
static void sort_set_justice_filter_fn(color_ostream &out, df::widget_unit_list *unitlist) {
if (unitlist && our_filter_idx(do_justice_filter, unitlist) == -1) {
DEBUG(log).print("adding justice filter function\n");
auto filter_vec = reinterpret_cast(&unitlist->filter_func);
filter_vec->emplace_back(do_justice_filter);
unitlist->sort_flags.bits.NEEDS_RESORTED = true;
}
}
static void sort_set_work_animal_assignment_filter_fn(color_ostream &out, df::widget_unit_list *unitlist) {
if (unitlist && our_filter_idx(do_work_animal_assignment_filter, unitlist) == -1) {
DEBUG(log).print("adding work animal assignment filter function\n");
auto filter_vec = reinterpret_cast(&unitlist->filter_func);
filter_vec->emplace_back(do_work_animal_assignment_filter);
unitlist->sort_flags.bits.NEEDS_RESORTED = true;
}
}
static void sort_set_sort_fn(color_ostream &out) {
auto unitlist = get_squad_unit_list();
if (!unitlist)
return;
DEBUG(log).print("adding squad sort function\n");
std::vector *sorting_by = reinterpret_cast<:vector> *>(&unitlist->sorting_by);
sorting_by->clear();
sorting_by->emplace_back(do_sort);
unitlist->sort_flags.bits.NEEDS_RESORTED = true;
}
static bool sort_get_sort_active(color_ostream &out) {
auto unitlist = get_squad_unit_list();
if (!unitlist)
return false;
std::vector *sorting_by = reinterpret_cast<:vector> *>(&unitlist->sorting_by);
return our_sort_idx(*sorting_by) >= 0;
}
static bool sort_is_interviewed(color_ostream &out, df::unit *unit) {
auto flag_map = reinterpret_cast<:unordered_map> *>(&game->main_interface.info.justice.crimeflag);
if (!flag_map->contains(unit))
return false;
return (*flag_map)[unit].bits.ALREADY_INTERVIEWED;
}
DFHACK_PLUGIN_LUA_FUNCTIONS{
DFHACK_LUA_FUNCTION(sort_set_squad_filter_fn),
DFHACK_LUA_FUNCTION(sort_set_justice_filter_fn),
DFHACK_LUA_FUNCTION(sort_set_work_animal_assignment_filter_fn),
DFHACK_LUA_FUNCTION(sort_set_sort_fn),
DFHACK_LUA_FUNCTION(sort_get_sort_active),
DFHACK_LUA_FUNCTION(sort_is_interviewed),
DFHACK_LUA_END
};