helper functions for creating jobs and checking whether units are ready to take jobs

Author: chdoc

docs/changelog.txt

@@ -63,6 +63,8 @@ Template for new versions:
 ## Documentation
 
 ## API
+- ``Job``: new functions ``createLinked`` and ``assignToWorkshop``
+- ``Units``: new functions ``getFocusPenalty``, ``unbailableSocialActivity``, ``isJobAvailable``
 
 ## Lua
 

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_types_set = std::bitset<ENUM_LAST_ITEM(need_type)+1UL>;
+DFHACK_EXPORT int32_t getFocusPenalty(df::unit* unit, need_types_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 unbailableSocialActivity(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 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>
@@ -2017,6 +2021,92 @@ df::activity_event *Units::getMainSocialEvent(df::unit *unit) {
     return entry->events[entry->events.size() - 1];
 }
 
+int32_t Units::getFocusPenalty(df::unit* unit, need_types_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_types_set().set(need_type);
+    return getFocusPenalty(unit, need_types);
+}
+
+// reverse engineered from unitst::have_unbailable_sp_activities (partial implementation)
+bool Units::unbailableSocialActivity(df::unit *unit)
+{
+    // these can become constexpr with C++23
+    static const need_types_set pray_needs = need_types_set()
+        .set(df::need_type::PrayOrMeditate);
+
+    static const need_types_set socialize_needs = need_types_set()
+        .set(df::need_type::Socialize)
+        .set(df::need_type::BeCreative)
+        .set(df::need_type::Excitement)
+        .set(df::need_type::AdmireArt);
+
+    static const need_types_set read_needs = need_types_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 (unbailableSocialActivity(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};