Add stress cutoffs to Units module, fix dwarfmonitor/manipulator

Fixes #1292

docs/Lua API.rst

@@ -1271,6 +1271,18 @@ Units module
 
   Retrieves the profession color for the given race/caste using raws.
 
+* ``dfhack.units.getStressCategory(unit)``
+
+  Returns a number from 0-6 indicating stress. 0 is most stressed; 6 is least.
+  Note that 0 is guaranteed to remain the most stressed but 6 could change in the future.
+
+* ``dfhack.units.getStressCategoryRaw(stress_level)``
+
+  Identical to ``getStressCategory`` but takes a raw stress level instead of a unit.
+
+* ``dfhack.units.getStressCutoffs()``
+
+  Returns a table of the cutoffs used by the above stress level functions.
 
 Items module
 ------------

docs/changelog.txt

@@ -40,8 +40,15 @@ changelog.txt uses a syntax similar to RST, with a few special sequences:
 ## Fixes
 -@ Console: fixed crash when entering long commands on Linux/macOS
 - Fixed special characters in `command-prompt` and other non-console in-game outputs on Linux/macOS (in tools using ``df2console``)
+- `dwarfmonitor`, `manipulator`: fixed stress cutoffs
 - `startdwarf`: fixed on 64-bit Linux
 
+## API
+- Added to ``Units`` module:
+    - ``getStressCategory(unit)``
+    - ``getStressCategoryRaw(level)``
+    - ``stress_cutoffs`` (Lua: ``getStressCutoffs()``)
+
 ## Structures
 - Added ``start_dwarf_count`` on 64-bit Linux again and fixed scanning script
 -@ ``viewscreen_civlistst``: split ``unk_20`` into 3 pointers

library/LuaApi.cpp

@@ -1641,6 +1641,8 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isDomesticated),
     WRAPM(Units, getMainSocialActivity),
     WRAPM(Units, getMainSocialEvent),
+    WRAPM(Units, getStressCategory),
+    WRAPM(Units, getStressCategoryRaw),
     { NULL, NULL }
 };
 
@@ -1691,10 +1693,19 @@ static int units_getUnitsInBox(lua_State *state)
     return 2;
 }
 
+static int units_getStressCutoffs(lua_State *L)
+{
+    lua_newtable(L);
+    for (size_t i = 0; i < Units::stress_cutoffs.size(); i++)
+        Lua::TableInsert(L, i, Units::stress_cutoffs[i]);
+    return 1;
+}
+
 static const luaL_Reg dfhack_units_funcs[] = {
     { "getPosition", units_getPosition },
     { "getNoblePositions", units_getNoblePositions },
     { "getUnitsInBox", units_getUnitsInBox },
+    { "getStressCutoffs", units_getStressCutoffs },
     { NULL, NULL }
 };
 

library/include/modules/Units.h

@@ -184,5 +184,10 @@ DFHACK_EXPORT std::string getSquadName(df::unit *unit);
 DFHACK_EXPORT df::activity_entry *getMainSocialActivity(df::unit *unit);
 DFHACK_EXPORT df::activity_event *getMainSocialEvent(df::unit *unit);
 
+// stress categories - 0 is highest stress
+DFHACK_EXPORT extern const std::vector<int32_t> stress_cutoffs;
+DFHACK_EXPORT int getStressCategory(df::unit *unit);
+DFHACK_EXPORT int getStressCategoryRaw(int32_t stress_level);
+
 }
 }

library/modules/Units.cpp

@@ -1626,3 +1626,30 @@ bool Units::isDomesticated(df::unit* unit)
     }
     return tame;
 }
+
+// 50000 and up is level 0, 25000 and up is level 1, etc.
+const vector<int32_t> Units::stress_cutoffs {50000, 25000, 10000, -10000, -25000, -50000, -100000};
+
+int Units::getStressCategory(df::unit *unit)
+{
+    CHECK_NULL_POINTER(unit);
+
+    if (!unit->status.current_soul)
+        return int(stress_cutoffs.size()) / 2;
+
+    return getStressCategoryRaw(unit->status.current_soul->personality.stress_level);
+}
+
+int Units::getStressCategoryRaw(int32_t stress_level)
+{
+    int max_level = int(stress_cutoffs.size()) - 1;
+    int level = max_level;
+    for (int i = max_level; i >= 0; i--)
+    {
+        if (stress_level >= stress_cutoffs[i])
+        {
+            level = i;
+        }
+    }
+    return level;
+}

plugins/dwarfmonitor.cpp

@@ -101,23 +101,10 @@ static color_value monitor_colors[] =
 
 static int get_happiness_cat(df::unit *unit)
 {
-    if (!unit || !unit->status.current_soul)
-        return 3;
-    int stress = unit->status.current_soul->personality.stress_level;
-    if (stress >= 500000)
-        return 0;
-    else if (stress >= 250000)
-        return 1;
-    else if (stress >= 100000)
-        return 2;
-    else if (stress >= 60000)
-        return 3;
-    else if (stress >= 30000)
-        return 4;
-    else if (stress >= 0)
-        return 5;
-    else
-        return 6;
+    int level = Units::getStressCategory(unit);
+    if (level < 0) level = 0;
+    if (level > 6) level = 6;
+    return level;
 }
 
 static int get_max_history()

plugins/manipulator.cpp

@@ -1932,20 +1932,20 @@ void viewscreen_unitlaborsst::render()
         int8_t fg = 15, bg = 0;
 
         int stress_lvl = unit->status.current_soul ? unit->status.current_soul->personality.stress_level : 0;
+        const vector<UIColor> stress_colors {
+            13, // 5:1
+            12, // 4:1
+            14, // 6:1
+            2,  // 2:0
+            10  // 2:1 (default)
+        };
+        fg = vector_get(stress_colors, Units::getStressCategoryRaw(stress_lvl), stress_colors.back());
+
         // cap at 6 digits
         if (stress_lvl < -99999) stress_lvl = -99999;
         if (stress_lvl > 999999) stress_lvl = 999999;
         string stress = stl_sprintf("%6i", stress_lvl);
-        if (stress_lvl >= 500000)
-            fg = 13;    // 5:1
-        else if (stress_lvl >= 250000)
-            fg = 12;    // 4:1
-        else if (stress_lvl >= 100000)
-            fg = 14;    // 6:1
-        else if (stress_lvl >= 0)
-            fg = 2;     // 2:0
-        else
-            fg = 10;    // 2:1
+
         Screen::paintString(Screen::Pen(' ', fg, bg), col_offsets[DISP_COLUMN_STRESS], 4 + row, stress);
 
         Screen::paintTile(