From 1763a16831e85051f277afb7c56e390b2c324bd0 Mon Sep 17 00:00:00 2001
From: Myk Taylor <myk002@yahoo.com>
Date: Sun, 20 Aug 2023 10:08:25 -0700
Subject: [PATCH] add training check functions to Units

---
 docs/dev/Lua API.rst            |  4 ++++
 library/LuaApi.cpp              |  4 ++++
 library/include/modules/Units.h |  4 ++++
 library/modules/Units.cpp       | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+)

diff --git a/docs/dev/Lua API.rst b/docs/dev/Lua API.rst
index f29df78c6..d14e56640 100644
--- a/docs/dev/Lua API.rst	
+++ b/docs/dev/Lua API.rst	
@@ -1380,6 +1380,10 @@ Units module
 * ``dfhack.units.isTame(unit)``
 * ``dfhack.units.isTamable(unit)``
 * ``dfhack.units.isDomesticated(unit)``
+* ``dfhack.units.isMarkedForTraining(unit)``
+* ``dfhack.units.isMarkedForTaming(unit)``
+* ``dfhack.units.isMarkedForWarTraining(unit)``
+* ``dfhack.units.isMarkedForHuntTraining(unit)``
 * ``dfhack.units.isMarkedForSlaughter(unit)``
 * ``dfhack.units.isMarkedForGelding(unit)``
 * ``dfhack.units.isGeldable(unit)``
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index beaf50579..f470e69c0 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -1777,6 +1777,10 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isTame),
     WRAPM(Units, isTamable),
     WRAPM(Units, isDomesticated),
+    WRAPM(Units, isMarkedForTraining),
+    WRAPM(Units, isMarkedForTaming),
+    WRAPM(Units, isMarkedForWarTraining),
+    WRAPM(Units, isMarkedForHuntTraining),
     WRAPM(Units, isMarkedForSlaughter),
     WRAPM(Units, isMarkedForGelding),
     WRAPM(Units, isGeldable),
diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
index 19af4ec9a..11a6c120a 100644
--- a/library/include/modules/Units.h
+++ b/library/include/modules/Units.h
@@ -109,6 +109,10 @@ DFHACK_EXPORT bool isWar(df::unit* unit);
 DFHACK_EXPORT bool isTame(df::unit* unit);
 DFHACK_EXPORT bool isTamable(df::unit* unit);
 DFHACK_EXPORT bool isDomesticated(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForTraining(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForTaming(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForWarTraining(df::unit* unit);
+DFHACK_EXPORT bool isMarkedForHuntTraining(df::unit* unit);
 DFHACK_EXPORT bool isMarkedForSlaughter(df::unit* unit);
 DFHACK_EXPORT bool isMarkedForGelding(df::unit* unit);
 DFHACK_EXPORT bool isGeldable(df::unit* unit);
diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
index 1605a3d06..ebfd13b5e 100644
--- a/library/modules/Units.cpp
+++ b/library/modules/Units.cpp
@@ -76,6 +76,7 @@ using namespace std;
 #include "df/tile_occupancy.h"
 #include "df/plotinfost.h"
 #include "df/syndrome.h"
+#include "df/training_assignment.h"
 #include "df/unit_inventory_item.h"
 #include "df/unit_misc_trait.h"
 #include "df/unit_relationship_type.h"
@@ -510,6 +511,38 @@ bool Units::isDomesticated(df::unit* unit)
     return tame;
 }
 
+static df::training_assignment * get_training_assignment(df::unit* unit) {
+    return binsearch_in_vector(df::global::plotinfo->equipment.training_assignments,
+        &df::training_assignment::animal_id, unit->id);
+}
+
+bool Units::isMarkedForTraining(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    return !!get_training_assignment(unit);
+}
+
+bool Units::isMarkedForTaming(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    auto assignment = get_training_assignment(unit);
+    return assignment && !assignment->flags.bits.train_war && !assignment->flags.bits.train_hunt;
+}
+
+bool Units::isMarkedForWarTraining(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    auto assignment = get_training_assignment(unit);
+    return assignment && assignment->flags.bits.train_war;
+}
+
+bool Units::isMarkedForHuntTraining(df::unit* unit)
+{
+    CHECK_NULL_POINTER(unit);
+    auto assignment = get_training_assignment(unit);
+    return assignment && assignment->flags.bits.train_hunt;
+}
+
 bool Units::isMarkedForSlaughter(df::unit* unit)
 {
     CHECK_NULL_POINTER(unit);