Merge pull request #2752 from myk002/myk_get_citizens

Add API for getting a list of citizens

docs/changelog.txt

@@ -55,11 +55,13 @@ changelog.txt uses a syntax similar to RST, with a few special sequences:
 
 ## API
 - ``Buildings::containsTile()``: no longer takes a ``room`` parameter since that's not how rooms work anymore. If the building has extents, the extents will be checked. otherwise, the result just depends on whether the tile is within the building's bounding box.
+- ``Units::getCitizens()``: gets a list of citizens, which otherwise you'd have to iterate over all units the world to discover
 
 ## Lua
 - `helpdb`: new function: ``helpdb.refresh()`` to force a refresh of the database. Call if you are a developer adding new scripts, loading new plugins, or changing help text during play
 - `helpdb`: changed from auto-refreshing every 60 seconds to only refreshing on explicit call to ``helpdb.refresh()``. docs very rarely change during a play session, and the automatic database refreshes were slowing down the startup of `gui/launcher`
 - ``widgets.Label``: ``label.scroll()`` now understands ``home`` and ``end`` keywords for scrolling to the top or bottom
+- ``dfhack.units.getCitizens()``: gets a list of citizens
 
 ## Removed
 

docs/dev/Lua API.rst

@@ -1412,6 +1412,11 @@ Units module
   Note that ``pos2xyz()`` cannot currently be used to convert coordinate objects to
   the arguments required by this function.
 
+* ``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.
+
 * ``dfhack.units.teleport(unit, pos)``
 
   Moves the specified unit and any riders to the target coordinates, setting

library/LuaApi.cpp

@@ -1883,6 +1883,17 @@ static int units_getUnitsInBox(lua_State *state)
     return 2;
 }
 
+static int units_getCitizens(lua_State *L) {
+    bool ignore_sanity = lua_toboolean(L, -1); // defaults to false
+
+    std::vector<df::unit *> citizens;
+    if (Units::getCitizens(citizens, ignore_sanity)) {
+        Lua::PushVector(L, citizens);
+        return 1;
+    }
+    return 0;
+}
+
 static int units_getStressCutoffs(lua_State *L)
 {
     lua_newtable(L);
@@ -1896,6 +1907,7 @@ static const luaL_Reg dfhack_units_funcs[] = {
     { "getOuterContainerRef", units_getOuterContainerRef },
     { "getNoblePositions", units_getNoblePositions },
     { "getUnitsInBox", units_getUnitsInBox },
+    { "getCitizens", units_getCitizens },
     { "getStressCutoffs", units_getStressCutoffs },
     { NULL, NULL }
 };

library/include/modules/Units.h

@@ -145,6 +145,7 @@ 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 getCitizens(std::vector<df::unit *> &citizens, bool ignore_sanity = false);
 
 DFHACK_EXPORT int32_t findIndexById(int32_t id);
 

library/modules/Units.cpp

@@ -756,6 +756,14 @@ bool Units::getUnitsInBox (std::vector<df::unit*> &units,
     return true;
 }
 
+bool Units::getCitizens(std::vector<df::unit *> &citizens, bool ignore_sanity) {
+    for (auto &unit : world->units.active) {
+        if (isCitizen(unit, ignore_sanity))
+            citizens.emplace_back(unit);
+    }
+    return true;
+}
+
 int32_t Units::findIndexById(int32_t creature_id)
 {
     return df::unit::binsearch_index(world->units.all, creature_id);

plugins/autochop.cpp

@@ -369,13 +369,6 @@ static void bucket_tree(df::plant *plant, bool designate_clearcut, bool *designa
         *can_chop = true;
 }
 
-static void get_citizens(vector<df::unit *> &vec) {
-    for (auto &unit : world->units.active) {
-        if (Units::isCitizen(unit))
-            vec.emplace_back(unit);
-    }
-}
-
 static void bucket_watched_burrows(color_ostream & out,
         map<int, PersistentDataItem *> &clearcut_burrows,
         map<int, PersistentDataItem *> &chop_burrows) {
@@ -546,7 +539,7 @@ static int32_t do_cycle(color_ostream &out, bool force_designate) {
     int32_t expected_yield;
     TreesBySize designatable_trees_by_size;
     vector<df::unit *> citizens;
-    get_citizens(citizens);
+    Units::getCitizens(citizens);
     int32_t newly_marked = scan_trees(out, &expected_yield,
             &designatable_trees_by_size, true, citizens);
 
@@ -634,7 +627,7 @@ static void autochop_printStatus(color_ostream &out) {
     int32_t designated_trees, expected_yield, accessible_yield;
     map<int32_t, int32_t> tree_counts, designated_tree_counts;
     vector<df::unit *> citizens;
-    get_citizens(citizens);
+    Units::getCitizens(citizens);
     scan_logs(&usable_logs, citizens, &inaccessible_logs);
     scan_trees(out, &expected_yield, NULL, false, citizens, &accessible_trees, &inaccessible_trees,
             &designated_trees, &accessible_yield, &tree_counts, &designated_tree_counts);
@@ -739,7 +732,7 @@ static int autochop_getLogCounts(lua_State *L) {
     DEBUG(status,*out).print("entering autochop_getNumLogs\n");
     int32_t usable_logs, inaccessible_logs;
     vector<df::unit *> citizens;
-    get_citizens(citizens);
+    Units::getCitizens(citizens);
     scan_logs(&usable_logs, citizens, &inaccessible_logs);
     Lua::Push(L, usable_logs);
     Lua::Push(L, inaccessible_logs);
@@ -778,7 +771,7 @@ static int autochop_getTreeCountsAndBurrowConfigs(lua_State *L) {
     int32_t designated_trees, expected_yield, accessible_yield;
     map<int32_t, int32_t> tree_counts, designated_tree_counts;
     vector<df::unit *> citizens;
-    get_citizens(citizens);
+    Units::getCitizens(citizens);
     scan_trees(*out, &expected_yield, NULL, false, citizens, &accessible_trees, &inaccessible_trees,
             &designated_trees, &accessible_yield, &tree_counts, &designated_tree_counts);