implement mechanism selection for lever linking

docs/changelog.txt

@@ -57,6 +57,7 @@ Template for new versions:
 ## New Features
 - `prospect`: can now give you an estimate of resources from the embark screen. hover the mouse over a potential embark area and run `prospect`.
 - `burrow`: integrated 3d box fill and 2d/3d flood fill extensions for burrow painting mode
+- `buildingplan`: allow specific mechanisms to be selected when linking levers
 
 ## Fixes
 - `stockpiles`: hide configure and help buttons when the overlay panel is minimized

plugins/buildingplan/buildingplan.cpp

@@ -674,13 +674,14 @@ static void scheduleCycle(color_ostream &out) {
 }
 
 static int scanAvailableItems(color_ostream &out, df::building_type type, int16_t subtype,
-        int32_t custom, int index, bool ignore_filters, vector<int> *item_ids = NULL,
-        map<MaterialInfo, int32_t> *counts = NULL) {
+        int32_t custom, int index, bool ignore_filters, bool ignore_quality, HeatSafety *heat_override = NULL,
+        vector<int> *item_ids = NULL, map<MaterialInfo, int32_t> *counts = NULL)
+{
     DEBUG(status,out).print(
             "entering scanAvailableItems building_type=%d subtype=%d custom=%d index=%d\n",
             type, subtype, custom, index);
     BuildingTypeKey key(type, subtype, custom);
-    HeatSafety heat = get_heat_safety_filter(key);
+    HeatSafety heat = heat_override ? *heat_override : get_heat_safety_filter(key);
     auto &job_items = get_job_items(out, key);
     if (index < 0 || job_items.size() <= (size_t)index)
         return 0;
@@ -703,6 +704,10 @@ static int scanAvailableItems(color_ostream &out, df::building_type type, int16_
                 filter.setMaterials(set<MaterialInfo>());
                 special.clear();
             }
+            if (ignore_quality) {
+                filter.setMinQuality(df::item_quality::Ordinary);
+                filter.setMaxQuality(df::item_quality::Artifact);
+            }
             if (itemPassesScreen(out, item) && matchesFilters(item, jitem, heat, filter, special)) {
                 if (item_ids)
                     item_ids->emplace_back(item->id);
@@ -732,7 +737,25 @@ static int getAvailableItems(lua_State *L) {
             "entering getAvailableItems building_type=%d subtype=%d custom=%d index=%d\n",
             type, subtype, custom, index);
     vector<int> item_ids;
-    scanAvailableItems(*out, type, subtype, custom, index, true, &item_ids);
+    scanAvailableItems(*out, type, subtype, custom, index, true, false, NULL, &item_ids);
+    Lua::PushVector(L, item_ids);
+    return 1;
+}
+
+static int getAvailableItemsByHeat(lua_State *L) {
+    color_ostream *out = Lua::GetOutput(L);
+    if (!out)
+        out = &Core::getInstance().getConsole();
+    df::building_type type = (df::building_type)luaL_checkint(L, 1);
+    int16_t subtype = luaL_checkint(L, 2);
+    int32_t custom = luaL_checkint(L, 3);
+    int index = luaL_checkint(L, 4);
+    HeatSafety heat = (HeatSafety)luaL_checkint(L, 5);
+    DEBUG(status,*out).print(
+            "entering getAvailableItemsByHeat building_type=%d subtype=%d custom=%d index=%d\n",
+            type, subtype, custom, index);
+    vector<int> item_ids;
+    scanAvailableItems(*out, type, subtype, custom, index, true, true, &heat, &item_ids);
     Lua::PushVector(L, item_ids);
     return 1;
 }
@@ -755,7 +778,7 @@ static int countAvailableItems(color_ostream &out, df::building_type type, int16
     DEBUG(status,out).print(
             "entering countAvailableItems building_type=%d subtype=%d custom=%d index=%d\n",
             type, subtype, custom, index);
-    int count = scanAvailableItems(out, type, subtype, custom, index, false);
+    int count = scanAvailableItems(out, type, subtype, custom, index, false, false);
     if (count)
         return count;
 
@@ -968,7 +991,7 @@ static int getMaterialFilter(lua_State *L) {
         return 0;
     const auto &mat_filter = filters[index].getMaterials();
     map<MaterialInfo, int32_t> counts;
-    scanAvailableItems(*out, type, subtype, custom, index, false, NULL, &counts);
+    scanAvailableItems(*out, type, subtype, custom, index, false, false, NULL, NULL, &counts);
     HeatSafety heat = get_heat_safety_filter(key);
     const df::job_item *jitem = get_job_items(*out, key)[index];
     // name -> {count=int, enabled=bool, category=string, heat=string}
@@ -1232,6 +1255,7 @@ DFHACK_PLUGIN_LUA_FUNCTIONS {
 DFHACK_PLUGIN_LUA_COMMANDS {
     DFHACK_LUA_COMMAND(getGlobalSettings),
     DFHACK_LUA_COMMAND(getAvailableItems),
+    DFHACK_LUA_COMMAND(getAvailableItemsByHeat),
     DFHACK_LUA_COMMAND(setMaterialMaskFilter),
     DFHACK_LUA_COMMAND(getMaterialMaskFilter),
     DFHACK_LUA_COMMAND(setMaterialFilter),

plugins/lua/buildingplan/itemselection.lua

@@ -21,6 +21,29 @@ function get_automaterial_selection(building_type)
     return tracker.list[#tracker.list]
 end
 
+local function get_artifact_name(item)
+    local gref = dfhack.items.getGeneralRef(item, df.general_ref_type.IS_ARTIFACT)
+    if not gref then return end
+    local artifact = df.artifact_record.find(gref.artifact_id)
+    if not artifact then return end
+    return dfhack.TranslateName(artifact.name)
+end
+
+function get_item_description(item_id, item, safety_label)
+    item = item or df.item.find(item_id)
+    if not item then
+        return ('No %s safe mechanisms available'):format(safety_label:lower())
+    end
+    local desc = item.flags.artifact and get_artifact_name(item) or
+        dfhack.items.getDescription(item, 0, true)
+    local wear_level = item:getWear()
+    if wear_level == 1 then desc = ('x%sx'):format(desc)
+    elseif wear_level == 2 then desc = ('X%sX'):format(desc)
+    elseif wear_level == 3 then desc = ('XX%sXX'):format(desc)
+    end
+    return desc
+end
+
 local function sort_by_type(a, b)
     local ad, bd = a.data, b.data
     return ad.item_type < bd.item_type or
@@ -57,7 +80,7 @@ ItemSelection.ATTRS{
     frame_title='Choose items',
     frame={w=56, h=24, l=4, t=7},
     resizable=true,
-    index=DEFAULT_NIL,
+    get_available_items_fn=DEFAULT_NIL,
     desc=DEFAULT_NIL,
     quantity=DEFAULT_NIL,
     autoselect=DEFAULT_NIL,
@@ -99,9 +122,9 @@ function ItemSelection:init()
                     text_pen=BUILD_TEXT_PEN,
                     text_hpen=BUILD_TEXT_HPEN,
                     text={
-                        '   Use filter  ', NEWLINE,
-                        ' for remaining ', NEWLINE,
-                        '     items     ',
+                        '               ', NEWLINE,
+                        '   Autoselect  ', NEWLINE,
+                        '               ',
                     },
                     on_click=self:callback('submit'),
                     visible=function() return self.num_selected < self.quantity end,
@@ -238,13 +261,12 @@ local function make_search_key(str)
 end
 
 function ItemSelection:get_choices(sort_fn)
-    local item_ids = require('plugins.buildingplan').getAvailableItems(uibs.building_type,
-            uibs.building_subtype, uibs.custom_type, self.index-1)
+    local item_ids = self.get_available_items_fn()
     local buckets = {}
     for _,item_id in ipairs(item_ids) do
         local item = df.item.find(item_id)
         if not item then goto continue end
-        local desc = dfhack.items.getDescription(item, 0, true)
+        local desc = get_item_description(item_id, item)
         if buckets[desc] then
             local bucket = buckets[desc]
             table.insert(bucket.data.item_ids, item_id)
@@ -396,7 +418,7 @@ ItemSelectionScreen.ATTRS {
     pass_pause=false,
     pass_mouse_clicks=false,
     defocusable=false,
-    index=DEFAULT_NIL,
+    get_available_items_fn=DEFAULT_NIL,
     desc=DEFAULT_NIL,
     quantity=DEFAULT_NIL,
     autoselect=DEFAULT_NIL,
@@ -407,7 +429,7 @@ ItemSelectionScreen.ATTRS {
 function ItemSelectionScreen:init()
     self:addviews{
         ItemSelection{
-            index=self.index,
+            get_available_items_fn=self.get_available_items_fn,
             desc=self.desc,
             quantity=self.quantity,
             autoselect=self.autoselect,

plugins/lua/buildingplan/mechanisms.lua

@@ -1,8 +1,11 @@
 local _ENV = mkmodule('plugins.buildingplan.mechanisms')
 
+local itemselection = require('plugins.buildingplan.itemselection')
 local overlay = require('plugins.overlay')
 local widgets = require('gui.widgets')
 
+local view_sheets = df.global.game.main_interface.view_sheets
+
 --------------------------------
 -- MechanismOverlay
 --
@@ -17,53 +20,148 @@ MechanismOverlay.ATTRS{
 
 function MechanismOverlay:init()
     self:addviews{
+        widgets.Label{
+            frame={t=4, l=2},
+            text='Mechanism safety:'
+        },
         widgets.CycleHotkeyLabel{
-            view_id='safety1',
-            frame={t=4, l=4, w=40},
+            view_id='safety_lever',
+            frame={t=4, l=21, w=15},
             key='CUSTOM_G',
-            label='Lever mechanism safety:',
+            label='Lever:',
             options={
                 {label='Any', value=0},
                 {label='Magma', value=2, pen=COLOR_RED},
                 {label='Fire', value=1, pen=COLOR_LIGHTRED},
             },
             initial_option=0,
+            on_change=self:callback('choose_mechanism', 'lever', true),
         },
         widgets.CycleHotkeyLabel{
-            view_id='safety2',
-            frame={t=5, l=4, w=40},
+            view_id='safety_target',
+            frame={t=4, l=38, w=16},
             key='CUSTOM_SHIFT_G',
-            label='Target mechanism safety:',
+            label='Target:',
             options={
                 {label='Any', value=0},
                 {label='Magma', value=2, pen=COLOR_RED},
                 {label='Fire', value=1, pen=COLOR_LIGHTRED},
             },
             initial_option=0,
+            on_change=self:callback('choose_mechanism', 'target', true),
         },
         widgets.HotkeyLabel{
-            frame={t=6, l=9, w=48, h=3},
+            frame={t=7, l=8, w=49, h=2},
             key='CUSTOM_M',
-            label='Choose',
+            label=function()
+                return itemselection.get_item_description(view_sheets.linking_lever_mech_lever_id,
+                    nil,
+                    self.subviews.safety_lever:getOptionLabel())
+            end,
             auto_height=false,
-            on_activate=function() end,
+            enabled=function() return view_sheets.linking_lever_mech_lever_id ~= -1 end,
+            on_activate=self:callback('choose_mechanism', 'lever', false),
         },
         widgets.HotkeyLabel{
-            frame={t=9, l=9, w=48, h=3},
+            frame={t=10, l=8, w=49, h=2},
             key='CUSTOM_SHIFT_M',
-            label='Choose',
+            label=function()
+                return itemselection.get_item_description(view_sheets.linking_lever_mech_target_id,
+                    nil,
+                    self.subviews.safety_target:getOptionLabel())
+            end,
             auto_height=false,
-            on_activate=function() end,
+            enabled=function() return view_sheets.linking_lever_mech_target_id ~= -1 end,
+            on_activate=self:callback('choose_mechanism', 'target', false),
         },
     }
 end
 
+local item_selection_dlg
+local function reset_dlg()
+    if item_selection_dlg then
+        if item_selection_dlg:isActive() then
+            item_selection_dlg:dismiss()
+        end
+        item_selection_dlg = nil
+    end
+end
+
+local function get_available_items(safety, other_mechanism)
+    local item_ids = require('plugins.buildingplan').getAvailableItemsByHeat(
+        df.building_type.Trap, df.trap_type.Lever, -1, 0, safety)
+    for idx,item_id in ipairs(item_ids) do
+        if item_id == other_mechanism then
+            table.remove(item_ids, idx)
+            break
+        end
+    end
+    return item_ids
+end
+
+function MechanismOverlay:save_id(which, item_id)
+    local saved_id = ('saved_%s_id'):format(which)
+    local ui_id = ('linking_lever_mech_%s_id'):format(which)
+    view_sheets[ui_id] = item_id
+    self[saved_id] = item_id
+end
+
+function MechanismOverlay:choose_mechanism(which, autoselect)
+    local widget_id = 'safety_' .. which
+    local safety = self.subviews[widget_id]:getOptionValue()
+    local ui_other_id = ('linking_lever_mech_%s_id'):format(which == 'lever' and 'target' or 'lever')
+    local available_item_ids = get_available_items(safety, view_sheets[ui_other_id])
+
+    if autoselect then
+        self:save_id(which, available_item_ids[1] or -1)
+        return
+    end
+
+    -- to integrate with ItemSelection's last used sorting
+    df.global.buildreq.building_type = df.building_type.Trap
+
+    local desc = self.subviews[widget_id]:getOptionLabel()
+    if desc ~= 'Any' then
+        desc = desc .. ' safe'
+    end
+    desc = desc .. ' mechanism'
+
+    item_selection_dlg = item_selection_dlg or itemselection.ItemSelectionScreen{
+        get_available_items_fn=function() return available_item_ids end,
+        desc=desc,
+        quantity=1,
+        autoselect=false,
+        on_cancel=reset_dlg,
+        on_submit=function(chosen_ids)
+            self:save_id(which, chosen_ids[1] or available_item_ids[1] -1)
+            reset_dlg()
+        end,
+    }:show()
+end
+
 function MechanismOverlay:onInput(keys)
-    if keys._MOUSE_L and self:getMousePos() then
-        MechanismOverlay.super.onInput(self, keys)
+    if MechanismOverlay.super.onInput(self, keys) then
+        return true
+    end
+    if keys._MOUSE_L then
+        if self:getMousePos() then
             -- don't let clicks bleed through the panel
             return true
         end
+        -- don't allow the lever to be linked if mechanisms are not set
+        return view_sheets.linking_lever_mech_lever_id == -1 or
+            view_sheets.linking_lever_mech_target_id == -1
+    end
+end
+
+function MechanismOverlay:onRenderFrame(dc, rect)
+    MechanismOverlay.super.onRenderFrame(self, dc, rect)
+    if self.saved_lever_id ~= view_sheets.linking_lever_mech_lever_id then
+        self:choose_mechanism('lever', true)
+    end
+    if self.saved_target_id ~= view_sheets.linking_lever_mech_target_id then
+        self:choose_mechanism('target', true)
+    end
 end
 
 return _ENV

plugins/lua/buildingplan/planneroverlay.lua

@@ -809,8 +809,12 @@ function PlannerOverlay:onInput(keys)
                     for idx = num_filters,1,-1 do
                         chosen_items[idx] = {}
                         local filter = filters[idx]
+                        local get_available_items_fn = function()
+                            return require('plugins.buildingplan').getAvailableItems(
+                                uibs.building_type, uibs.building_subtype, uibs.custom_type, idx-1)
+                        end
                         local selection_screen = itemselection.ItemSelectionScreen{
-                            index=idx,
+                            get_available_items_fn=get_available_items_fn,
                             desc=require('plugins.buildingplan').get_desc(filter),
                             quantity=get_quantity(filter, is_hollow, bounds),
                             autoselect=autoselect,