From 34de030ba9c9c78a10fdeb19e479aff4ba4e6972 Mon Sep 17 00:00:00 2001
From: Josh Cooper <cooper.s.josh@outlook.com>
Date: Wed, 9 Nov 2022 22:03:39 -0800
Subject: [PATCH] Units module

- Updates `Lua API.rst`
  - Only adds the most important additions (complicated enough to need explaining)
- Adds new functions to LuaApi.cpp
- Revises isUndead to accommodate `dfhack.units.isUndead(u)` => `Units::isUndead(u, false)` instead of taking the default value
---
 docs/Lua API.rst                | 29 +++++++++++++++++++++++++++--
 library/LuaApi.cpp              | 11 +++++++++++
 library/include/modules/Units.h |  2 +-
 library/modules/Units.cpp       |  6 +++---
 4 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/docs/Lua API.rst b/docs/Lua API.rst
index 5b6a30d93..6b8290d12 100644
--- a/docs/Lua API.rst	
+++ b/docs/Lua API.rst	
@@ -1287,10 +1287,35 @@ Units module
 
   The unit is of the correct race of the fortress.
 
-* ``dfhack.units.isCitizen(unit)``
+* ``dfhack.units.isCitizen(unit[,ignore_sanity])``
 
   The unit is an alive sane citizen of the fortress; wraps the
-  same checks the game uses to decide game-over by extinction.
+  same checks the game uses to decide game-over by extinction
+  (except for the sanity check).
+
+* ``dfhack.units.isInvader(unit)``
+
+  The unit is an active invader or marauder.
+
+* ``dfhack.units.isVisiting(unit)``
+
+  The unit is either a merchant, diplomat, or plain visitor.
+
+* ``dfhack.units.isVisitor(unit)``
+
+  The unit is strictly a visitor. Merchants and diplomats do not count here.
+
+* ``dfhack.units.isUndead(unit[,include_vamps])``
+
+  The unit is undead, but not a vampire.
+
+* ``dfhack.units.isGreatDanger(unit)``
+
+  The unit is of Great Danger. This include demons, titans, and megabeasts.
+
+* ``dfhack.units.isDanger(unit)``
+
+  The unit is dangerous, and probably hostile. This includes Great Dangers, semi-megabeasts, night creatures, undead, and invaders.
 
 * ``dfhack.units.isFortControlled(unit)``
 
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index 9f713244d..f99d65c65 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -1615,9 +1615,11 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, getRaceBabyNameById),
     WRAPM(Units, getRaceChildName),
     WRAPM(Units, getRaceChildNameById),
+    WRAPM(Units, isInvader),
     WRAPM(Units, isBaby),
     WRAPM(Units, isChild),
     WRAPM(Units, isAdult),
+    WRAPM(Units, isAnimal),
     WRAPM(Units, isEggLayer),
     WRAPM(Units, isGrazer),
     WRAPM(Units, isMilkable),
@@ -1626,8 +1628,10 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isTamable),
     WRAPM(Units, isMale),
     WRAPM(Units, isFemale),
+    WRAPM(Units, isVisiting),
     WRAPM(Units, isMerchant),
     WRAPM(Units, isDiplomat),
+    WRAPM(Units, isVisitor),
     WRAPM(Units, isForest),
     WRAPM(Units, isMarkedForSlaughter),
     WRAPM(Units, isTame),
@@ -1640,6 +1644,13 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     WRAPM(Units, isKilled),
     WRAPM(Units, isGelded),
     WRAPM(Units, isDomesticated),
+    WRAPM(Units, isDemon),
+    WRAPM(Units, isTitan),
+    WRAPM(Units, isMegabeast),
+    WRAPM(Units, isGreatDanger),
+    WRAPM(Units, isSemiMegabeast),
+    WRAPM(Units, isNightCreature),
+    WRAPM(Units, isDanger),
     WRAPM(Units, getMainSocialActivity),
     WRAPM(Units, getMainSocialEvent),
     WRAPM(Units, getStressCategory),
diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
index d25ee8d0f..84fcf5bfc 100644
--- a/library/include/modules/Units.h
+++ b/library/include/modules/Units.h
@@ -161,7 +161,7 @@ DFHACK_EXPORT bool isTame(df::unit* unit);
 DFHACK_EXPORT bool isTrained(df::unit* unit);
 DFHACK_EXPORT bool isGay(df::unit* unit);
 DFHACK_EXPORT bool isNaked(df::unit* unit);
-DFHACK_EXPORT bool isUndead(df::unit* unit, bool ignore_vamps = true);
+DFHACK_EXPORT bool isUndead(df::unit* unit, bool include_vamps = false);
 DFHACK_EXPORT bool isGhost(df::unit *unit);
 DFHACK_EXPORT bool isActive(df::unit *unit);
 DFHACK_EXPORT bool isKilled(df::unit *unit);
diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
index 0bb6235f9..b0aad49fb 100644
--- a/library/modules/Units.cpp
+++ b/library/modules/Units.cpp
@@ -1787,13 +1787,13 @@ bool Units::isNaked(df::unit* unit)
     return (unit->inventory.empty());
 }
 
-bool Units::isUndead(df::unit* unit, bool ignore_vamps)
+bool Units::isUndead(df::unit* unit, bool include_vamps)
 {
     CHECK_NULL_POINTER(unit);
 
     const auto &cb = unit->curse.add_tags1.bits;
     return unit->flags3.bits.ghostly ||
-           ((cb.OPPOSED_TO_LIFE || cb.NOT_LIVING) && (!ignore_vamps || !cb.BLOODSUCKER));
+           ((cb.OPPOSED_TO_LIFE || cb.NOT_LIVING) && (include_vamps || !cb.BLOODSUCKER));
 }
 
 bool Units::isGhost(df::unit *unit)
@@ -1895,7 +1895,7 @@ bool Units::isNightCreature(df::unit* unit)
 bool Units::isDanger(df::unit* unit) {
     CHECK_NULL_POINTER(unit);
     return isInvader(unit) ||
-           isUndead(unit, false) ||
+           isUndead(unit, true) ||
            isSemiMegabeast(unit) ||
            isNightCreature(unit) ||
            isGreatDanger(unit);