From 37cfb1fdcd01e4a7dce1a2b759cc2655bd42cf03 Mon Sep 17 00:00:00 2001
From: Alexander Gavrilov <angavrilov@gmail.com>
Date: Fri, 13 Apr 2012 16:10:19 +0400
Subject: [PATCH] Add unit position and container api.

---
 LUA_API.rst                     | 10 +++++++++-
 Lua API.html                    |  8 +++++++-
 library/LuaApi.cpp              | 13 ++++++++++++-
 library/include/modules/Units.h |  5 +++++
 library/modules/Items.cpp       |  8 +++++++-
 library/modules/Units.cpp       | 29 +++++++++++++++++++++++++++++
 6 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/LUA_API.rst b/LUA_API.rst
index df997008c..673053189 100644
--- a/LUA_API.rst
+++ b/LUA_API.rst
@@ -663,6 +663,14 @@ Job module
 Units module
 ------------
 
+* ``dfhack.units.getPosition(unit)``
+
+  Returns true *x,y,z* of the unit; may be not equal to unit.pos if caged.
+
+* ``dfhack.units.getContainer(unit)``
+
+  Returns the container (cage) item or *nil*.
+
 * ``dfhack.units.setNickname(unit,nick)``
 
   Sets the unit's nickname properly.
@@ -701,7 +709,7 @@ Items module
 
 * ``dfhack.items.getPosition(item)``
 
-  Returns true *x,y,z* of the item.
+  Returns true *x,y,z* of the item; may be not equal to item.pos if in inventory.
 
 * ``dfhack.items.getOwner(item)``
 
diff --git a/Lua API.html b/Lua API.html
index 4be1050fb..32a3b7f04 100644
--- a/Lua API.html	
+++ b/Lua API.html	
@@ -900,6 +900,12 @@ a lua list containing them.</p>
 <div class="section" id="units-module">
 <h3><a class="toc-backref" href="#id16">Units module</a></h3>
 <ul>
+<li><p class="first"><tt class="docutils literal">dfhack.units.getPosition(unit)</tt></p>
+<p>Returns true <em>x,y,z</em> of the unit; may be not equal to unit.pos if caged.</p>
+</li>
+<li><p class="first"><tt class="docutils literal">dfhack.units.getContainer(unit)</tt></p>
+<p>Returns the container (cage) item or <em>nil</em>.</p>
+</li>
 <li><p class="first"><tt class="docutils literal">dfhack.units.setNickname(unit,nick)</tt></p>
 <p>Sets the unit's nickname properly.</p>
 </li>
@@ -930,7 +936,7 @@ a lua list containing them.</p>
 <h3><a class="toc-backref" href="#id17">Items module</a></h3>
 <ul>
 <li><p class="first"><tt class="docutils literal">dfhack.items.getPosition(item)</tt></p>
-<p>Returns true <em>x,y,z</em> of the item.</p>
+<p>Returns true <em>x,y,z</em> of the item; may be not equal to item.pos if in inventory.</p>
 </li>
 <li><p class="first"><tt class="docutils literal">dfhack.items.getOwner(item)</tt></p>
 <p>Returns the owner unit or <em>nil</em>.</p>
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index 25f654b8c..6792c921e 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -627,6 +627,7 @@ static const luaL_Reg dfhack_job_funcs[] = {
 
 
 static const LuaWrapper::FunctionReg dfhack_units_module[] = {
+    WRAPM(Units, getContainer),
     WRAPM(Units, setNickname),
     WRAPM(Units, getVisibleName),
     WRAPM(Units, getNemesis),
@@ -638,6 +639,16 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     { NULL, NULL }
 };
 
+static int units_getPosition(lua_State *state)
+{
+    return push_pos(state, Units::getPosition(get_checked_arg<df::unit>(state,1)));
+}
+
+static const luaL_Reg dfhack_units_funcs[] = {
+    { "getPosition", units_getPosition },
+    { NULL, NULL }
+};
+
 static bool items_moveToGround(df::item *item, df::coord pos)
 {
     MapExtras::MapCache mc;
@@ -729,7 +740,7 @@ void OpenDFHackApi(lua_State *state)
     LuaWrapper::SetFunctionWrappers(state, dfhack_module);
     OpenModule(state, "gui", dfhack_gui_module);
     OpenModule(state, "job", dfhack_job_module, dfhack_job_funcs);
-    OpenModule(state, "units", dfhack_units_module);
+    OpenModule(state, "units", dfhack_units_module, dfhack_units_funcs);
     OpenModule(state, "items", dfhack_items_module, dfhack_items_funcs);
     OpenModule(state, "maps", dfhack_maps_module, dfhack_maps_funcs);
 }
diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
index 12e687112..d7dadb21b 100644
--- a/library/include/modules/Units.h
+++ b/library/include/modules/Units.h
@@ -193,6 +193,11 @@ DFHACK_EXPORT int32_t GetDwarfCivId ( void );
 
 DFHACK_EXPORT void CopyNameTo(df::unit *creature, df::language_name * target);
 
+/// Returns the true position of the unit (non-trivial in case of caged).
+DFHACK_EXPORT df::coord getPosition(df::unit *unit);
+
+DFHACK_EXPORT df::item *getContainer(df::unit *unit);
+
 DFHACK_EXPORT void setNickname(df::unit *unit, std::string nick);
 DFHACK_EXPORT df::language_name *getVisibleName(df::unit *unit);
 
diff --git a/library/modules/Items.cpp b/library/modules/Items.cpp
index ca15ed278..9d0657e19 100644
--- a/library/modules/Items.cpp
+++ b/library/modules/Items.cpp
@@ -529,7 +529,7 @@ df::coord Items::getPosition(df::item *item)
 
             case general_ref_type::UNIT_HOLDER:
                 if (auto unit = ref->getUnit())
-                    return unit->pos;
+                    return Units::getPosition(unit);
                 break;
 
             case general_ref_type::BUILDING_HOLDER:
@@ -579,7 +579,11 @@ static bool detachItem(MapExtras::MapCache &mc, df::item *item)
             {
             case general_ref_type::CONTAINED_IN_ITEM:
                 if (auto item2 = ref->getItem())
+                {
+                    item2->flags.bits.weight_computed = false;
+
                     removeRef(item2->itemrefs, general_ref_type::CONTAINS_ITEM, item->id);
+                }
                 break;
 
             case general_ref_type::UNIT_HOLDER:
@@ -649,6 +653,8 @@ bool DFHack::Items::moveToContainer(MapExtras::MapCache &mc, df::item *item, df:
     item->flags.bits.in_inventory = true;
     container->flags.bits.container = true;
 
+    container->flags.bits.weight_computed = false;
+
     ref1->item_id = item->id;
     container->itemrefs.push_back(ref1);
     ref2->item_id = container->id;
diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
index 3e58e0d5c..7dbb40c97 100644
--- a/library/modules/Units.cpp
+++ b/library/modules/Units.cpp
@@ -40,6 +40,7 @@ using namespace std;
 
 // we connect to those
 #include "modules/Units.h"
+#include "modules/Items.h"
 #include "modules/Materials.h"
 #include "modules/Translation.h"
 #include "ModuleFactory.h"
@@ -500,6 +501,34 @@ void Units::CopyNameTo(df::unit * creature, df::language_name * target)
     Translation::copyName(&creature->name, target);
 }
 
+df::coord Units::getPosition(df::unit *unit)
+{
+    CHECK_NULL_POINTER(unit);
+
+    if (unit->flags1.bits.caged)
+    {
+        auto cage = getContainer(unit);
+        if (cage)
+            return Items::getPosition(cage);
+    }
+
+    return unit->pos;
+}
+
+df::item *Units::getContainer(df::unit *unit)
+{
+    CHECK_NULL_POINTER(unit);
+
+    for (size_t i = 0; i < unit->refs.size(); i++)
+    {
+        df::general_ref *ref = unit->refs[i];
+        if (ref->getType() == general_ref_type::CONTAINED_IN_ITEM)
+            return ref->getItem();
+    }
+
+    return NULL;
+}
+
 void Units::setNickname(df::unit *unit, std::string nick)
 {
     CHECK_NULL_POINTER(unit);