diff --git a/docs/changelog.txt b/docs/changelog.txt index bf8182c73..814bb3909 100644 --- a/docs/changelog.txt +++ b/docs/changelog.txt @@ -42,11 +42,13 @@ changelog.txt uses a syntax similar to RST, with a few special sequences: ## Documentation ## API +- ``Units::getUnitByNobleRole``, ``Units::getUnitsByNobleRole``: unit lookup API by role ## Internals ## Lua - ``dfhack.items.markForTrade``: new API for marking items for trade +- ``dfhack.units.getUnitByNobleRole``, ``dfhack.units.getUnitsByNobleRole``: unit lookup API by role ## Removed diff --git a/docs/dev/Lua API.rst b/docs/dev/Lua API.rst index 33e27be2d..57e63dabb 100644 --- a/docs/dev/Lua API.rst +++ b/docs/dev/Lua API.rst @@ -1438,10 +1438,25 @@ Units module Note that ``pos2xyz()`` cannot currently be used to convert coordinate objects to the arguments required by this function. +* ``dfhack.units.getUnitByNobleRole(role_name)`` + + Returns the unit assigned to the given noble role, if any. ``role_name`` must + be one of the position codes associated with the active fort government. + Normally, this includes: ``MILITIA_COMMANDER``, ``MILITIA_CAPTAIN``, + ``SHERIFF``, ``CAPTAIN_OF_THE_GUARD``, ``EXPEDITION_LEADER``, ``MAYOR``, + ``MANAGER``, ``CHIEF_MEDICAL_DWARF``, ``BROKER``, ``BOOKKEEPER``, + ``CHAMPION``, ``HAMMERER``, ``DUNGEON_MASTER``, and ``MESSENGER``. Note that + if more than one unit has the role, only the first will be returned. See + ``getUnitsByNobleRole`` below for retrieving all units with a particular role. + +* ``dfhack.units.getUnitsByNobleRole(role_name)`` + + Returns a list of units (possibly empty) assigned to the given noble role. + * ``dfhack.units.getCitizens([ignore_sanity])`` - Returns a table (list) of all citizens, which you would otherwise have to loop over all - units in world and test against ``isCitizen()`` to discover. + Returns a table (list) of all citizens, which you would otherwise have to + loop over all units in world and test against ``isCitizen()`` to discover. * ``dfhack.units.teleport(unit, pos)`` diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp index 5d9411434..65a899814 100644 --- a/library/LuaApi.cpp +++ b/library/LuaApi.cpp @@ -1833,6 +1833,7 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = { WRAPM(Units, multiplyGroupActionTimers), WRAPM(Units, setActionTimers), WRAPM(Units, setGroupActionTimers), + WRAPM(Units, getUnitByNobleRole), { NULL, NULL } }; @@ -1921,6 +1922,14 @@ static int units_getCitizens(lua_State *L) { return 0; } +static int units_getUnitsByNobleRole(lua_State *L) { + std::string role_name = lua_tostring(L, -1); + std::vector units; + Units::getUnitsByNobleRole(units, role_name); + Lua::PushVector(L, units); + return 1; +} + static int units_getStressCutoffs(lua_State *L) { lua_newtable(L); @@ -1935,6 +1944,7 @@ static const luaL_Reg dfhack_units_funcs[] = { { "getNoblePositions", units_getNoblePositions }, { "getUnitsInBox", units_getUnitsInBox }, { "getCitizens", units_getCitizens }, + { "getUnitsByNobleRole", units_getUnitsByNobleRole}, { "getStressCutoffs", units_getStressCutoffs }, { NULL, NULL } }; diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h index 4fd9246aa..3c56d0890 100644 --- a/library/include/modules/Units.h +++ b/library/include/modules/Units.h @@ -148,6 +148,8 @@ DFHACK_EXPORT df::unit *getUnit(const int32_t index); DFHACK_EXPORT bool getUnitsInBox(std::vector &units, int16_t x1, int16_t y1, int16_t z1, int16_t x2, int16_t y2, int16_t z2); +DFHACK_EXPORT bool getUnitsByNobleRole(std::vector &units, std::string noble); +DFHACK_EXPORT df::unit *getUnitByNobleRole(std::string noble); DFHACK_EXPORT bool getCitizens(std::vector &citizens, bool ignore_sanity = false); DFHACK_EXPORT int32_t findIndexById(int32_t id); diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp index 533b40ca8..9fb19cfb3 100644 --- a/library/modules/Units.cpp +++ b/library/modules/Units.cpp @@ -59,6 +59,7 @@ using namespace std; #include "df/entity_position_assignment.h" #include "df/entity_raw.h" #include "df/entity_raw_flags.h" +#include "df/entity_site_link.h" #include "df/identity_type.h" #include "df/game_mode.h" #include "df/histfig_entity_link_positionst.h" @@ -80,6 +81,8 @@ using namespace std; #include "df/unit_soul.h" #include "df/unit_wound.h" #include "df/world.h" +#include "df/world_data.h" +#include "df/world_site.h" #include "df/unit_action.h" #include "df/unit_action_type_group.h" @@ -770,6 +773,58 @@ bool Units::getUnitsInBox (std::vector &units, return true; } +static int32_t get_noble_position_id(const df::historical_entity::T_positions &positions, const string &noble) { + string target_id = toUpper(noble); + for (auto &position : positions.own) { + if (position->code == target_id) + return position->id; + } + return -1; +} + +static void get_assigned_noble_units(vector &units, const df::historical_entity::T_positions &positions, int32_t noble_position_id, size_t limit) { + units.clear(); + for (auto &assignment : positions.assignments) { + if (assignment->position_id != noble_position_id) + continue; + auto histfig = df::historical_figure::find(assignment->histfig); + if (!histfig) + continue; + auto unit = df::unit::find(histfig->unit_id); + if (!unit) + continue; + units.emplace_back(unit); + if (limit > 0 && units.size() >= limit) + break; + } +} + +static void get_units_by_noble_role(vector &units, string noble, size_t limit = 0) { + auto &site = df::global::world->world_data->active_site[0]; + for (auto &link : site->entity_links) { + auto gov = df::historical_entity::find(link->entity_id); + if (!gov || gov->type != df::historical_entity_type::SiteGovernment) + continue; + int32_t noble_position_id = get_noble_position_id(gov->positions, noble); + if (noble_position_id < 0) + return; + get_assigned_noble_units(units, gov->positions, noble_position_id, limit); + } +} + +bool Units::getUnitsByNobleRole(vector &units, std::string noble) { + get_units_by_noble_role(units, noble); + return !units.empty(); +} + +df::unit *Units::getUnitByNobleRole(string noble) { + vector units; + get_units_by_noble_role(units, noble, 1); + if (units.empty()) + return NULL; + return units[0]; +} + bool Units::getCitizens(std::vector &citizens, bool ignore_sanity) { for (auto &unit : world->units.active) { if (isCitizen(unit, ignore_sanity))