Merge pull request #4223 from myk002/myk_assign_trainer

move trainer assignment logic into core

docs/changelog.txt

@@ -65,8 +65,13 @@ Template for new versions:
 ## Documentation
 
 ## API
+- ``Units::assignTrainer``: assign a trainer to a trainable animal
+- ``Units::unassignTrainer``: unassign a trainer from an animal
 
 ## Lua
+- ``dfhack.units.isTamable``: return false for invaders to match vanilla logic
+- ``dfhack.units.assignTrainer``: expose new API to Lua
+- ``dfhack.units.unassignTrainer``: expose new API to Lua
 
 ## Removed
 - ``nopause``: functionality has moved to `spectate`

docs/dev/Lua API.rst

@@ -1485,6 +1485,13 @@ Units module
   Moves the specified unit and any riders to the target coordinates, setting
   tile occupancy flags appropriately. Returns true if successful.
 
+* ``dfhack.units.assignTrainer(unit[, trainer_id])``
+* ``dfhack.units.unassignTrainer(unit)``
+
+  Assignes (or unassigns) a trainer for the specified trainable unit. The
+  trainer ID can be omitted if "any trainer" is desired. Returns a boolean
+  indicating whether the operation was successful.
+
 * ``dfhack.units.getGeneralRef(unit, type)``
 
   Searches for a general_ref with the given type.

library/LuaApi.cpp

@@ -1823,6 +1823,7 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, setActionTimers),
     WRAPM(Units, setGroupActionTimers),
     WRAPM(Units, getUnitByNobleRole),
+    WRAPM(Units, unassignTrainer),
     { NULL, NULL }
 };
 
@@ -1927,6 +1928,16 @@ static int units_getStressCutoffs(lua_State *L)
     return 1;
 }
 
+static int units_assignTrainer(lua_State *L) {
+    df::unit * unit = Lua::CheckDFObject<df::unit>(L, 1);
+    int isNum = 0;
+    int trainer_id = lua_tointegerx(L, 2, &isNum);
+    if (!isNum || trainer_id < 0)
+        trainer_id = -1;
+    Lua::Push(L, Units::assignTrainer(unit, trainer_id));
+    return 1;
+}
+
 static const luaL_Reg dfhack_units_funcs[] = {
     { "getPosition", units_getPosition },
     { "getOuterContainerRef", units_getOuterContainerRef },
@@ -1935,6 +1946,7 @@ static const luaL_Reg dfhack_units_funcs[] = {
     { "getCitizens", units_getCitizens },
     { "getUnitsByNobleRole", units_getUnitsByNobleRole},
     { "getStressCutoffs", units_getStressCutoffs },
+    { "assignTrainer", units_assignTrainer },
     { NULL, NULL }
 };
 

library/include/modules/Units.h

@@ -176,6 +176,9 @@ DFHACK_EXPORT inline df::specific_ref getOuterContainerRef(df::unit *unit) { df:
 DFHACK_EXPORT void setNickname(df::unit *unit, std::string nick);
 DFHACK_EXPORT df::language_name *getVisibleName(df::unit *unit);
 
+DFHACK_EXPORT bool assignTrainer(df::unit *unit, int32_t trainer_id = -1);
+DFHACK_EXPORT bool unassignTrainer(df::unit *unit);
+
 DFHACK_EXPORT df::identity *getIdentity(df::unit *unit);
 DFHACK_EXPORT df::nemesis_record *getNemesis(df::unit *unit);
 

library/modules/Units.cpp

@@ -484,12 +484,35 @@ bool Units::isTame(df::unit* unit)
 bool Units::isTamable(df::unit* unit)
 {
     CHECK_NULL_POINTER(unit);
+    if (isInvader(unit))
+        return false;
     df::creature_raw *raw = world->raws.creatures.all[unit->race];
     df::caste_raw *caste = raw->caste.at(unit->caste);
     return caste->flags.is_set(caste_raw_flags::PET)
            || caste->flags.is_set(caste_raw_flags::PET_EXOTIC);
 }
 
+bool Units::assignTrainer(df::unit* unit, int32_t trainer_id) {
+    CHECK_NULL_POINTER(unit);
+    if (!isTamable(unit) || isDomesticated(unit) || isMarkedForTraining(unit))
+        return false;
+    if (trainer_id != -1 && !df::unit::find(trainer_id))
+        return false;
+    df::training_assignment *assignment = new df::training_assignment();
+    assignment->animal_id = unit->id;
+    assignment->trainer_id = trainer_id;
+    assignment->flags.bits.any_trainer = trainer_id == -1;
+    insert_into_vector(plotinfo->equipment.training_assignments,
+        &df::training_assignment::animal_id, assignment);
+    return true;
+}
+
+bool Units::unassignTrainer(df::unit* unit) {
+    CHECK_NULL_POINTER(unit);
+    return erase_from_vector(plotinfo->equipment.training_assignments,
+        &df::training_assignment::animal_id, unit->id);
+}
+
 // check if creature is domesticated
 // seems to be the only way to really tell if it's completely safe to autonestbox it (training can revert)
 bool Units::isDomesticated(df::unit* unit)
@@ -512,7 +535,7 @@ bool Units::isDomesticated(df::unit* unit)
 }
 
 static df::training_assignment * get_training_assignment(df::unit* unit) {
-    return binsearch_in_vector(df::global::plotinfo->equipment.training_assignments,
+    return binsearch_in_vector(plotinfo->equipment.training_assignments,
         &df::training_assignment::animal_id, unit->id);
 }
 

plugins/logistics.cpp

@@ -14,7 +14,6 @@
 #include "df/caravan_state.h"
 #include "df/general_ref_building_holderst.h"
 #include "df/plotinfost.h"
-#include "df/training_assignment.h"
 #include "df/world.h"
 
 using std::string;
@@ -359,22 +358,17 @@ class TrainStockProcessor : public StockProcessor {
 
     bool can_designate(color_ostream& out, df::item* item) override {
         auto unit = get_caged_unit(item);
-        return unit && !Units::isInvader(unit) &&
-            Units::isTamable(unit) && !Units::isTame(unit) &&
+        return unit &&
+            Units::isTamable(unit) &&
+            !Units::isDomesticated(unit) &&
             !Units::isMarkedForTraining(unit);
     }
 
     bool designate(color_ostream& out, df::item* item) override {
         auto unit = get_caged_unit(item);
         if (!unit)
             return false;
-        df::training_assignment *assignment = new df::training_assignment();
-        assignment->animal_id = unit->id;
-        assignment->trainer_id = -1;
-        assignment->flags.bits.any_trainer = true;
-        insert_into_vector(df::global::plotinfo->equipment.training_assignments,
-            &df::training_assignment::animal_id, assignment);
-        return true;
+        return Units::assignTrainer(unit);
     }
 
 private: