From 5d5502ae341030c58478dae2d04ba31ab2bfb93e Mon Sep 17 00:00:00 2001
From: Alexander Gavrilov <angavrilov@gmail.com>
Date: Wed, 11 Apr 2012 20:10:31 +0400
Subject: [PATCH] Update the item owner modification api and export it to lua.

---
 LUA_API.rst                     | 13 ++++++
 Lua API.html                    | 17 ++++++-
 library/LuaApi.cpp              |  8 ++++
 library/include/DataDefs.h      |  3 ++
 library/include/modules/Items.h | 11 ++---
 library/include/modules/Units.h |  6 ---
 library/modules/Items.cpp       | 78 +++++++++++++++++++--------------
 library/modules/Units.cpp       | 31 -------------
 plugins/cleanowned.cpp          |  7 ++-
 9 files changed, 93 insertions(+), 81 deletions(-)

diff --git a/LUA_API.rst b/LUA_API.rst
index 8698e1b23..fc86070c3 100644
--- a/LUA_API.rst
+++ b/LUA_API.rst
@@ -696,6 +696,19 @@ Units module
   Adds or removes the unit from the burrow.
 
 
+Items module
+------------
+
+* ``dfhack.items.getOwner(item)``
+
+  Returns the owner unit or *nil*.
+
+* ``dfhack.items.setOwner(item,unit)``
+
+  Replaces the owner of the item. If unit is *nil*, removes ownership.
+  Returns *false* in case of error.
+
+
 Maps module
 -----------
 
diff --git a/Lua API.html b/Lua API.html
index 0ca88966c..7c3eeddb9 100644
--- a/Lua API.html	
+++ b/Lua API.html	
@@ -340,7 +340,8 @@ ul.auto-toc {
 <li><a class="reference internal" href="#gui-module" id="id14">Gui module</a></li>
 <li><a class="reference internal" href="#job-module" id="id15">Job module</a></li>
 <li><a class="reference internal" href="#units-module" id="id16">Units module</a></li>
-<li><a class="reference internal" href="#maps-module" id="id17">Maps module</a></li>
+<li><a class="reference internal" href="#items-module" id="id17">Items module</a></li>
+<li><a class="reference internal" href="#maps-module" id="id18">Maps module</a></li>
 </ul>
 </li>
 </ul>
@@ -925,8 +926,20 @@ a lua list containing them.</p>
 </li>
 </ul>
 </div>
+<div class="section" id="items-module">
+<h3><a class="toc-backref" href="#id17">Items module</a></h3>
+<ul>
+<li><p class="first"><tt class="docutils literal">dfhack.items.getOwner(item)</tt></p>
+<p>Returns the owner unit or <em>nil</em>.</p>
+</li>
+<li><p class="first"><tt class="docutils literal">dfhack.items.setOwner(item,unit)</tt></p>
+<p>Replaces the owner of the item. If unit is <em>nil</em>, removes ownership.
+Returns <em>false</em> in case of error.</p>
+</li>
+</ul>
+</div>
 <div class="section" id="maps-module">
-<h3><a class="toc-backref" href="#id17">Maps module</a></h3>
+<h3><a class="toc-backref" href="#id18">Maps module</a></h3>
 <ul>
 <li><p class="first"><tt class="docutils literal">dfhack.maps.getSize()</tt></p>
 <p>Returns map size in blocks: <em>x, y, z</em></p>
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
index fb818c206..febc59026 100644
--- a/library/LuaApi.cpp
+++ b/library/LuaApi.cpp
@@ -42,6 +42,7 @@ distribution.
 #include "modules/Job.h"
 #include "modules/Translation.h"
 #include "modules/Units.h"
+#include "modules/Items.h"
 #include "modules/Materials.h"
 #include "modules/Maps.h"
 
@@ -609,6 +610,12 @@ static const LuaWrapper::FunctionReg dfhack_units_module[] = {
     { NULL, NULL }
 };
 
+static const LuaWrapper::FunctionReg dfhack_items_module[] = {
+    WRAPM(Items, getOwner),
+    WRAPM(Items, setOwner),
+    { NULL, NULL }
+};
+
 static bool maps_isBlockBurrowTile(df::burrow *burrow, df::map_block *block, int x, int y)
 {
     return Maps::isBlockBurrowTile(burrow, block, df::coord2d(x,y));
@@ -666,5 +673,6 @@ void OpenDFHackApi(lua_State *state)
     OpenModule(state, "gui", dfhack_gui_module);
     OpenModule(state, "job", dfhack_job_module, dfhack_job_funcs);
     OpenModule(state, "units", dfhack_units_module);
+    OpenModule(state, "items", dfhack_items_module);
     OpenModule(state, "maps", dfhack_maps_module, dfhack_maps_funcs);
 }
diff --git a/library/include/DataDefs.h b/library/include/DataDefs.h
index b9b4707d9..d4d757d94 100644
--- a/library/include/DataDefs.h
+++ b/library/include/DataDefs.h
@@ -425,6 +425,9 @@ namespace df
         static compound_identity *get() { return &T::_identity; }
     };
 
+    template<class T>
+    inline T* allocate() { return (T*)identity_traits<T>::get()->allocate(); }
+
     template<class T>
     struct enum_traits {};
 
diff --git a/library/include/modules/Items.h b/library/include/modules/Items.h
index e58c9214b..9bab02cc8 100644
--- a/library/include/modules/Items.h
+++ b/library/include/modules/Items.h
@@ -122,16 +122,17 @@ DFHACK_EXPORT bool copyItem(df::item * source, dfh_item & target);
 /// write copied item back to its origin
 DFHACK_EXPORT bool writeItem(const dfh_item & item);
 
-/// who owns this item we already read?
-DFHACK_EXPORT int32_t getItemOwnerID(const df::item * item);
-DFHACK_EXPORT df::unit *getItemOwner(const df::item * item);
+/// Retrieve the owner of the item.
+DFHACK_EXPORT df::unit *getOwner(df::item *item);
+/// Set the owner of the item. Pass NULL as unit to remove the owner.
+DFHACK_EXPORT bool setOwner(df::item *item, df::unit *unit);
+
 /// which item is it contained in?
 DFHACK_EXPORT int32_t getItemContainerID(const df::item * item);
 DFHACK_EXPORT df::item *getItemContainer(const df::item * item);
 /// which items does it contain?
 DFHACK_EXPORT bool getContainedItems(const df::item * item, /*output*/ std::vector<int32_t> &items);
-/// wipe out the owner records
-DFHACK_EXPORT bool removeItemOwner(df::item * item);
+
 /// read item references, filtered by class
 DFHACK_EXPORT bool readItemRefs(const df::item * item, const df::general_ref_type type,
                   /*output*/ std::vector<int32_t> &values);
diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
index 35701b6ec..12e687112 100644
--- a/library/include/modules/Units.h
+++ b/library/include/modules/Units.h
@@ -185,9 +185,6 @@ DFHACK_EXPORT bool ReadJob(const df::unit * unit, std::vector<t_material> & mat)
 DFHACK_EXPORT bool ReadInventoryByIdx(const uint32_t index, std::vector<df::item *> & item);
 DFHACK_EXPORT bool ReadInventoryByPtr(const df::unit * unit, std::vector<df::item *> & item);
 
-DFHACK_EXPORT bool ReadOwnedItemsByIdx(const uint32_t index, std::vector<int32_t> & item);
-DFHACK_EXPORT bool ReadOwnedItemsByPtr(const df::unit * unit, std::vector<int32_t> & item);
-
 DFHACK_EXPORT int32_t FindIndexById(int32_t id);
 
 /* Getters */
@@ -196,9 +193,6 @@ DFHACK_EXPORT int32_t GetDwarfCivId ( void );
 
 DFHACK_EXPORT void CopyNameTo(df::unit *creature, df::language_name * target);
 
-DFHACK_EXPORT bool RemoveOwnedItemByIdx(const uint32_t index, int32_t id);
-DFHACK_EXPORT bool RemoveOwnedItemByPtr(df::unit * unit, int32_t id);
-
 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 19618fc10..d884a734e 100644
--- a/library/modules/Items.cpp
+++ b/library/modules/Items.cpp
@@ -41,6 +41,7 @@ using namespace std;
 #include "modules/Units.h"
 #include "ModuleFactory.h"
 #include "Core.h"
+#include "Error.h"
 #include "MiscUtils.h"
 
 #include "df/world.h"
@@ -63,6 +64,7 @@ using namespace std;
 #include "df/trapcomp_flags.h"
 #include "df/job_item.h"
 #include "df/general_ref.h"
+#include "df/general_ref_unit_itemownerst.h"
 
 using namespace DFHack;
 using namespace df::enums;
@@ -415,28 +417,62 @@ bool Items::copyItem(df::item * itembase, DFHack::dfh_item &item)
     return true;
 }
 
-int32_t Items::getItemOwnerID(const df::item * item)
+df::unit *Items::getOwner(df::item * item)
 {
+    CHECK_NULL_POINTER(item);
+
     for (size_t i = 0; i < item->itemrefs.size(); i++)
     {
         df::general_ref *ref = item->itemrefs[i];
-        if (ref->getType() == general_ref_type::UNIT_ITEMOWNER)
-            return ref->getID();
+        if (strict_virtual_cast<df::general_ref_unit_itemownerst>(ref))
+            return ref->getUnit();
     }
-    return -1;
+
+    return NULL;
 }
 
-df::unit *Items::getItemOwner(const df::item * item)
+bool Items::setOwner(df::item *item, df::unit *unit)
 {
-    for (size_t i = 0; i < item->itemrefs.size(); i++)
+    CHECK_NULL_POINTER(item);
+
+    for (int i = item->itemrefs.size()-1; i >= 0; i--)
     {
         df::general_ref *ref = item->itemrefs[i];
-        if (ref->getType() == general_ref_type::UNIT_ITEMOWNER)
-            return ref->getUnit();
+
+        if (!strict_virtual_cast<df::general_ref_unit_itemownerst>(ref))
+            continue;
+
+        if (auto cur = ref->getUnit())
+        {
+            if (cur == unit)
+                return true;
+
+            erase_from_vector(cur->owned_items, item->id);
+        }
+
+        delete ref;
+        vector_erase_at(item->itemrefs, i);
     }
-    return NULL;
+
+    if (unit)
+    {
+        auto ref = df::allocate<df::general_ref_unit_itemownerst>();
+        if (!ref)
+            return false;
+
+        item->flags.bits.owned = true;
+        ref->unit_id = unit->id;
+
+        insert_into_vector(unit->owned_items, item->id);
+        item->itemrefs.push_back(ref);
+    }
+    else
+        item->flags.bits.owned = false;
+
+    return true;
 }
 
+
 int32_t Items::getItemContainerID(const df::item * item)
 {
     for (size_t i = 0; i < item->itemrefs.size(); i++)
@@ -477,27 +513,3 @@ bool Items::readItemRefs(const df::item * item, df::general_ref_type type, std::
 
     return !values.empty();
 }
-
-bool Items::removeItemOwner(df::item * item)
-{
-    for (size_t i = 0; i < item->itemrefs.size(); i++)
-    {
-        df::general_ref *ref = item->itemrefs[i];
-        if (ref->getType() != general_ref_type::UNIT_ITEMOWNER)
-            continue;
-
-        df::unit *unit = ref->getUnit();
-
-        if (unit == NULL || !Units::RemoveOwnedItemByPtr(unit, item->id))
-        {
-            cerr << "RemoveOwnedItemIdx: CREATURE " << ref->getID() << " ID " << item->id << " FAILED!" << endl;
-            return false;
-        }
-        delete ref;
-        item->itemrefs.erase(item->itemrefs.begin() + i--);
-    }
-
-    item->flags.bits.owned = 0;
-
-    return true;
-}
diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
index 47208dcbc..3e58e0d5c 100644
--- a/library/modules/Units.cpp
+++ b/library/modules/Units.cpp
@@ -495,37 +495,6 @@ bool Units::ReadInventoryByPtr(const df::unit * unit, std::vector<df::item *> &
     return true;
 }
 
-bool Units::ReadOwnedItemsByIdx(const uint32_t index, std::vector<int32_t> & item)
-{
-    if(index >= world->units.all.size()) return false;
-    df::unit * temp = world->units.all[index];
-    return ReadOwnedItemsByPtr(temp, item);
-}
-
-bool Units::ReadOwnedItemsByPtr(const df::unit * unit, std::vector<int32_t> & items)
-{
-    if(!isValid()) return false;
-    if(!unit) return false;
-    items = unit->owned_items;
-    return true;
-}
-
-bool Units::RemoveOwnedItemByIdx(const uint32_t index, int32_t id)
-{
-    if(index >= world->units.all.size()) return false;
-    df::unit * temp = world->units.all[index];
-    return RemoveOwnedItemByPtr(temp, id);
-}
-
-bool Units::RemoveOwnedItemByPtr(df::unit * unit, int32_t id)
-{
-    if(!isValid()) return false;
-    if(!unit) return false;
-    vector <int32_t> & vec = unit->owned_items;
-    vec.erase(std::remove(vec.begin(), vec.end(), id), vec.end());
-    return true;
-}
-
 void Units::CopyNameTo(df::unit * creature, df::language_name * target)
 {
     Translation::copyName(&creature->name, target);
diff --git a/plugins/cleanowned.cpp b/plugins/cleanowned.cpp
index 280f2e1c5..c1521b8de 100644
--- a/plugins/cleanowned.cpp
+++ b/plugins/cleanowned.cpp
@@ -98,8 +98,7 @@ command_result df_cleanowned (color_ostream &out, vector <string> & parameters)
 
         if (!item->flags.bits.owned)
         {
-            int32_t owner = Items::getItemOwnerID(item);
-            if (owner >= 0)
+            if (Items::getOwner(item))
             {
                 out.print("Fixing a misflagged item: \t");
                 confiscate = true;
@@ -168,14 +167,14 @@ command_result df_cleanowned (color_ostream &out, vector <string> & parameters)
                 item->getWear()
             );
 
-            df::unit *owner = Items::getItemOwner(item);
+            df::unit *owner = Items::getOwner(item);
 
             if (owner)
                 out.print(", owner %s", Translation::TranslateName(&owner->name,false).c_str());
 
             if (!dry_run)
             {
-                if (!Items::removeItemOwner(item))
+                if (!Items::setOwner(item,NULL))
                     out.print("(unsuccessfully) ");
                 if (dump)
                     item->flags.bits.dump = 1;