add resident-related functions

myk002 · myk002 · commit 4b66cbb954df · 2024-03-14T17:26:21.000-07:00
diff --git a/docs/changelog.txt b/docs/changelog.txt
@@ -86,12 +86,15 @@ Template for new versions:
 ## API
 - ``Units::isForgottenBeast``: property check for forgotten beasts
 - ``Units::isGreatDanger``: now includes forgotten beasts
+- ``Units::isResident``: property check for residents (as opposed to citizens)
+- ``Units::getCitizens``: now includes residents by default
 
 ## Lua
 - ``widgets.Label``: ``*pen`` attributes can now either be a pen or a function that dynamically returns a pen
 - `helpdb`: ``search_entries`` now returns a match if *all* filters in the ``include`` list are matched. previous behavior was to match if *any* ``include`` filter matched.
 - ``matinfo.decode``: now directly handles plant objects
 - ``dfhack.units.isForgottenBeast``: make new units method available to Lua
+- ``dfhack.units.getCitizens``: now includes residents by default
 
 ## Removed
 
diff --git a/docs/dev/Lua API.rst b/docs/dev/Lua API.rst
@@ -1542,9 +1542,11 @@ Units module
 
   Returns a list of units (possibly empty) assigned to the given noble role.
 
-* ``dfhack.units.getCitizens([ignore_sanity])``
+* ``dfhack.units.getCitizens([exclude_residents, [include_insane]])``
 
-  Returns a list of all living citizens that are currently on the map.
+  Returns a list of all living, sane, citizens and residents that are currently
+  on the  map. Pass ``exclude_residents`` and ``include_insane`` both default
+  to ``false`` but can be overridden.
 
 * ``dfhack.units.teleport(unit, pos)``
 
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
@@ -2150,10 +2150,11 @@ static int units_getUnitsInBox(lua_State *state)
 }
 
 static int units_getCitizens(lua_State *L) {
-    bool ignore_sanity = lua_toboolean(L, -1); // defaults to false
+    bool exclude_residents = lua_toboolean(L, 1); // defaults to false
+    bool include_insane = lua_toboolean(L, 2); // defaults to false
 
     std::vector<df::unit *> citizens;
-    if (Units::getCitizens(citizens, ignore_sanity)) {
+    if (Units::getCitizens(citizens, exclude_residents, include_insane)) {
         Lua::PushVector(L, citizens);
         return 1;
     }
diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
@@ -78,7 +78,8 @@ DFHACK_EXPORT bool isUnitInBox(df::unit* u,
 
 DFHACK_EXPORT bool isActive(df::unit *unit);
 DFHACK_EXPORT bool isVisible(df::unit* unit);
-DFHACK_EXPORT bool isCitizen(df::unit *unit, bool ignore_sanity = false);
+DFHACK_EXPORT bool isCitizen(df::unit *unit, bool include_insane = false);
+DFHACK_EXPORT bool isResident(df::unit *unit, bool include_insane = false);
 DFHACK_EXPORT bool isFortControlled(df::unit *unit);
 DFHACK_EXPORT bool isOwnCiv(df::unit* unit);
 DFHACK_EXPORT bool isOwnGroup(df::unit* unit);
@@ -158,7 +159,7 @@ DFHACK_EXPORT bool getUnitsInBox(std::vector<df::unit*> &units,
     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 bool getCitizens(std::vector<df::unit *> &citizens, bool exclude_residents = false, bool include_insane = false);
 
 DFHACK_EXPORT int32_t findIndexById(int32_t id);
 
diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
@@ -129,7 +129,7 @@ bool Units::isVisible(df::unit* unit)
     return Maps::isTileVisible(unit->pos);
 }
 
-bool Units::isCitizen(df::unit *unit, bool ignore_sanity)
+bool Units::isCitizen(df::unit *unit, bool include_insane)
 {
     CHECK_NULL_POINTER(unit);
 
@@ -149,12 +149,24 @@ bool Units::isCitizen(df::unit *unit, bool ignore_sanity)
         unit->flags2.bits.resident)
         return false;
 
-    if (!ignore_sanity && !isSane(unit))
+    if (!include_insane && !isSane(unit))
         return false;
 
     return isOwnGroup(unit);
 }
 
+bool Units::isResident(df::unit *unit, bool include_insane){
+    CHECK_NULL_POINTER(unit);
+
+    if (!include_insane && !isSane(unit))
+        return false;
+
+    return isOwnCiv(unit) &&
+        !isAnimal(unit) &&
+        !isVisitor(unit) &&
+        !isCitizen(unit, true);
+}
+
 bool Units::isFortControlled(df::unit *unit)
 {   // Reverse-engineered from ambushing unit code
     CHECK_NULL_POINTER(unit);
@@ -908,9 +920,14 @@ df::unit *Units::getUnitByNobleRole(string noble) {
     return units[0];
 }
 
-bool Units::getCitizens(std::vector<df::unit *> &citizens, bool ignore_sanity) {
+bool Units::getCitizens(std::vector<df::unit *> &citizens, bool exclude_residents, bool include_insane) {
     for (auto &unit : world->units.active) {
-        if (isCitizen(unit, ignore_sanity) && isAlive(unit) && isActive(unit))
+        if (!isAlive(unit) || !isActive(unit))
+            continue;
+
+        if (isCitizen(unit, include_insane))
+            citizens.emplace_back(unit);
+        else if (!exclude_residents && isResident(unit, include_insane))
             citizens.emplace_back(unit);
     }
     return true;
diff --git a/plugins/autochop.cpp b/plugins/autochop.cpp
@@ -574,7 +574,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;
-    Units::getCitizens(citizens);
+    Units::getCitizens(citizens, true);
     int32_t newly_marked = scan_trees(out, &expected_yield,
             &designatable_trees_by_size, true, citizens);
 
@@ -662,7 +662,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;
-    Units::getCitizens(citizens);
+    Units::getCitizens(citizens, true);
     scan_logs(out, &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);
@@ -767,7 +767,7 @@ static int autochop_getLogCounts(lua_State *L) {
     DEBUG(control,*out).print("entering autochop_getNumLogs\n");
     int32_t usable_logs, inaccessible_logs;
     vector<df::unit *> citizens;
-    Units::getCitizens(citizens);
+    Units::getCitizens(citizens, true);
     scan_logs(*out, &usable_logs, citizens, &inaccessible_logs);
     Lua::Push(L, usable_logs);
     Lua::Push(L, inaccessible_logs);
@@ -830,7 +830,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;
-    Units::getCitizens(citizens);
+    Units::getCitizens(citizens, true);
     scan_trees(*out, &expected_yield, NULL, false, citizens, &accessible_trees, &inaccessible_trees,
             &designated_trees, &accessible_yield, &tree_counts, &designated_tree_counts);
 
diff --git a/plugins/seedwatch.cpp b/plugins/seedwatch.cpp
@@ -309,7 +309,7 @@ static void scan_seeds(color_ostream &out, unordered_map<int32_t, int32_t> *acce
     static const BadFlags bad_flags;
 
     vector<df::unit *> citizens;
-    Units::getCitizens(citizens);
+    Units::getCitizens(citizens, false);
 
     for (auto &item : world->items.other[items_other_id::SEEDS]) {
         MaterialInfo mat(item);

Original file line number	Diff line number	Diff line change
`@@ -2150,10 +2150,11 @@ static int units_getUnitsInBox(lua_State *state)`
`2150`	`2150`	`}`
`2151`	`2151`
`2152`	`2152`	`static int units_getCitizens(lua_State *L) {`
`2153`		`- bool ignore_sanity = lua_toboolean(L, -1); // defaults to false`
	`2153`	`+ bool exclude_residents = lua_toboolean(L, 1); // defaults to false`
	`2154`	`+ bool include_insane = lua_toboolean(L, 2); // defaults to false`
`2154`	`2155`
`2155`	`2156`	`std::vector<df::unit *> citizens;`
`2156`		`- if (Units::getCitizens(citizens, ignore_sanity)) {`
	`2157`	`+ if (Units::getCitizens(citizens, exclude_residents, include_insane)) {`
`2157`	`2158`	`Lua::PushVector(L, citizens);`
`2158`	`2159`	`return 1;`
`2159`	`2160`	`}`