Support linking siege engines to stockpiles.

Since they can't do that natively, the links object has to
be maintained in dfhack memory, and with dfhack persistence.

plugins/devel/siege-engine.cpp

@@ -46,6 +46,8 @@
 #include "df/job_item.h"
 #include "df/item.h"
 #include "df/items_other_id.h"
+#include "df/building_stockpilest.h"
+#include "df/stockpile_links.h"
 
 #include "MiscUtils.h"
 
@@ -158,6 +160,9 @@ struct EngineInfo {
 
     int operator_id, operator_frame;
 
+    std::set<int> stockpiles;
+    df::stockpile_links links;
+
     bool hasTarget() { return is_range_valid(target); }
     bool onTarget(df::coord pos) { return is_in_range(target, pos); }
     df::coord getTargetSize() { return target.second - target.first; }
@@ -170,6 +175,8 @@ struct EngineInfo {
 static std::map<df::building*, EngineInfo> engines;
 static std::map<df::coord, df::building*> coord_engines;
 
+static std::set<df::building*> recheck_piles;
+
 static EngineInfo *find_engine(df::building *bld, bool create = false)
 {
     auto ebld = strict_virtual_cast<df::building_siegeenginest>(bld);
@@ -209,12 +216,12 @@ static EngineInfo *find_engine(df::building *bld, bool create = false)
     return obj;
 }
 
-static EngineInfo *find_engine(lua_State *L, int idx, bool create = false)
+static EngineInfo *find_engine(lua_State *L, int idx, bool create = false, bool silent = false)
 {
     auto bld = Lua::CheckDFObject<df::building_siegeenginest>(L, idx);
 
     auto engine = find_engine(bld, create);
-    if (!engine)
+    if (!engine && !silent)
         luaL_error(L, "no such engine");
 
     return engine;
@@ -241,6 +248,7 @@ static void clear_engines()
 {
     engines.clear();
     coord_engines.clear();
+    recheck_piles.clear();
 }
 
 static void load_engines()
@@ -267,13 +275,33 @@ static void load_engines()
         engine->ammo_vector_id = (df::job_item_vector_id)it->ival(1);
         engine->ammo_item_type = (df::item_type)it->ival(2);
     }
+
+    pworld->GetPersistentData(&vec, "siege-engine/stockpiles/", true);
+    for (auto it = vec.begin(); it != vec.end(); ++it)
+    {
+        auto engine = find_engine(df::building::find(it->ival(0)), true);
+        if (!engine)
+            continue;
+        auto pile = df::building::find(it->ival(1));
+        if (!pile || pile->getType() != building_type::Stockpile)
+        {
+            pworld->DeletePersistentData(*it);
+            continue;;
+        }
+        auto plinks = pile->getStockpileLinks();
+        if (!plinks)
+            continue;
+
+        engine->stockpiles.insert(it->ival(1));
+
+        insert_into_vector(engine->links.take_from_pile, &df::building::id, pile);
+        insert_into_vector(plinks->give_to_workshop, &df::building::id, (df::building*)engine->bld);
+    }
 }
 
 static int getTargetArea(lua_State *L)
 {
-    auto bld = Lua::CheckDFObject<df::building>(L, 1);
-    if (!bld) luaL_argerror(L, 1, "null building");
-    auto engine = find_engine(bld);
+    auto engine = find_engine(L, 1, false, true);
 
     if (engine && engine->hasTarget())
     {
@@ -335,7 +363,10 @@ static bool setTargetArea(df::building_siegeenginest *bld, df::coord target_min,
 
 static int getAmmoItem(lua_State *L)
 {
-    auto engine = find_engine(L, 1, true);
+    auto engine = find_engine(L, 1, false, true);
+    if (!engine)
+        Lua::Push(L, item_type::BOULDER);
+    else
         Lua::Push(L, engine->ammo_item_type);
     return 1;
 }
@@ -378,6 +409,123 @@ static int setAmmoItem(lua_State *L)
     return 1;
 }
 
+static int getStockpileLinks(lua_State *L)
+{
+    auto engine = find_engine(L, 1, false, true);
+    if (!engine || engine->stockpiles.empty())
+        return 0;
+
+    int idx = 1;
+    lua_createtable(L, engine->stockpiles.size(), 0);
+
+    for (auto it = engine->stockpiles.begin(); it != engine->stockpiles.end(); ++it)
+    {
+        auto pile = df::building::find(*it);
+        if (!pile) continue;
+        Lua::Push(L, pile);
+        lua_rawseti(L, -2, idx++);
+    }
+
+    return 1;
+}
+
+static bool isLinkedToPile(df::building_siegeenginest *bld, df::building_stockpilest *pile)
+{
+    CHECK_NULL_POINTER(bld);
+    CHECK_NULL_POINTER(pile);
+
+    auto engine = find_engine(bld);
+
+    return engine && engine->stockpiles.count(pile->id);
+}
+
+static bool addStockpileLink(df::building_siegeenginest *bld, df::building_stockpilest *pile)
+{
+    CHECK_NULL_POINTER(bld);
+    CHECK_NULL_POINTER(pile);
+
+    auto plinks = pile->getStockpileLinks();
+    CHECK_NULL_POINTER(plinks);
+
+    if (!enable_plugin())
+        return false;
+
+    auto pworld = Core::getInstance().getWorld();
+    auto key = stl_sprintf("siege-engine/stockpiles/%d/%d", bld->id, pile->id);
+    auto entry = pworld->GetPersistentData(key, NULL);
+    if (!entry.isValid())
+        return false;
+
+    auto engine = find_engine(bld, true);
+
+    entry.ival(0) = bld->id;
+    entry.ival(1) = pile->id;
+
+    engine->stockpiles.insert(pile->id);
+
+    insert_into_vector(engine->links.take_from_pile, &df::building::id, (df::building*)pile);
+    insert_into_vector(plinks->give_to_workshop, &df::building::id, (df::building*)engine->bld);
+    return true;
+}
+
+static void forgetStockpileLink(EngineInfo *engine, int pile_id)
+{
+    engine->stockpiles.erase(pile_id);
+
+    auto pworld = Core::getInstance().getWorld();
+    auto key = stl_sprintf("siege-engine/stockpiles/%d/%d", engine->id, pile_id);
+    pworld->DeletePersistentData(pworld->GetPersistentData(key));
+}
+
+static bool removeStockpileLink(df::building_siegeenginest *bld, df::building_stockpilest *pile)
+{
+    CHECK_NULL_POINTER(bld);
+    CHECK_NULL_POINTER(pile);
+
+    if (auto engine = find_engine(bld))
+    {
+        forgetStockpileLink(engine, pile->id);
+
+        auto plinks = pile->getStockpileLinks();
+        erase_from_vector(engine->links.take_from_pile, &df::building::id, pile->id);
+        erase_from_vector(plinks->give_to_workshop, &df::building::id, bld->id);
+        return true;
+    }
+
+    return false;
+}
+
+static void recheck_pile_links(color_ostream &out, EngineInfo *engine)
+{
+    auto removed = engine->stockpiles;
+
+    out.print("rechecking piles in %d\n", engine->id);
+
+    // Detect and save changes in take links
+    for (size_t i = 0; i < engine->links.take_from_pile.size(); i++)
+    {
+        auto pile = engine->links.take_from_pile[i];
+
+        removed.erase(pile->id);
+
+        if (!engine->stockpiles.count(pile->id))
+            addStockpileLink(engine->bld, (df::building_stockpilest*)pile);
+    }
+
+    for (auto it = removed.begin(); it != removed.end(); it++)
+        forgetStockpileLink(engine, *it);
+
+    // Remove give links
+    for (size_t i = 0; i < engine->links.give_to_pile.size(); i++)
+    {
+        auto pile = engine->links.give_to_pile[i];
+        auto plinks = pile->getStockpileLinks();
+        erase_from_vector(plinks->take_from_workshop, &df::building::id, engine->id);
+    }
+
+    engine->links.give_to_pile.clear();
+}
+
 static int getShotSkill(df::building_siegeenginest *bld)
 {
     CHECK_NULL_POINTER(bld);
@@ -1122,6 +1270,25 @@ IMPLEMENT_VMETHOD_INTERPOSE(projectile_hook, checkMovement);
 struct building_hook : df::building_siegeenginest {
     typedef df::building_siegeenginest interpose_base;
 
+    DEFINE_VMETHOD_INTERPOSE(bool, canLinkToStockpile, ())
+    {
+        if (find_engine(this, true))
+            return true;
+
+        return INTERPOSE_NEXT(canLinkToStockpile)();
+    }
+
+    DEFINE_VMETHOD_INTERPOSE(df::stockpile_links*, getStockpileLinks, ())
+    {
+        if (auto engine = find_engine(this))
+        {
+            recheck_piles.insert(this);
+            return &engine->links;
+        }
+
+        return INTERPOSE_NEXT(getStockpileLinks)();
+    }
+
     DEFINE_VMETHOD_INTERPOSE(void, updateAction, ())
     {
         INTERPOSE_NEXT(updateAction)();
@@ -1186,6 +1353,8 @@ struct building_hook : df::building_siegeenginest {
     }
 };
 
+IMPLEMENT_VMETHOD_INTERPOSE(building_hook, canLinkToStockpile);
+IMPLEMENT_VMETHOD_INTERPOSE(building_hook, getStockpileLinks);
 IMPLEMENT_VMETHOD_INTERPOSE(building_hook, updateAction);
 
 /*
@@ -1195,6 +1364,9 @@ IMPLEMENT_VMETHOD_INTERPOSE(building_hook, updateAction);
 DFHACK_PLUGIN_LUA_FUNCTIONS {
     DFHACK_LUA_FUNCTION(clearTargetArea),
     DFHACK_LUA_FUNCTION(setTargetArea),
+    DFHACK_LUA_FUNCTION(isLinkedToPile),
+    DFHACK_LUA_FUNCTION(addStockpileLink),
+    DFHACK_LUA_FUNCTION(removeStockpileLink),
     DFHACK_LUA_FUNCTION(getTileStatus),
     DFHACK_LUA_FUNCTION(paintAimScreen),
     DFHACK_LUA_FUNCTION(canTargetUnit),
@@ -1208,6 +1380,7 @@ DFHACK_PLUGIN_LUA_COMMANDS {
     DFHACK_LUA_COMMAND(getTargetArea),
     DFHACK_LUA_COMMAND(getAmmoItem),
     DFHACK_LUA_COMMAND(setAmmoItem),
+    DFHACK_LUA_COMMAND(getStockpileLinks),
     DFHACK_LUA_COMMAND(projPosAtStep),
     DFHACK_LUA_COMMAND(projPathMetrics),
     DFHACK_LUA_COMMAND(adjustToTarget),
@@ -1224,6 +1397,9 @@ static void enable_hooks(bool enable)
     is_enabled = enable;
 
     INTERPOSE_HOOK(projectile_hook, checkMovement).apply(enable);
+
+    INTERPOSE_HOOK(building_hook, canLinkToStockpile).apply(enable);
+    INTERPOSE_HOOK(building_hook, getStockpileLinks).apply(enable);
     INTERPOSE_HOOK(building_hook, updateAction).apply(enable);
 
     if (enable)
@@ -1246,7 +1422,7 @@ static bool enable_plugin()
     return true;
 }
 
-static void clear_caches()
+static void clear_caches(color_ostream &out)
 {
     if (!UnitPath::cache.empty())
     {
@@ -1255,6 +1431,15 @@ static void clear_caches()
 
         UnitPath::cache.clear();
     }
+
+    if (!recheck_piles.empty())
+    {
+        for (auto it = recheck_piles.begin(); it != recheck_piles.end(); ++it)
+            if (auto engine = find_engine(*it))
+                recheck_pile_links(out, engine);
+
+        recheck_piles.clear();
+    }
 }
 
 DFhackCExport command_result plugin_onstatechange(color_ostream &out, state_change_event event)
@@ -1300,6 +1485,6 @@ DFhackCExport command_result plugin_shutdown ( color_ostream &out )
 
 DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 {
-    clear_caches();
+    clear_caches(out);
     return CR_OK;
 }

scripts/gui/siege-engine.lua

@@ -35,7 +35,7 @@ function SiegeEngine:init(building)
     self:init_fields{
         building = building,
         center = utils.getBuildingCenter(building),
-        links = {}, selected = 1,
+        selected_pile = 1,
     }
     guidm.MenuOverlay.init(self)
     self.mode_main = {
@@ -46,6 +46,10 @@ function SiegeEngine:init(building)
         render = self:callback 'onRenderBody_aim',
         input = self:callback 'onInput_aim',
     }
+    self.mode_pile = {
+        render = self:callback 'onRenderBody_pile',
+        input = self:callback 'onInput_pile',
+    }
     return self
 end
 
@@ -157,6 +161,26 @@ function SiegeEngine:renderTargetView(target_min, target_max)
     end
 end
 
+function SiegeEngine:scrollPiles(delta)
+    local links = plugin.getStockpileLinks(self.building)
+    if links then
+        self.selected_pile = 1+(self.selected_pile+delta-1) % #links
+        return links[self.selected_pile]
+    end
+end
+
+function SiegeEngine:renderStockpiles(dc, links, nlines)
+    local idx = (self.selected_pile-1) % #links
+    local page = math.floor(idx/nlines)
+    for i = page*nlines,math.min(#links,(page+1)*nlines)-1 do
+        local color = COLOR_BROWN
+        if i == idx then
+            color = COLOR_YELLOW
+        end
+        dc:newline(2):string(utils.getBuildingName(links[i+1]), color)
+    end
+end
+
 function SiegeEngine:onRenderBody_main(dc)
     dc:newline(1):pen(COLOR_WHITE):string("Target: ")
 
@@ -194,6 +218,15 @@ function SiegeEngine:onRenderBody_main(dc)
         end
     end
 
+    dc:newline():newline(1)
+    dc:string("t",COLOR_LIGHTGREEN):string(": Take from stockpile"):newline(3)
+    local links = plugin.getStockpileLinks(self.building)
+    if links then
+        dc:string("d",COLOR_LIGHTGREEN):string(": Delete, ")
+        dc:string("o",COLOR_LIGHTGREEN):string(": Zoom"):newline()
+        self:renderStockpiles(dc, links, 19-dc:localY())
+    end
+
     if self.target_select_first then
         self:renderTargetView(self.target_select_first, guidm.getCursorPos())
     else
@@ -243,6 +276,25 @@ function SiegeEngine:onInput_main(keys)
         self:zoomToTarget()
     elseif keys.CUSTOM_X then
         plugin.clearTargetArea(self.building)
+    elseif keys.SECONDSCROLL_UP then
+        self:scrollPiles(-1)
+    elseif keys.SECONDSCROLL_DOWN then
+        self:scrollPiles(1)
+    elseif keys.CUSTOM_D then
+        local pile = self:scrollPiles(0)
+        if pile then
+            plugin.removeStockpileLink(self.building, pile)
+        end
+    elseif keys.CUSTOM_O then
+        local pile = self:scrollPiles(0)
+        if pile then
+            self:centerViewOn(utils.getBuildingCenter(pile))
+        end
+    elseif keys.CUSTOM_T then
+        self:showCursor(true)
+        self.mode = self.mode_pile
+        self:sendInputToParent('CURSOR_DOWN_Z')
+        self:sendInputToParent('CURSOR_UP_Z')
     elseif self:simulateViewScroll(keys) then
         self.cursor = nil
     else
@@ -316,6 +368,57 @@ function SiegeEngine:onInput_aim(keys)
     return true
 end
 
+function SiegeEngine:onRenderBody_pile(dc)
+    dc:newline(1):string('Select pile to take from'):newline():newline(2)
+
+    local sel = df.global.world.selected_building
+
+    if df.building_stockpilest:is_instance(sel) then
+        dc:string(utils.getBuildingName(sel), COLOR_GREEN):newline():newline(1)
+
+        if plugin.isLinkedToPile(self.building, sel) then
+            dc:string("Already taking from here"):newline():newline(2)
+            dc:string("d", COLOR_LIGHTGREEN):string(": Delete link")
+        else
+            dc:string("Enter",COLOR_LIGHTGREEN):string(": Take from this pile")
+        end
+    elseif sel then
+        dc:string(utils.getBuildingName(sel), COLOR_DARKGREY)
+        dc:newline():newline(1)
+        dc:string("Not a stockpile",COLOR_LIGHTRED)
+    else
+        dc:string("No building selected", COLOR_DARKGREY)
+    end
+end
+
+function SiegeEngine:onInput_pile(keys)
+    if keys.SELECT then
+        local sel = df.global.world.selected_building
+        if df.building_stockpilest:is_instance(sel)
+        and not plugin.isLinkedToPile(self.building, sel) then
+            plugin.addStockpileLink(self.building, sel)
+
+            df.global.world.selected_building = self.building
+            self.mode = self.mode_main
+            self:showCursor(false)
+        end
+    elseif keys.CUSTOM_D then
+        local sel = df.global.world.selected_building
+        if df.building_stockpilest:is_instance(sel) then
+            plugin.removeStockpileLink(self.building, sel)
+        end
+    elseif keys.LEAVESCREEN then
+        df.global.world.selected_building = self.building
+        self.mode = self.mode_main
+        self:showCursor(false)
+    elseif self:propagateMoveKeys(keys) then
+        --
+    else
+        return false
+    end
+    return true
+end
+
 function SiegeEngine:onRenderBody(dc)
     dc:clear()
     dc:seek(1,1):pen(COLOR_WHITE):string(utils.getBuildingName(self.building)):newline()