dynamically count available materials when placing

plugins/buildingplan/buildingplan.cpp

@@ -47,11 +47,37 @@ void set_config_bool(PersistentDataItem &c, int index, bool value) {
     set_config_val(c, index, value ? 1 : 0);
 }
 
+// building type, subtype, custom
+typedef std::tuple<df::building_type, int16_t, int32_t> BuildingTypeKey;
+
+// rotates a size_t value left by count bits
+// assumes count is not 0 or >= size_t_bits
+// replace this with std::rotl when we move to C++20
+static std::size_t rotl_size_t(size_t val, uint32_t count)
+{
+    static const int size_t_bits = CHAR_BIT * sizeof(std::size_t);
+    return val << count | val >> (size_t_bits - count);
+}
+
+struct BuildingTypeKeyHash {
+    std::size_t operator() (const BuildingTypeKey & key) const {
+        // cast first param to appease gcc-4.8, which is missing the enum
+        // specializations for std::hash
+        std::size_t h1 = std::hash<int32_t>()(static_cast<int32_t>(std::get<0>(key)));
+        std::size_t h2 = std::hash<int16_t>()(std::get<1>(key));
+        std::size_t h3 = std::hash<int32_t>()(std::get<2>(key));
+
+        return h1 ^ rotl_size_t(h2, 8) ^ rotl_size_t(h3, 16);
+    }
+};
+
 static PersistentDataItem config;
+// for use in counting available materials for the UI
+static unordered_map<BuildingTypeKey, vector<df::job_item *>, BuildingTypeKeyHash> job_item_repo;
 // building id -> PlannedBuilding
-unordered_map<int32_t, PlannedBuilding> planned_buildings;
+static unordered_map<int32_t, PlannedBuilding> planned_buildings;
 // vector id -> filter bucket -> queue of (building id, job_item index)
-Tasks tasks;
+static Tasks tasks;
 
 // note that this just removes the PlannedBuilding. the tasks will get dropped
 // as we discover them in the tasks queues and they fail to be found in planned_buildings.
@@ -61,7 +87,7 @@ Tasks tasks;
 // no chance of duplicate tasks getting added to the tasks queues.
 void PlannedBuilding::remove(color_ostream &out) {
     DEBUG(status,out).print("removing persistent data for building %d\n", id);
-    World::DeletePersistentData(config);
+    World::DeletePersistentData(bld_config);
     if (planned_buildings.count(id) > 0)
         planned_buildings.erase(id);
 }
@@ -106,6 +132,31 @@ DFhackCExport command_result plugin_shutdown (color_ostream &out) {
     return CR_OK;
 }
 
+static void validate_config(color_ostream &out, bool verbose = false) {
+    if (get_config_bool(config, CONFIG_BLOCKS)
+            || get_config_bool(config, CONFIG_BOULDERS)
+            || get_config_bool(config, CONFIG_LOGS)
+            || get_config_bool(config, CONFIG_BARS))
+        return;
+
+    if (verbose)
+        out.printerr("all contruction materials disabled; resetting config\n");
+
+    set_config_bool(config, CONFIG_BLOCKS, true);
+    set_config_bool(config, CONFIG_BOULDERS, true);
+    set_config_bool(config, CONFIG_LOGS, true);
+    set_config_bool(config, CONFIG_BARS, false);
+}
+
+static void clear_job_item_repo() {
+    for (auto &entry : job_item_repo) {
+        for (auto &jitem : entry.second) {
+            delete jitem;
+        }
+    }
+    job_item_repo.clear();
+}
+
 DFhackCExport command_result plugin_load_data (color_ostream &out) {
     cycle_timestamp = 0;
     config = World::GetPersistentData(CONFIG_KEY);
@@ -113,15 +164,13 @@ DFhackCExport command_result plugin_load_data (color_ostream &out) {
     if (!config.isValid()) {
         DEBUG(status,out).print("no config found in this save; initializing\n");
         config = World::AddPersistentData(CONFIG_KEY);
-        set_config_bool(config, CONFIG_BLOCKS, true);
-        set_config_bool(config, CONFIG_BOULDERS, true);
-        set_config_bool(config, CONFIG_LOGS, true);
-        set_config_bool(config, CONFIG_BARS, false);
     }
+    validate_config(out);
 
     DEBUG(status,out).print("loading persisted state\n");
     planned_buildings.clear();
     tasks.clear();
+    clear_job_item_repo();
     vector<PersistentDataItem> building_configs;
     World::GetPersistentData(&building_configs, BLD_CONFIG_KEY);
     const size_t num_building_configs = building_configs.size();
@@ -138,30 +187,11 @@ DFhackCExport command_result plugin_onstatechange(color_ostream &out, state_chan
         DEBUG(status,out).print("world unloaded; clearing state for %s\n", plugin_name);
         planned_buildings.clear();
         tasks.clear();
+        clear_job_item_repo();
     }
     return CR_OK;
 }
 
-static bool cycle_requested = false;
-
-static void do_cycle(color_ostream &out) {
-    // mark that we have recently run
-    cycle_timestamp = world->frame_counter;
-    cycle_requested = false;
-
-    buildingplan_cycle(out, tasks, planned_buildings);
-}
-
-DFhackCExport command_result plugin_onupdate(color_ostream &out) {
-    if (!Core::getInstance().isWorldLoaded())
-        return CR_OK;
-
-    if (is_enabled &&
-            (cycle_requested || world->frame_counter - cycle_timestamp >= CYCLE_TICKS))
-        do_cycle(out);
-    return CR_OK;
-}
-
 static bool call_buildingplan_lua(color_ostream *out, const char *fn_name,
         int nargs = 0, int nres = 0,
         Lua::LuaLambda && args_lambda = Lua::DEFAULT_LUA_LAMBDA,
@@ -182,6 +212,27 @@ static bool call_buildingplan_lua(color_ostream *out, const char *fn_name,
             std::forward<Lua::LuaLambda&&>(res_lambda));
 }
 
+static bool cycle_requested = false;
+
+static void do_cycle(color_ostream &out) {
+    // mark that we have recently run
+    cycle_timestamp = world->frame_counter;
+    cycle_requested = false;
+
+    buildingplan_cycle(out, tasks, planned_buildings);
+    call_buildingplan_lua(&out, "reset_counts");
+}
+
+DFhackCExport command_result plugin_onupdate(color_ostream &out) {
+    if (!Core::getInstance().isWorldLoaded())
+        return CR_OK;
+
+    if (is_enabled &&
+            (cycle_requested || world->frame_counter - cycle_timestamp >= CYCLE_TICKS))
+        do_cycle(out);
+    return CR_OK;
+}
+
 static command_result do_command(color_ostream &out, vector<string> &parameters) {
     CoreSuspender suspend;
 
@@ -228,8 +279,7 @@ static string getBucket(const df::job_item & ji) {
 }
 
 // get a list of item vectors that we should search for matches
-static vector<df::job_item_vector_id> getVectorIds(color_ostream &out, df::job_item *job_item)
-{
+static vector<df::job_item_vector_id> getVectorIds(color_ostream &out, df::job_item *job_item) {
     std::vector<df::job_item_vector_id> ret;
 
     // if the filter already has the vector_id set to something specific, use it
@@ -344,7 +394,7 @@ static void printStatus(color_ostream &out) {
         }
     }
 
-    out.print("Waiting for %d item(s) to be produced or %zd building(s):\n",
+    out.print("Waiting for %d item(s) to be produced for %zd building(s):\n",
               total, planned_buildings.size());
     for (auto &count : counts)
         out.print("  %3d %s\n", count.second, count.first.c_str());
@@ -365,6 +415,9 @@ static bool setSetting(color_ostream &out, string name, bool value) {
         out.printerr("unrecognized setting: '%s'\n", name.c_str());
         return false;
     }
+
+    validate_config(out, true);
+    call_buildingplan_lua(&out, "reset_counts");
     return true;
 }
 
@@ -412,7 +465,49 @@ static void scheduleCycle(color_ostream &out) {
 
 static int countAvailableItems(color_ostream &out, df::building_type type, int16_t subtype, int32_t custom, int index) {
     DEBUG(status,out).print("entering countAvailableItems\n");
-    return 10;
+    DEBUG(status,out).print(
+            "entering countAvailableItems building_type=%d subtype=%d custom=%d index=%d\n",
+            type, subtype, custom, index);
+    BuildingTypeKey key(type, subtype, custom);
+    auto &job_items = job_item_repo[key];
+    if (index >= job_items.size()) {
+        for (int i = job_items.size(); i <= index; ++i) {
+            bool failed = false;
+            if (!call_buildingplan_lua(&out, "get_job_item", 4, 1,
+                    [&](lua_State *L) {
+                        Lua::Push(L, type);
+                        Lua::Push(L, subtype);
+                        Lua::Push(L, custom);
+                        Lua::Push(L, index+1);
+                    },
+                    [&](lua_State *L) {
+                        df::job_item *jitem = Lua::GetDFObject<df::job_item>(L, -1);
+                        DEBUG(status,out).print("retrieving job_item for index=%d: %p\n",
+                                index, jitem);
+                        if (!jitem)
+                            failed = true;
+                        else
+                            job_items.emplace_back(jitem);
+                    }) || failed) {
+                return 0;
+            }
+        }
+    }
+
+    auto &jitem = job_items[index];
+    auto vector_ids = getVectorIds(out, jitem);
+
+    int count = 0;
+    for (auto vector_id : vector_ids) {
+        auto other_id = ENUM_ATTR(job_item_vector_id, other, vector_id);
+        for (auto &item : df::global::world->items.other[other_id]) {
+            if (itemPassesScreen(item) && matchesFilters(item, jitem))
+                ++count;
+        }
+    }
+
+    DEBUG(status,out).print("found matches %d\n", count);
+    return count;
 }
 
 DFHACK_PLUGIN_LUA_FUNCTIONS {

plugins/buildingplan/buildingplan.h

@@ -2,6 +2,7 @@
 
 #include "modules/Persistence.h"
 
+#include "df/job_item.h"
 #include "df/job_item_vector_id.h"
 
 #include <deque>
@@ -26,3 +27,6 @@ int get_config_val(DFHack::PersistentDataItem &c, int index);
 bool get_config_bool(DFHack::PersistentDataItem &c, int index);
 void set_config_val(DFHack::PersistentDataItem &c, int index, int value);
 void set_config_bool(DFHack::PersistentDataItem &c, int index, bool value);
+
+bool itemPassesScreen(df::item * item);
+bool matchesFilters(df::item * item, df::job_item * job_item);

plugins/buildingplan/buildingplan_cycle.cpp

@@ -10,7 +10,6 @@
 #include "df/building_design.h"
 #include "df/item.h"
 #include "df/job.h"
-#include "df/job_item.h"
 #include "df/world.h"
 
 #include <unordered_map>
@@ -41,13 +40,13 @@ struct BadFlags {
     }
 };
 
-static bool itemPassesScreen(df::item * item) {
+bool itemPassesScreen(df::item * item) {
     static const BadFlags bad_flags;
     return !(item->flags.whole & bad_flags.whole)
         && !item->isAssignedToStockpile();
 }
 
-static bool matchesFilters(df::item * item, df::job_item * job_item) {
+bool matchesFilters(df::item * item, df::job_item * job_item) {
     // check the properties that are not checked by Job::isSuitableItem()
     if (job_item->item_type > -1 && job_item->item_type != item->getType())
         return false;

plugins/lua/buildingplan.lua

@@ -56,6 +56,19 @@ function get_num_filters(btype, subtype, custom)
     return filters and #filters or 0
 end
 
+function get_job_item(btype, subtype, custom, index)
+    local filters = dfhack.buildings.getFiltersByType({}, btype, subtype, custom)
+    if not filters or not filters[index] then return nil end
+    local obj = df.job_item:new()
+    obj:assign(filters[index])
+    return obj
+end
+
+local reset_counts_flag = false
+function reset_counts()
+    reset_counts_flag = true
+end
+
 --------------------------------
 -- Planner Overlay
 --
@@ -143,8 +156,32 @@ local function to_title_case(str)
     return str
 end
 
-function get_item_line_text(idx)
-    local filter = get_cur_filters()[idx]
+ItemLine = defclass(ItemLine, widgets.Panel)
+ItemLine.ATTRS{
+    idx=DEFAULT_NIL,
+}
+
+function ItemLine:init()
+    self.frame.h = 1
+    self.visible = function() return #get_cur_filters() >= self.idx end
+    self:addviews{
+        widgets.Label{
+            frame={t=0, l=0},
+            text={{text=function() return self:get_item_line_text() end}},
+        },
+        widgets.Label{
+            frame={t=0, l=22},
+            text='[filter][x]',
+        },
+    }
+end
+
+function ItemLine:reset()
+    self.desc = nil
+    self.available = nil
+end
+
+local function get_desc(filter)
     local desc = 'Unknown'
     if filter.has_tool_use then
         desc = to_title_case(df.tool_uses[filter.has_tool_use])
@@ -170,7 +207,10 @@ function get_item_line_text(idx)
     if desc == 'Trappart' then
         desc = 'Mechanism'
     end
+    return desc
+end
 
+local function get_quantity(filter)
     local quantity = filter.quantity or 1
     local dimx, dimy, dimz = get_cur_area_dims()
     if quantity < 1 then
@@ -178,34 +218,29 @@ function get_item_line_text(idx)
     else
         quantity = quantity * dimx * dimy * dimz
     end
-    desc = ('%d %s%s'):format(quantity, desc, quantity == 1 and '' or 's')
+    return quantity
+end
 
-    local available = countAvailableItems(uibs.building_type,
+function ItemLine:get_item_line_text()
+    local idx = self.idx
+    local filter = get_cur_filters()[idx]
+    local quantity = get_quantity(filter)
+
+    self.desc = self.desc or get_desc(filter)
+    local line = ('%d %s%s'):format(quantity, self.desc, quantity == 1 and '' or 's')
+
+    self.available = self.available or countAvailableItems(uibs.building_type,
             uibs.building_subtype, uibs.custom_type, idx - 1)
-    local note = available >= quantity and
+    local note = self.available >= quantity and
             'Can build now' or 'Will wait for item'
 
-    return ('%-21s%s%s'):format(desc:sub(1,21), (' '):rep(13), note)
+    return ('%-21s%s%s'):format(line:sub(1,21), (' '):rep(13), note)
 end
 
-ItemLine = defclass(ItemLine, widgets.Panel)
-ItemLine.ATTRS{
-    idx=DEFAULT_NIL,
-}
-
-function ItemLine:init()
-    self.frame.h = 1
-    self.visible = function() return #get_cur_filters() >= self.idx end
-    self:addviews{
-        widgets.Label{
-            frame={t=0, l=0},
-            text={{text=function() return get_item_line_text(self.idx) end}},
-        },
-        widgets.Label{
-            frame={t=0, l=22},
-            text='[filter][x]',
-        },
-    }
+function ItemLine:reduce_quantity()
+    if not self.available then return end
+    local filter = get_cur_filters()[self.idx]
+    self.available = math.max(0, self.available - get_quantity(filter))
 end
 
 PlannerOverlay = defclass(PlannerOverlay, overlay.OverlayWidget)
@@ -226,10 +261,10 @@ function PlannerOverlay:init()
             text='No items required.',
             visible=function() return #get_cur_filters() == 0 end,
         },
-        ItemLine{frame={t=0, l=0}, idx=1},
-        ItemLine{frame={t=2, l=0}, idx=2},
-        ItemLine{frame={t=4, l=0}, idx=3},
-        ItemLine{frame={t=6, l=0}, idx=4},
+        ItemLine{view_id='item1', frame={t=0, l=0}, idx=1},
+        ItemLine{view_id='item2', frame={t=2, l=0}, idx=2},
+        ItemLine{view_id='item3', frame={t=4, l=0}, idx=3},
+        ItemLine{view_id='item4', frame={t=6, l=0}, idx=4},
         widgets.CycleHotkeyLabel{
             view_id="stairs_top_subtype",
             frame={t=3, l=0},
@@ -268,13 +303,22 @@ function PlannerOverlay:init()
     }
 end
 
-function PlannerOverlay:do_config()
-    dfhack.run_script('gui/buildingplan')
+function PlannerOverlay:reset()
+    self.subviews.item1:reset()
+    self.subviews.item2:reset()
+    self.subviews.item3:reset()
+    self.subviews.item4:reset()
+    reset_counts_flag = false
 end
 
 function PlannerOverlay:onInput(keys)
     if not is_plannable() then return false end
     if keys.LEAVESCREEN or keys._MOUSE_R_DOWN then
+        if uibs.selection_pos:isValid() then
+            uibs.selection_pos:clear()
+            return true
+        end
+        self:reset()
         return false
     end
     if PlannerOverlay.super.onInput(self, keys) then
@@ -290,6 +334,10 @@ function PlannerOverlay:onInput(keys)
                 if #get_cur_filters() == 0 then
                     return false -- we don't add value; let the game place it
                 end
+                self.subviews.item1:reduce_quantity()
+                self.subviews.item2:reduce_quantity()
+                self.subviews.item3:reduce_quantity()
+                self.subviews.item4:reduce_quantity()
                 self:place_building()
                 uibs.selection_pos:clear()
                 return true
@@ -315,6 +363,10 @@ local BAD_PEN = to_pen{ch='X', fg=COLOR_RED,
 function PlannerOverlay:onRenderFrame(dc, rect)
     PlannerOverlay.super.onRenderFrame(self, dc, rect)
 
+    if reset_counts_flag then
+        self:reset()
+    end
+
     if not is_choosing_area() then return end
 
     local bounds = {