helper functions for creating and assigning jobs

docs/changelog.txt

@@ -63,9 +63,13 @@ Template for new versions:
 ## Documentation
 
 ## API
+- ``Job``: new functions ``createLinked`` and ``assignToWorkshop``
+- ``Units``: new functions ``getFocusPenalty``, ``unbailableSocialActivity``, ``isJobAvailable``
 
 ## Lua
 
+- New functions: ``dfhack.jobs.createLinked``, ``dfhack.jobs.assignToWorkshop``, ``dfhack.units.getFocusPenalty``, ``dfhack.units.unbailableSocialActivity``, and ``dfhack.units.isJobAvailable``
+
 ## Removed
 
 # 52.03-r1.1

docs/dev/Lua API.rst

@@ -1313,6 +1313,10 @@ Job module
 
   Creates a deep copy of the given job.
 
+* ``dfhack.job.createLinked()``
+
+  Create a job and immediately link it into the global job list.
+
 * ``dfhack.job.printJobDetails(job)``
 
   Prints info about the job.
@@ -1338,6 +1342,12 @@ Job module
 
   Searches for a specific_ref with the given type.
 
+* ``dfhack.job.assignToWorkshop(job, workshop)``
+
+  Assign job to workshop (i.e. establish the bidirectional link between the job
+  and the workshop). Does nothing and returns ``false`` if the workshop already
+  has the maximum of ten jobs.
+
 * ``dfhack.job.getHolder(job)``
 
   Returns the building holding the job.
@@ -1628,7 +1638,7 @@ Units module
   Returns true if the unit is within a box defined by the
   specified coordinates.
 
-``dfhack.units.getUnitsInBox(pos1, pos2[, filter])``
+* ``dfhack.units.getUnitsInBox(pos1, pos2[, filter])``
 * ``dfhack.units.getUnitsInBox(x1,y1,z1,x2,y2,z2[,filter])``
 
   Returns a table of all units within the specified coordinates.
@@ -1894,6 +1904,22 @@ Units module
   Return the ``df.activity_entry`` or ``df.activity_event`` representing the
   unit's current social activity.
 
+* ``dfhack.units.hasUnbailableSocialActivity(unit)``
+
+  Unit has an uninterruptible social activity (e.g. a purple "Socialize!").
+
+* ``dfhack.units.isJobAvailable(unit [, interrupt_social])``
+
+  Check whether a unit can be assigned to (i.e. is looking for) a job. Will
+  return ``true`` if the unit is engaged in "green" social activities, unless
+  the boolean ``interrupt_social`` is true.
+
+* ``dfhack.units.getFocusPenalty(unit, need_type [, need_type, ...])``
+
+  Get largest (i.e. most negative) focus penalty associated to a collection of
+  ``df.need_type`` arguments. Returns a number strictly greater than 400 if the
+  unit does not have any of the requested needs.
+
 * ``dfhack.units.getStressCategory(unit)``
 
   Returns a number from 0-6 indicating stress. 0 is most stressed; 6 is least.

library/LuaApi.cpp

@@ -69,6 +69,7 @@ distribution.
 #include "df/building_civzonest.h"
 #include "df/building_stockpilest.h"
 #include "df/building_tradedepotst.h"
+#include "df/building_workshopst.h"
 #include "df/burrow.h"
 #include "df/caravan_state.h"
 #include "df/construction.h"
@@ -116,6 +117,7 @@ distribution.
 #include <string>
 #include <vector>
 #include <filesystem>
+#include <stdexcept>
 
 namespace DFHack {
     DBG_DECLARE(core, luaapi, DebugCategory::LINFO);
@@ -1902,6 +1904,8 @@ static const LuaWrapper::FunctionReg dfhack_job_module[] = {
     WRAPM(Job,disconnectJobItem),
     WRAPM(Job,disconnectJobGeneralRef),
     WRAPM(Job,removeJob),
+    WRAPM(Job,createLinked),
+    WRAPM(Job,assignToWorkshop),
     WRAPN(is_equal, jobEqual),
     WRAPN(is_item_equal, jobItemEqual),
     { NULL, NULL }
@@ -2128,6 +2132,8 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, setGroupActionTimers),
     WRAPM(Units, getUnitByNobleRole),
     WRAPM(Units, unassignTrainer),
+    WRAPM(Units, hasUnbailableSocialActivity),
+    WRAPM(Units, isJobAvailable),
     { NULL, NULL }
 };
 
@@ -2321,6 +2327,25 @@ static int units_getProfessionName(lua_State *L) {
     return 1;
 }
 
+int32_t units_getFocusPenalty(lua_State *L) {
+  auto unit = Lua::GetDFObject<df::unit>(L, 1);
+  Units::need_type_set needs;
+  auto top = lua_gettop(L);
+  if (top < 2) {
+    luaL_argerror(L, 2, "Expected at least one need type");
+  } else {
+    for (int i = 2; i <= top; ++i) {
+      try {
+        needs.set(luaL_checkint(L, i));
+      } catch ([[maybe_unused]] const std::out_of_range &e) {
+        luaL_argerror(L, i, "Expected a need type");
+      }
+    }
+    Lua::Push(L, Units::getFocusPenalty(unit, needs));
+  }
+  return 1;
+}
+
 static const luaL_Reg dfhack_units_funcs[] = {
     { "getPosition", units_getPosition },
     { "getOuterContainerRef", units_getOuterContainerRef },
@@ -2335,6 +2360,7 @@ static const luaL_Reg dfhack_units_funcs[] = {
     { "getReadableName", units_getReadablename },
     { "getVisibleName", units_getVisibleName },
     { "getProfessionName", units_getProfessionName },
+    { "getFocusPenalty", units_getFocusPenalty },
     { NULL, NULL }
 };
 

library/include/modules/Job.h

@@ -31,6 +31,7 @@ distribution.
 #include "Types.h"
 #include "DataDefs.h"
 
+#include "df/building_workshopst.h"
 #include "df/item_type.h"
 #include "df/job_item_ref.h"
 
@@ -94,7 +95,13 @@ namespace DFHack
         DFHACK_EXPORT void checkBuildingsNow();
         DFHACK_EXPORT void checkDesignationsNow();
 
+        // link the job into the global job list, passing ownership to DF
         DFHACK_EXPORT bool linkIntoWorld(df::job *job, bool new_id = true);
+        // create a job and immediately link it into the global job list
+        DFHACK_EXPORT df::job* createLinked();
+
+        // assign job to workshop, returns false if workshop already has the maximum of ten jobs
+        DFHACK_EXPORT bool assignToWorkshop(df::job *job, df::building_workshopst *workshop);
 
         // Flag this job's posting as "dead" and set its posting_index to -1
         // If remove_all is true, flag all postings pointing to this job

library/include/modules/Units.h

@@ -39,12 +39,14 @@ distribution.
 #include "df/job_skill.h"
 #include "df/mental_attribute_type.h"
 #include "df/misc_trait_type.h"
+#include "df/need_type.h"
 #include "df/physical_attribute_type.h"
 #include "df/unit_action.h"
 #include "df/unit_action_type_group.h"
 #include "df/unit_path_goal.h"
 
 #include <ranges>
+#include <bitset>
 
 namespace df {
     struct activity_entry;
@@ -339,6 +341,17 @@ DFHACK_EXPORT bool isGoalAchieved(df::unit *unit, size_t goalIndex = 0);
 DFHACK_EXPORT df::activity_entry *getMainSocialActivity(df::unit *unit);
 DFHACK_EXPORT df::activity_event *getMainSocialEvent(df::unit *unit);
 
+// get largest (i.e. most negative) focus penalty for a set of needs
+using need_type_set = std::bitset<ENUM_LAST_ITEM(need_type)+1UL>;
+DFHACK_EXPORT int32_t getFocusPenalty(df::unit* unit, need_type_set need_types);
+// get focused penalty for a single need
+DFHACK_EXPORT int32_t getFocusPenalty(df::unit* unit, df::need_type need_type);
+
+// unit has an unbailable social activity (e.g. "Socialize!")
+DFHACK_EXPORT bool hasUnbailableSocialActivity(df::unit *unit);
+// unit can be assigned a job
+DFHACK_EXPORT bool isJobAvailable(df::unit *unit, bool interrupt_social);
+
 // Stress categories. 0 is highest stress, 6 is lowest.
 DFHACK_EXPORT extern const std::vector<int32_t> stress_cutoffs;
 DFHACK_EXPORT int getStressCategory(df::unit *unit);

library/modules/Job.cpp

@@ -38,6 +38,7 @@ distribution.
 #include "modules/References.h"
 
 #include "df/building.h"
+#include "df/building_workshopst.h"
 #include "df/general_ref.h"
 #include "df/general_ref_unit_workerst.h"
 #include "df/general_ref_building_holderst.h"
@@ -517,6 +518,27 @@ bool DFHack::Job::linkIntoWorld(df::job *job, bool new_id)
     }
 }
 
+df::job* DFHack::Job::createLinked()
+{
+    auto job = new df::job();
+    DFHack::Job::linkIntoWorld(job, true);
+    return job;
+}
+
+bool DFHack::Job::assignToWorkshop(df::job *job, df::building_workshopst *workshop)
+{
+    CHECK_NULL_POINTER(job);
+    CHECK_NULL_POINTER(workshop);
+
+    if (workshop->jobs.size() >= 10) {
+        return false;
+    }
+    job->pos = df::coord(workshop->centerx, workshop->centery, workshop->z);
+    DFHack::Job::addGeneralRef(job, df::general_ref_type::BUILDING_HOLDER, workshop->id);
+    workshop->jobs.push_back(job);
+    return true;
+}
+
 bool DFHack::Job::removePostings(df::job *job, bool remove_all)
 {
     using df::global::world;

library/modules/Units.cpp

@@ -64,12 +64,15 @@ distribution.
 #include "df/interaction_profilest.h"
 #include "df/item.h"
 #include "df/job.h"
+#include "df/need_type.h"
 #include "df/nemesis_record.h"
 #include "df/personality_goalst.h"
+#include "df/personality_needst.h"
 #include "df/plotinfost.h"
 #include "df/proj_unitst.h"
 #include "df/reputation_profilest.h"
 #include "df/syndrome.h"
+#include "df/squad.h"
 #include "df/tile_occupancy.h"
 #include "df/training_assignment.h"
 #include "df/unit.h"
@@ -89,6 +92,7 @@ distribution.
 #include "df/world_site.h"
 
 #include <algorithm>
+#include <bitset>
 #include <cstring>
 #include <functional>
 #include <map>
@@ -2019,6 +2023,92 @@ df::activity_event *Units::getMainSocialEvent(df::unit *unit) {
     return entry->events[entry->events.size() - 1];
 }
 
+int32_t Units::getFocusPenalty(df::unit* unit, need_type_set need_types) {
+    CHECK_NULL_POINTER(unit);
+
+    int max_penalty = INT_MAX;
+    auto& needs = unit->status.current_soul->personality.needs;
+    for (auto const need : needs) {
+        if (need_types.test(need->id)) {
+            max_penalty = min(max_penalty, need->focus_level);
+        }
+    }
+    return max_penalty;
+}
+
+int32_t Units::getFocusPenalty(df::unit* unit, df::need_type need_type) {
+    auto need_types = need_type_set().set(need_type);
+    return getFocusPenalty(unit, need_types);
+}
+
+// reverse engineered from unitst::have_unbailable_sp_activities (partial implementation)
+bool Units::hasUnbailableSocialActivity(df::unit *unit)
+{
+    // these can become constexpr with C++23
+    static const need_type_set pray_needs = need_type_set()
+        .set(df::need_type::PrayOrMeditate);
+
+    static const need_type_set socialize_needs = need_type_set()
+        .set(df::need_type::Socialize)
+        .set(df::need_type::BeCreative)
+        .set(df::need_type::Excitement)
+        .set(df::need_type::AdmireArt);
+
+    static const need_type_set read_needs = need_type_set()
+        .set(df::need_type::ThinkAbstractly)
+        .set(df::need_type::LearnSomething);
+
+    CHECK_NULL_POINTER(unit);
+
+    if (unit->social_activities.empty()) {
+        return false;
+    } else if (unit->social_activities.size() > 1) {
+        return true; // is this even possible?
+    }
+
+    auto activity = df::activity_entry::find(unit->social_activities[0]);
+    if (activity) {
+        using df::activity_entry_type;
+        switch (activity->type) {
+            case activity_entry_type::Socialize:
+                return getFocusPenalty(unit, socialize_needs) <= -10000;
+            case activity_entry_type::Prayer:
+                return getFocusPenalty(unit, pray_needs) <= -10000;
+            case activity_entry_type::Read:
+                return getFocusPenalty(unit, read_needs) <= -10000;
+            default:
+                // consider unhandled activities as uninterruptible
+                return true;
+        }
+    }
+    // this should never happen
+    return false;
+}
+
+bool Units::isJobAvailable(df::unit *unit, bool preserve_social = false){
+    if (unit->job.current_job)
+        return false;
+    if (unit->flags1.bits.caged || unit->flags1.bits.chained)
+        return false;
+    if (unit->individual_drills.size() > 0) {
+        if (unit->individual_drills.size() > 1)
+            return false; // this is even possible
+        auto activity = df::activity_entry::find(unit->individual_drills[0]);
+        if (activity && (activity->type == df::activity_entry_type::FillServiceOrder))
+            return false;
+    }
+    if (hasUnbailableSocialActivity(unit))
+        return false;
+    if (preserve_social && unit->social_activities.size() > 0)
+        return false;
+    if (unit->military.squad_id != -1) {
+        auto squad = df::squad::find(unit->military.squad_id);
+        if (squad)
+            return squad->orders.size() == 0 && squad->activity == -1;
+    }
+    return true;
+}
+
 // 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};