port API for unit lookup by noble role from quickfort

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
 

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
+  Returns a table (list) of all citizens, which you would otherwise have to
-  units in world and test against ``isCitizen()`` to discover.
+  loop over all units in world and test against ``isCitizen()`` to discover.
 
 * ``dfhack.units.teleport(unit, pos)``
 

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<df::unit *> 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 }
 };

library/include/modules/Units.h

@@ -148,6 +148,8 @@ DFHACK_EXPORT df::unit *getUnit(const int32_t index);
 DFHACK_EXPORT bool getUnitsInBox(std::vector<df::unit*> &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<df::unit *> &units, std::string noble);
+DFHACK_EXPORT df::unit *getUnitByNobleRole(std::string noble);
 DFHACK_EXPORT bool getCitizens(std::vector<df::unit *> &citizens, bool ignore_sanity = false);
 
 DFHACK_EXPORT int32_t findIndexById(int32_t id);

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<df::unit*> &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<df::unit *> &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<df::unit *> &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<df::unit *> &units, std::string noble) {
+    get_units_by_noble_role(units, noble);
+    return !units.empty();
+}
+
+df::unit *Units::getUnitByNobleRole(string noble) {
+    vector<df::unit *> units;
+    get_units_by_noble_role(units, noble, 1);
+    if (units.empty())
+        return NULL;
+    return units[0];
+}
+
 bool Units::getCitizens(std::vector<df::unit *> &citizens, bool ignore_sanity) {
     for (auto &unit : world->units.active) {
         if (isCitizen(unit, ignore_sanity))