military and burrow membership filters for burrow assignment

docs/changelog.txt

@@ -58,6 +58,7 @@ Template for new versions:
 - `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
+- `sort`: military and burrow membership filters for the burrow assignment screen
 
 ## Fixes
 - `stockpiles`: hide configure and help buttons when the overlay panel is minimized

plugins/lua/sort.lua

@@ -1295,6 +1295,7 @@ OVERLAY_WIDGETS = {
     location_selector=require('plugins.sort.locationselector').LocationSelectorOverlay,
     unit_selector=require('plugins.sort.unitselector').UnitSelectorOverlay,
     worker_assignment=require('plugins.sort.unitselector').WorkerAssignmentOverlay,
+    burrow_assignment=require('plugins.sort.unitselector').BurrowAssignmentOverlay,
     slab=require('plugins.sort.slab').SlabOverlay,
     world=require('plugins.sort.world').WorldOverlay,
 }

plugins/lua/sort/info.lua

@@ -165,7 +165,7 @@ InfoOverlay.ATTRS{
     frame={w=40, h=6},
 }
 
-local function get_squad_options()
+function get_squad_options()
     local options = {{label='Any', value='all', pen=COLOR_GREEN}}
     local fort = df.historical_entity.find(df.global.plotinfo.group_id)
     if not fort then return options end
@@ -179,7 +179,7 @@ local function get_squad_options()
     return options
 end
 
-local function get_burrow_options()
+function get_burrow_options()
     local options = {
         {label='Any', value='all', pen=COLOR_GREEN},
         {label='Unburrowed', value='none', pen=COLOR_LIGHTRED},
@@ -194,6 +194,25 @@ local function get_burrow_options()
     return options
 end
 
+function matches_squad_burrow_filters(unit, subset, target_squad_id, target_burrow_id)
+    if subset == 'all' then
+        return true
+    elseif subset == 'civilian' then
+        return unit.military.squad_id == -1
+    elseif subset == 'military' then
+        local squad_id = unit.military.squad_id
+        if squad_id == -1 then return false end
+        if target_squad_id == 'all' then return true end
+        return target_squad_id == squad_id
+    elseif subset == 'burrow' then
+        if target_burrow_id == 'all' then return #unit.burrows + #unit.inactive_burrows > 0 end
+        if target_burrow_id == 'none' then return #unit.burrows + #unit.inactive_burrows == 0 end
+        return utils.binsearch(unit.burrows, target_burrow_id) or
+            utils.binsearch(unit.inactive_burrows, target_burrow_id)
+    end
+    return true
+end
+
 function InfoOverlay:init()
     self:addviews{
         widgets.BannerPanel{
@@ -217,7 +236,7 @@ function InfoOverlay:init()
             subviews={
                 widgets.CycleHotkeyLabel{
                     view_id='subset',
-                    frame={l=1, t=0},
+                    frame={l=1, t=0, r=1},
                     key='CUSTOM_SHIFT_F',
                     label='Show:',
                     options={
@@ -255,7 +274,7 @@ function InfoOverlay:init()
             subviews={
                 widgets.CycleHotkeyLabel{
                     view_id='squad',
-                    frame={l=1, t=0},
+                    frame={l=1, t=0, r=1},
                     key='CUSTOM_SHIFT_S',
                     label='Squad:',
                     options={
@@ -266,7 +285,7 @@ function InfoOverlay:init()
                 },
                 widgets.CycleHotkeyLabel{
                     view_id='burrow',
-                    frame={l=1, t=0},
+                    frame={l=1, t=0, r=1},
                     key='CUSTOM_SHIFT_B',
                     label='Burrow:',
                     options={
@@ -400,24 +419,8 @@ function InfoOverlay:onInput(keys)
 end
 
 function InfoOverlay:matches_filters(unit)
-    local subset = self.subviews.subset:getOptionValue()
-    if subset == 'all' then
-        return true
-    elseif subset == 'civilian' then
-        return unit.military.squad_id == -1
-    elseif subset == 'military' then
-        local squad_id = unit.military.squad_id
-        if squad_id == -1 then return false end
-        local target_id = self.subviews.squad:getOptionValue()
-        if target_id == 'all' then return true end
-        return target_id == squad_id
-    elseif subset == 'burrow' then
-        local target_id = self.subviews.burrow:getOptionValue()
-        if target_id == 'all' then return #unit.burrows + #unit.inactive_burrows > 0 end
-        if target_id == 'none' then return #unit.burrows + #unit.inactive_burrows == 0 end
-        return utils.binsearch(unit.burrows, target_id) or utils.binsearch(unit.inactive_burrows, target_id)
-    end
-    return true
+    return matches_squad_burrow_filters(unit, self.subviews.subset:getOptionValue(),
+        self.subviews.squad:getOptionValue(), self.subviews.burrow:getOptionValue())
 end
 
 -- ----------------------

plugins/lua/sort/sortoverlay.lua

@@ -5,9 +5,9 @@ local utils = require('utils')
 
 function get_unit_search_key(unit)
     return ('%s %s %s'):format(
-        dfhack.units.getReadableName(unit),  -- last name is in english
+        dfhack.units.getReadableName(unit),
         dfhack.units.getProfessionName(unit),
-        dfhack.TranslateName(unit.name, false, true))  -- get untranslated last name
+        dfhack.TranslateName(unit.name, true, true))  -- get English last name
 end
 
 local function copy_to_lua_table(vec)

plugins/lua/sort/unitselector.lua

@@ -1,6 +1,9 @@
 local _ENV = mkmodule('plugins.sort.unitselector')
 
+local info = require('plugins.sort.info')
+local gui = require('gui')
 local sortoverlay = require('plugins.sort.sortoverlay')
+local utils = require('utils')
 local widgets = require('gui.widgets')
 
 local unit_selector = df.global.game.main_interface.unit_selector
@@ -9,12 +12,25 @@ local unit_selector = df.global.game.main_interface.unit_selector
 -- UnitSelectorOverlay
 --
 
+local WIDGET_WIDTH = 31
+
 UnitSelectorOverlay = defclass(UnitSelectorOverlay, sortoverlay.SortOverlay)
 UnitSelectorOverlay.ATTRS{
     default_pos={x=62, y=6},
     viewscreens='dwarfmode/UnitSelector',
     frame={w=31, h=1},
-    handled_screens=DEFAULT_NIL,
+    -- pen, pit, chain, and cage assignment are handled by dedicated screens
+    -- squad fill position screen has a specialized overlay
+    -- we *could* add search functionality to vanilla screens for pit and cage,
+    -- but then we'd have to handle the itemid vector
+    handled_screens={
+        ZONE_BEDROOM_ASSIGNMENT='already',
+        ZONE_OFFICE_ASSIGNMENT='already',
+        ZONE_DINING_HALL_ASSIGNMENT='already',
+        ZONE_TOMB_ASSIGNMENT='already',
+        OCCUPATION_ASSIGNMENT='selected',
+        SQUAD_KILL_ORDER='selected',
+    },
 }
 
 local function get_unit_id_search_key(unit_id)
@@ -26,7 +42,7 @@ end
 function UnitSelectorOverlay:init()
     self:addviews{
         widgets.BannerPanel{
-            frame={l=0, t=0, r=0, h=1},
+            frame={l=0, t=0, w=WIDGET_WIDTH, h=1},
             visible=self:callback('get_key'),
             subviews={
                 widgets.EditField{
@@ -40,20 +56,10 @@ function UnitSelectorOverlay:init()
         },
     }
 
-    -- pen, pit, chain, and cage assignment are handled by dedicated screens
-    -- squad fill position screen has a specialized overlay
-    -- we *could* add search functionality to vanilla screens for pit and cage,
-    -- but then we'd have to handle the itemid vector
-    self.handled_screens = self.handled_screens or {
-        ZONE_BEDROOM_ASSIGNMENT='already',
-        ZONE_OFFICE_ASSIGNMENT='already',
-        ZONE_DINING_HALL_ASSIGNMENT='already',
-        ZONE_TOMB_ASSIGNMENT='already',
-        OCCUPATION_ASSIGNMENT='selected',
-        BURROW_ASSIGNMENT='selected',
-        SQUAD_KILL_ORDER='selected',
-    }
+    self:register_handlers()
+end
 
+function UnitSelectorOverlay:register_handlers()
     for name,flags_vec in pairs(self.handled_screens) do
         self:register_handler(name, unit_selector.unid,
             curry(sortoverlay.flags_vector_search, {get_search_key_fn=get_unit_id_search_key},
@@ -77,13 +83,15 @@ function UnitSelectorOverlay:onRenderBody(dc)
 end
 
 function UnitSelectorOverlay:onInput(keys)
+    if UnitSelectorOverlay.super.onInput(self, keys) then
+        return true
+    end
     if keys._MOUSE_L then
         self.refresh_search = true
     end
-    return UnitSelectorOverlay.super.onInput(self, keys)
 end
 
--- ----------------------
+-- -----------------------
 -- WorkerAssignmentOverlay
 --
 
@@ -95,4 +103,178 @@ WorkerAssignmentOverlay.ATTRS{
     handled_screens={WORKER_ASSIGNMENT='selected'},
 }
 
+-- -----------------------
+-- BurrowAssignmentOverlay
+--
+
+local DEFAULT_OVERLAY_WIDTH = 58
+
+BurrowAssignmentOverlay = defclass(BurrowAssignmentOverlay, UnitSelectorOverlay)
+BurrowAssignmentOverlay.ATTRS{
+    viewscreens='dwarfmode/UnitSelector',
+    frame={w=DEFAULT_OVERLAY_WIDTH, h=5},
+    handled_screens={BURROW_ASSIGNMENT='selected'},
+}
+
+local function get_screen_width()
+    local sw = dfhack.screen.getWindowSize()
+    return sw
+end
+
+local function toggle_all()
+    if #unit_selector.unid == 0 then return end
+    local burrow = df.burrow.find(unit_selector.burrow_id)
+    if not burrow then return end
+    local target_state = unit_selector.selected[0] == 0
+    local target_val = target_state and 1 or 0
+    for i,unit_id in ipairs(unit_selector.unid) do
+        local unit = df.unit.find(unit_id)
+        if unit then
+            dfhack.burrows.setAssignedUnit(burrow, unit, target_state)
+            unit_selector.selected[i] = target_val
+        end
+    end
+end
+
+function BurrowAssignmentOverlay:init()
+    self:addviews{
+        widgets.Panel{
+            view_id='top_mask',
+            frame={l=WIDGET_WIDTH, r=0, t=0, h=1},
+            frame_background=gui.CLEAR_PEN,
+            visible=function() return get_screen_width() >= 144 end,
+        },
+        widgets.Panel{
+            view_id='wide_mask',
+            frame={r=0, t=1, h=2, w=DEFAULT_OVERLAY_WIDTH},
+            frame_background=gui.CLEAR_PEN,
+            visible=function() return get_screen_width() >= 144 end,
+        },
+        widgets.Panel{
+            view_id='narrow_mask',
+            frame={l=0, t=1, h=2, w=24},
+            frame_background=gui.CLEAR_PEN,
+            visible=function() return get_screen_width() < 144 end,
+        },
+        widgets.BannerPanel{
+            view_id='subset_panel',
+            frame={l=0, t=1, w=WIDGET_WIDTH, h=1},
+            subviews={
+                widgets.CycleHotkeyLabel{
+                    view_id='subset',
+                    frame={l=1, t=0, r=1},
+                    key='CUSTOM_SHIFT_F',
+                    label='Show:',
+                    options={
+                        {label='All', value='all', pen=COLOR_GREEN},
+                        {label='Military', value='military', pen=COLOR_YELLOW},
+                        {label='Civilians', value='civilian', pen=COLOR_CYAN},
+                        {label='Burrowed', value='burrow', pen=COLOR_MAGENTA},
+                    },
+                    on_change=function(value)
+                        local squad = self.subviews.squad
+                        local burrow = self.subviews.burrow
+                        squad.visible = false
+                        burrow.visible = false
+                        if value == 'military' then
+                            squad.options = info.get_squad_options()
+                            squad:setOption('all')
+                            squad.visible = true
+                        elseif value == 'burrow' then
+                            burrow.options = info.get_burrow_options()
+                            burrow:setOption('all')
+                            burrow.visible = true
+                        end
+                        self:do_search(self.subviews.search.text, true)
+                    end,
+                },
+            },
+        },
+        widgets.BannerPanel{
+            view_id='subfilter_panel',
+            frame={l=0, t=2, w=WIDGET_WIDTH, h=1},
+            visible=function()
+                local subset = self.subviews.subset:getOptionValue()
+                return subset == 'military' or subset == 'burrow'
+            end,
+            subviews={
+                widgets.CycleHotkeyLabel{
+                    view_id='squad',
+                    frame={l=1, t=0, r=1},
+                    key='CUSTOM_SHIFT_S',
+                    label='Squad:',
+                    options={
+                        {label='Any', value='all', pen=COLOR_GREEN},
+                    },
+                    visible=false,
+                    on_change=function() self:do_search(self.subviews.search.text, true) end,
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='burrow',
+                    frame={l=1, t=0, r=1},
+                    key='CUSTOM_SHIFT_B',
+                    label='Burrow:',
+                    options={
+                        {label='Any', value='all', pen=COLOR_GREEN},
+                    },
+                    visible=false,
+                    on_change=function() self:do_search(self.subviews.search.text, true) end,
+                },
+            },
+        },
+        widgets.BannerPanel{
+            frame={r=0, t=4, w=25, h=1},
+            subviews={
+                widgets.HotkeyLabel{
+                    frame={l=1, t=0, r=1},
+                    label='Select all/none',
+                    key='CUSTOM_CTRL_A',
+                    on_activate=toggle_all,
+                },
+            },
+        },
+    }
+end
+
+function BurrowAssignmentOverlay:register_handlers()
+    for name,flags_vec in pairs(self.handled_screens) do
+        self:register_handler(name, unit_selector.unid,
+            curry(sortoverlay.flags_vector_search, {
+                get_search_key_fn=get_unit_id_search_key,
+                matches_filters_fn=self:callback('matches_filters'),
+            },
+            unit_selector[flags_vec]))
+    end
+end
+
+function BurrowAssignmentOverlay:matches_filters(unit_id)
+    local unit = df.unit.find(unit_id)
+    if not unit then return false end
+    return info.matches_squad_burrow_filters(unit, self.subviews.subset:getOptionValue(),
+        self.subviews.squad:getOptionValue(), self.subviews.burrow:getOptionValue())
+end
+
+local function clicked_on_mask(self, keys)
+    if not keys._MOUSE_L then return false end
+    for _,mask in ipairs{'top_mask', 'wide_mask', 'narrow_mask'} do
+        if utils.getval(self.subviews[mask].visible) then
+            if self.subviews[mask]:getMousePos() then
+                return true
+            end
+        end
+    end
+    return false
+end
+
+function BurrowAssignmentOverlay:onInput(keys)
+    return BurrowAssignmentOverlay.super.onInput(self, keys) or
+        clicked_on_mask(self, keys)
+end
+
+function BurrowAssignmentOverlay:onRenderFrame(dc, rect)
+    local sw = get_screen_width()
+    self.frame.w = math.min(DEFAULT_OVERLAY_WIDTH, sw - 94)
+    BurrowAssignmentOverlay.super.onRenderFrame(self, dc, rect)
+end
+
 return _ENV