Merge pull request #3678 from myk002/myk_sort

[sort] search and sort for squad assignment screen

docs/changelog.txt

@@ -36,8 +36,6 @@ Template for new versions:
 
 ## API
 
-## Internals
-
 ## Lua
 
 ## Removed
@@ -56,10 +54,11 @@ Template for new versions:
 ## New Tools
 
 ## New Features
-- `caravan`: advanced unit assignment screens for cages, restraints, and pits/ponds
+- `sort`: search and sort for squad assignment screen
+- `zone`: advanced unit assignment screens for cages, restraints, and pits/ponds
 
 ## Fixes
-- Core: properly reload scripts in mods when a world is unloaded and immediately loaded again
+- Core: reload scripts in mods when a world is unloaded and immediately loaded again
 - Core: fix text getting added to DFHack text entry widgets when Alt- or Ctrl- keys are hit
 - `orders`: prevent import/export overlay from appearing on the create workorder screen
 
@@ -73,14 +72,12 @@ Template for new versions:
 - ``Items::getValue()``: remove ``caravan_buying`` parameter since the identity of the selling party doesn't actually affect the item value
 - `RemoteFortressReader`: add a ``force_reload`` option to the GetBlockList RPC API to return blocks regardless of whether they have changed since the last request
 
-## Internals
-
 ## Lua
 - ``new()``: improved error handling so that certain errors that were previously uncatchable (creating objects with members with unknown vtables) are now catchable with ``pcall()``
 - ``dfhack.items.getValue()``: remove ``caravan_buying`` param as per C++ API change
 - ``widgets.BannerPanel``: panel with distinctive border for marking DFHack UI elements on otherwise vanilla screens
 - ``widgets.Panel``: new functions to override instead of setting corresponding properties (useful when subclassing instead of just setting attributes): ``onDragBegin``, ``onDragEnd``, ``onResizeBegin``, ``onResizeEnd``
-- ``dfhack.screen.readTile()``: now populates extended tile property fields (like top_of_text) in the returned Pen object
+- ``dfhack.screen.readTile()``: now populates extended tile property fields (like ``top_of_text``) in the returned ``Pen`` object
 
 ## Removed
 

docs/plugins/sort.rst

@@ -2,59 +2,37 @@ sort
 ====
 
 .. dfhack-tool::
-    :summary: Sort lists shown in the DF interface.
-    :tags: unavailable fort productivity interface
+    :summary: Search and sort lists shown in the DF interface.
+    :tags: fort productivity interface
     :no-command:
 
-.. dfhack-command:: sort-items
-   :summary: Sort the visible item list.
+The ``sort`` tool provides overlays that sorting and searching options for
+lists displayed in the DF interface.
 
-.. dfhack-command:: sort-units
-   :summary: Sort the visible unit list.
+Searching and sorting functionality is provided by `overlay` widgets, and widgets for individual lists can be moved via `gui/overlay` or turned on or off via `gui/control-panel`.
 
-Usage
------
+Squad assignment overlay
+------------------------
 
-::
+The squad assignment screen can be sorted by name, by migrant wave, or by various military-related skills.
 
-    sort-items <property> [<property> ...]
-    sort-units <property> [<property> ...]
+If sorted by "any melee", then the citizen is sorted according to the highest
+skill they have in axes, short swords, maces, warhammers, spears, or general
+fighting.
 
-Both commands sort the visible list using the given sequence of comparisons.
-Each property can be prefixed with a ``<`` or ``>`` character to indicate
-whether elements that don't have the given property defined go first or last
-(respectively) in the sorted list.
+If sorted by "any ranged", then the citizen is sorted according to the highest
+skill they have in crossbows or general ranged combat.
 
-Examples
---------
+If sorted by "leadership", then the citizen is sorted according to the highest
+skill they have in leader, teacher, or military tactics.
 
-``sort-items material type quality``
-    Sort a list of items by material, then by type, then by quality
-``sort-units profession name``
-    Sort a list of units by profession, then by name
+You can search for a dwarf by name by typing in the Search field. You can also
+type in the name of any job skill (military-related or not) and dwarves with
+any experience in that skill will be shown. For example, to only see citizens
+with military tactics skill, type in "tactics".
 
-Properties
-----------
+You can see all the job skill names that you can search for by running::
 
-Items can be sorted by the following properties:
+    :lua @df.job_skill
 
-- ``type``
-- ``description``
-- ``base_quality``
-- ``quality``
-- ``improvement``
-- ``wear``
-- ``material``
-
-Units can be sorted by the following properties:
-
-- ``name``
-- ``age``
-- ``arrival``
-- ``noble``
-- ``profession``
-- ``profession_class``
-- ``race``
-- ``squad``
-- ``squad_position``
-- ``happiness``
+in `gui/launcher`.

plugins/CMakeLists.txt

@@ -156,7 +156,7 @@ if(BUILD_SUPPORTED)
     dfhack_plugin(seedwatch seedwatch.cpp LINK_LIBRARIES lua)
     dfhack_plugin(showmood showmood.cpp)
     #dfhack_plugin(siege-engine siege-engine.cpp LINK_LIBRARIES lua)
-    #dfhack_plugin(sort sort.cpp LINK_LIBRARIES lua)
+    dfhack_plugin(sort sort.cpp LINK_LIBRARIES lua)
     #dfhack_plugin(steam-engine steam-engine.cpp)
     #add_subdirectory(spectate)
     #dfhack_plugin(stockflow stockflow.cpp LINK_LIBRARIES lua)

plugins/lua/sort.lua

@@ -1,5 +1,506 @@
 local _ENV = mkmodule('plugins.sort')
 
+local gui = require('gui')
+local overlay = require('plugins.overlay')
+local utils = require('utils')
+local widgets = require('gui.widgets')
+
+local CH_UP = string.char(30)
+local CH_DN = string.char(31)
+
+local MELEE_WEAPON_SKILLS = {
+    df.job_skill.AXE,
+    df.job_skill.SWORD,
+    df.job_skill.MACE,
+    df.job_skill.HAMMER,
+    df.job_skill.SPEAR,
+    df.job_skill.MELEE_COMBAT, --Fighter
+}
+
+local RANGED_WEAPON_SKILLS = {
+    df.job_skill.CROSSBOW,
+    df.job_skill.RANGED_COMBAT,
+}
+
+local LEADERSHIP_SKILLS = {
+    df.job_skill.MILITARY_TACTICS,
+    df.job_skill.LEADERSHIP,
+    df.job_skill.TEACHING,
+}
+
+local function sort_noop(a, b)
+    -- this function is used as a marker and never actually gets called
+    error('sort_noop should not be called')
+end
+
+local function sort_by_name_desc(unit_id_1, unit_id_2)
+    if unit_id_1 == unit_id_2 then return 0 end
+    local unit1 = df.unit.find(unit_id_1)
+    local unit2 = df.unit.find(unit_id_2)
+    if not unit1 then return -1 end
+    if not unit2 then return 1 end
+    local name1 = dfhack.TranslateName(dfhack.units.getVisibleName(unit1))
+    local name2 = dfhack.TranslateName(dfhack.units.getVisibleName(unit2))
+    return utils.compare_name(name1, name2)
+end
+
+local function sort_by_name_asc(unit_id_1, unit_id_2)
+    if unit_id_1 == unit_id_2 then return 0 end
+    local unit1 = df.unit.find(unit_id_1)
+    local unit2 = df.unit.find(unit_id_2)
+    if not unit1 then return -1 end
+    if not unit2 then return 1 end
+    local name1 = dfhack.TranslateName(dfhack.units.getVisibleName(unit1))
+    local name2 = dfhack.TranslateName(dfhack.units.getVisibleName(unit2))
+    return utils.compare_name(name2, name1)
+end
+
+local function sort_by_migrant_wave_desc(unit_id_1, unit_id_2)
+    if unit_id_1 == unit_id_2 then return 0 end
+    local unit1 = df.unit.find(unit_id_1)
+    local unit2 = df.unit.find(unit_id_2)
+    if not unit1 then return -1 end
+    if not unit2 then return 1 end
+    return utils.compare(unit2.id, unit1.id)
+end
+
+local function sort_by_migrant_wave_asc(unit_id_1, unit_id_2)
+    if unit_id_1 == unit_id_2 then return 0 end
+    local unit1 = df.unit.find(unit_id_1)
+    local unit2 = df.unit.find(unit_id_2)
+    if not unit1 then return -1 end
+    if not unit2 then return 1 end
+    return utils.compare(unit1.id, unit2.id)
+end
+
+local function get_skill(unit_id, skill, unit)
+    unit = unit or df.unit.find(unit_id)
+    return unit and
+        unit.status.current_soul and
+        utils.binsearch(unit.status.current_soul.skills, skill, 'id')
+end
+
+local function get_max_skill(unit_id, list)
+    local unit = df.unit.find(unit_id)
+    if not unit or not unit.status.current_soul then return end
+    local max
+    for _,skill in ipairs(list) do
+        local s = get_skill(unit_id, skill, unit)
+        if s then
+            if not max then
+                max = s
+            else
+                if max.rating == s.rating and max.experience < s.experience then
+                    max = s
+                elseif max.rating < s.rating then
+                    max = s
+                end
+            end
+        end
+    end
+    return max
+end
+
+local function make_sort_by_skill_list_desc(list)
+    return function(unit_id_1, unit_id_2)
+        if unit_id_1 == unit_id_2 then return 0 end
+        if unit_id_1 == -1 then return -1 end
+        if unit_id_2 == -1 then return 1 end
+        local s1 = get_max_skill(unit_id_1, list)
+        local s2 = get_max_skill(unit_id_2, list)
+        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
+        if not s2 then return -1 end
+        if not s1 then return 1 end
+        if s1.rating ~= s2.rating then
+            return utils.compare(s2.rating, s1.rating)
+        end
+        if s1.experience ~= s2.experience then
+            return utils.compare(s2.experience, s1.experience)
+        end
+        return sort_by_name_desc(unit_id_1, unit_id_2)
+    end
+end
+
+local function make_sort_by_skill_list_asc(list)
+    return function(unit_id_1, unit_id_2)
+        if unit_id_1 == unit_id_2 then return 0 end
+        if unit_id_1 == -1 then return -1 end
+        if unit_id_2 == -1 then return 1 end
+        local s1 = get_max_skill(unit_id_1, list)
+        local s2 = get_max_skill(unit_id_2, list)
+        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
+        if not s2 then return 1 end
+        if not s1 then return -1 end
+        if s1.rating ~= s2.rating then
+            return utils.compare(s1.rating, s2.rating)
+        end
+        if s1.experience ~= s2.experience then
+            return utils.compare(s1.experience, s2.experience)
+        end
+        return sort_by_name_desc(unit_id_1, unit_id_2)
+    end
+end
+
+local function make_sort_by_skill_desc(sort_skill)
+    return function(unit_id_1, unit_id_2)
+        if unit_id_1 == unit_id_2 then return 0 end
+        if unit_id_1 == -1 then return -1 end
+        if unit_id_2 == -1 then return 1 end
+        local s1 = get_skill(unit_id_1, sort_skill)
+        local s2 = get_skill(unit_id_2, sort_skill)
+        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
+        if not s2 then return -1 end
+        if not s1 then return 1 end
+        if s1.rating ~= s2.rating then
+            return utils.compare(s2.rating, s1.rating)
+        end
+        if s1.experience ~= s2.experience then
+            return utils.compare(s2.experience, s1.experience)
+        end
+        return sort_by_name_desc(unit_id_1, unit_id_2)
+    end
+end
+
+local function make_sort_by_skill_asc(sort_skill)
+    return function(unit_id_1, unit_id_2)
+        if unit_id_1 == unit_id_2 then return 0 end
+        if unit_id_1 == -1 then return -1 end
+        if unit_id_2 == -1 then return 1 end
+        local s1 = get_skill(unit_id_1, sort_skill)
+        local s2 = get_skill(unit_id_2, sort_skill)
+        if s1 == s2 then return sort_by_name_desc(unit_id_1, unit_id_2) end
+        if not s2 then return 1 end
+        if not s1 then return -1 end
+        if s1.rating ~= s2.rating then
+            return utils.compare(s1.rating, s2.rating)
+        end
+        if s1.experience ~= s2.experience then
+            return utils.compare(s1.experience, s2.experience)
+        end
+        return sort_by_name_desc(unit_id_1, unit_id_2)
+    end
+end
+
+local SORT_FNS = {
+    sort_by_any_melee_desc=make_sort_by_skill_list_desc(MELEE_WEAPON_SKILLS),
+    sort_by_any_melee_asc=make_sort_by_skill_list_asc(MELEE_WEAPON_SKILLS),
+    sort_by_any_ranged_desc=make_sort_by_skill_list_desc(RANGED_WEAPON_SKILLS),
+    sort_by_any_ranged_asc=make_sort_by_skill_list_asc(RANGED_WEAPON_SKILLS),
+    sort_by_leadership_desc=make_sort_by_skill_list_desc(LEADERSHIP_SKILLS),
+    sort_by_leadership_asc=make_sort_by_skill_list_asc(LEADERSHIP_SKILLS),
+    sort_by_axe_desc=make_sort_by_skill_desc(df.job_skill.AXE),
+    sort_by_axe_asc=make_sort_by_skill_asc(df.job_skill.AXE),
+    sort_by_sword_desc=make_sort_by_skill_desc(df.job_skill.SWORD),
+    sort_by_sword_asc=make_sort_by_skill_asc(df.job_skill.SWORD),
+    sort_by_mace_desc=make_sort_by_skill_desc(df.job_skill.MACE),
+    sort_by_mace_asc=make_sort_by_skill_asc(df.job_skill.MACE),
+    sort_by_hammer_desc=make_sort_by_skill_desc(df.job_skill.HAMMER),
+    sort_by_hammer_asc=make_sort_by_skill_asc(df.job_skill.HAMMER),
+    sort_by_spear_desc=make_sort_by_skill_desc(df.job_skill.SPEAR),
+    sort_by_spear_asc=make_sort_by_skill_asc(df.job_skill.SPEAR),
+    sort_by_crossbow_desc=make_sort_by_skill_desc(df.job_skill.CROSSBOW),
+    sort_by_crossbow_asc=make_sort_by_skill_asc(df.job_skill.CROSSBOW),
+}
+
+-- ----------------------
+-- SquadAssignmentOverlay
+--
+
+SquadAssignmentOverlay = defclass(SquadAssignmentOverlay, overlay.OverlayWidget)
+SquadAssignmentOverlay.ATTRS{
+    default_pos={x=-33, y=40},
+    default_enabled=true,
+    viewscreens='dwarfmode/UnitSelector/SQUAD_FILL_POSITION',
+    frame={w=63, h=7},
+    frame_style=gui.FRAME_PANEL,
+    frame_background=gui.CLEAR_PEN,
+}
+
+function SquadAssignmentOverlay:init()
+    self.dirty = true
+
+    self:addviews{
+        widgets.CycleHotkeyLabel{
+            view_id='sort',
+            frame={l=0, t=0, w=29},
+            label='Sort by:',
+            key='CUSTOM_SHIFT_S',
+            options={
+                {label='any melee skill'..CH_DN, value=SORT_FNS.sort_by_any_melee_desc, pen=COLOR_GREEN},
+                {label='any melee skill'..CH_UP, value=SORT_FNS.sort_by_any_melee_asc, pen=COLOR_YELLOW},
+                {label='any ranged skill'..CH_DN, value=SORT_FNS.sort_by_any_ranged_desc, pen=COLOR_GREEN},
+                {label='any ranged skill'..CH_UP, value=SORT_FNS.sort_by_any_ranged_asc, pen=COLOR_YELLOW},
+                {label='any leader skill'..CH_DN, value=SORT_FNS.sort_by_leadership_desc, pen=COLOR_GREEN},
+                {label='any leader skill'..CH_UP, value=SORT_FNS.sort_by_leadership_asc, pen=COLOR_YELLOW},
+                {label='name'..CH_DN, value=sort_by_name_desc, pen=COLOR_GREEN},
+                {label='name'..CH_UP, value=sort_by_name_asc, pen=COLOR_YELLOW},
+                {label='migrant wave'..CH_DN, value=sort_by_migrant_wave_desc, pen=COLOR_GREEN},
+                {label='migrant wave'..CH_UP, value=sort_by_migrant_wave_asc, pen=COLOR_YELLOW},
+                {label='axe skill'..CH_DN, value=SORT_FNS.sort_by_axe_desc, pen=COLOR_GREEN},
+                {label='axe skill'..CH_UP, value=SORT_FNS.sort_by_axe_asc, pen=COLOR_YELLOW},
+                {label='sword skill'..CH_DN, value=SORT_FNS.sort_by_sword_desc, pen=COLOR_GREEN},
+                {label='sword skill'..CH_UP, value=SORT_FNS.sort_by_sword_asc, pen=COLOR_YELLOW},
+                {label='mace skill'..CH_DN, value=SORT_FNS.sort_by_mace_desc, pen=COLOR_GREEN},
+                {label='mace skill'..CH_UP, value=SORT_FNS.sort_by_mace_asc, pen=COLOR_YELLOW},
+                {label='hammer skill'..CH_DN, value=SORT_FNS.sort_by_hammer_desc, pen=COLOR_GREEN},
+                {label='hammer skill'..CH_UP, value=SORT_FNS.sort_by_hammer_asc, pen=COLOR_YELLOW},
+                {label='spear skill'..CH_DN, value=SORT_FNS.sort_by_spear_desc, pen=COLOR_GREEN},
+                {label='spear skill'..CH_UP, value=SORT_FNS.sort_by_spear_asc, pen=COLOR_YELLOW},
+                {label='crossbow skill'..CH_DN, value=SORT_FNS.sort_by_crossbow_desc, pen=COLOR_GREEN},
+                {label='crossbow skill'..CH_UP, value=SORT_FNS.sort_by_crossbow_asc, pen=COLOR_YELLOW},
+            },
+            initial_option=SORT_FNS.sort_by_any_melee_desc,
+            on_change=self:callback('refresh_list', 'sort'),
+        },
+        widgets.EditField{
+            view_id='search',
+            frame={l=32, t=0},
+            label_text='Search: ',
+            on_char=function(ch) return ch:match('[%l _-]') end,
+            on_change=function() self:refresh_list() end,
+        },
+        widgets.Panel{
+            frame={t=2, l=0, r=0, b=0},
+            subviews={
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_any_melee',
+                    frame={t=0, l=0, w=10},
+                    options={
+                        {label='any melee', value=sort_noop},
+                        {label='any melee'..CH_DN, value=SORT_FNS.sort_by_any_melee_desc, pen=COLOR_GREEN},
+                        {label='any melee'..CH_UP, value=SORT_FNS.sort_by_any_melee_asc, pen=COLOR_YELLOW},
+                    },
+                    initial_option=SORT_FNS.sort_by_any_melee_desc,
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_any_melee'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_any_ranged',
+                    frame={t=0, l=13, w=11},
+                    options={
+                        {label='any ranged', value=sort_noop},
+                        {label='any ranged'..CH_DN, value=SORT_FNS.sort_by_any_ranged_desc, pen=COLOR_GREEN},
+                        {label='any ranged'..CH_UP, value=SORT_FNS.sort_by_any_ranged_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_any_ranged'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_leadership',
+                    frame={t=0, l=27, w=11},
+                    options={
+                        {label='leadership', value=sort_noop},
+                        {label='leadership'..CH_DN, value=SORT_FNS.sort_by_leadership_desc, pen=COLOR_GREEN},
+                        {label='leadership'..CH_UP, value=SORT_FNS.sort_by_leadership_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_leadership'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_name',
+                    frame={t=0, l=41, w=5},
+                    options={
+                        {label='name', value=sort_noop},
+                        {label='name'..CH_DN, value=sort_by_name_desc, pen=COLOR_GREEN},
+                        {label='name'..CH_UP, value=sort_by_name_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_name'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_migrant_wave',
+                    frame={t=0, l=48, w=13},
+                    options={
+                        {label='migrant wave', value=sort_noop},
+                        {label='migrant wave'..CH_DN, value=sort_by_migrant_wave_desc, pen=COLOR_GREEN},
+                        {label='migrant wave'..CH_UP, value=sort_by_migrant_wave_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_migrant_wave'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_axe',
+                    frame={t=2, l=2, w=4},
+                    options={
+                        {label='axe', value=sort_noop},
+                        {label='axe'..CH_DN, value=SORT_FNS.sort_by_axe_desc, pen=COLOR_GREEN},
+                        {label='axe'..CH_UP, value=SORT_FNS.sort_by_axe_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_axe'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_sword',
+                    frame={t=2, l=9, w=6},
+                    options={
+                        {label='sword', value=sort_noop},
+                        {label='sword'..CH_DN, value=SORT_FNS.sort_by_sword_desc, pen=COLOR_GREEN},
+                        {label='sword'..CH_UP, value=SORT_FNS.sort_by_sword_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_sword'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_mace',
+                    frame={t=2, l=18, w=5},
+                    options={
+                        {label='mace', value=sort_noop},
+                        {label='mace'..CH_DN, value=SORT_FNS.sort_by_mace_desc, pen=COLOR_GREEN},
+                        {label='mace'..CH_UP, value=SORT_FNS.sort_by_mace_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_mace'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_hammer',
+                    frame={t=2, l=25, w=7},
+                    options={
+                        {label='hammer', value=sort_noop},
+                        {label='hammer'..CH_DN, value=SORT_FNS.sort_by_hammer_desc, pen=COLOR_GREEN},
+                        {label='hammer'..CH_UP, value=SORT_FNS.sort_by_hammer_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_hammer'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_spear',
+                    frame={t=2, l=36, w=6},
+                    options={
+                        {label='spear', value=sort_noop},
+                        {label='spear'..CH_DN, value=SORT_FNS.sort_by_spear_desc, pen=COLOR_GREEN},
+                        {label='spear'..CH_UP, value=SORT_FNS.sort_by_spear_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_spear'),
+                },
+                widgets.CycleHotkeyLabel{
+                    view_id='sort_crossbow',
+                    frame={t=2, l=45, w=9},
+                    options={
+                        {label='crossbow', value=sort_noop},
+                        {label='crossbow'..CH_DN, value=SORT_FNS.sort_by_crossbow_desc, pen=COLOR_GREEN},
+                        {label='crossbow'..CH_UP, value=SORT_FNS.sort_by_crossbow_asc, pen=COLOR_YELLOW},
+                    },
+                    option_gap=0,
+                    on_change=self:callback('refresh_list', 'sort_crossbow'),
+                },
+            }
+        },
+    }
+end
+
+local function normalize_search_key(search_key)
+    local out = ''
+    for c in dfhack.toSearchNormalized(search_key):gmatch("[%w%s]") do
+        out = out .. c:lower()
+    end
+    return out
+end
+
+local function filter_matches(unit_id, search)
+    if unit_id == -1 then return true end
+    local unit = df.unit.find(unit_id)
+    if not unit then return true end
+    local search_key = dfhack.TranslateName(dfhack.units.getVisibleName(unit))
+    if unit.status.current_soul then
+        for _,skill in ipairs(unit.status.current_soul.skills) do
+            search_key = (search_key or '') .. ' ' .. (df.job_skill[skill.id] or '')
+        end
+    end
+    return normalize_search_key(search_key):find(dfhack.toSearchNormalized(search))
+end
+
+local unit_selector = df.global.game.main_interface.unit_selector
+
+-- this function uses the unused itemid and selected vectors to keep state,
+-- taking advantage of the fact that they are reset by DF when the list of units changes
+local function filter_vector(search, prev_search)
+    local unid_is_filtered = #unit_selector.selected >= 0 and unit_selector.selected[0] ~= 0
+    if #search == 0 or #unit_selector.selected == 0 then
+        if not unid_is_filtered then
+            -- we haven't modified the unid vector; nothing to do here
+            return
+        end
+        -- restore the unid vector
+        unit_selector.unid:assign(unit_selector.itemid)
+        -- clear our "we meddled" flag
+        unit_selector.selected[0] = 0
+        return
+    end
+    if unid_is_filtered and search == prev_search then
+        -- prev filter still stands
+        return
+    end
+    if unid_is_filtered then
+        -- restore the unid vector
+        unit_selector.unid:assign(unit_selector.itemid)
+    else
+        -- save the unid vector and set our meddle flag
+        unit_selector.itemid:assign(unit_selector.unid)
+        unit_selector.selected[0] = 1
+    end
+    -- do the actual filtering
+    for idx=#unit_selector.unid-1,0,-1 do
+        if not filter_matches(unit_selector.unid[idx], search) then
+            unit_selector.unid:erase(idx)
+        end
+    end
+end
+
+local SORT_WIDGET_NAMES = {
+    'sort',
+    'sort_any_melee',
+    'sort_any_ranged',
+    'sort_leadership',
+    'sort_name',
+    'sort_migrant_wave',
+    'sort_axe',
+    'sort_sword',
+    'sort_mace',
+    'sort_hammer',
+    'sort_spear',
+    'sort_crossbow',
+}
+
+function SquadAssignmentOverlay:refresh_list(sort_widget, sort_fn)
+    sort_widget = sort_widget or 'sort'
+    sort_fn = sort_fn or self.subviews.sort:getOptionValue()
+    if sort_fn == sort_noop then
+        self.subviews[sort_widget]:cycle()
+        return
+    end
+    for _,widget_name in ipairs(SORT_WIDGET_NAMES) do
+        self.subviews[widget_name]:setOption(sort_fn)
+    end
+    local search = self.subviews.search.text
+    filter_vector(search, self.prev_search)
+    self.prev_search = search
+    utils.sort_vector(unit_selector.unid, nil, sort_fn)
+end
+
+function SquadAssignmentOverlay:onInput(keys)
+    if keys._MOUSE_R_DOWN or
+        keys._MOUSE_L_DOWN and not self:getMouseFramePos()
+    then
+        -- if any click is made outside of our window, we may need to refresh our list
+        self.dirty = true
+    end
+    return SquadAssignmentOverlay.super.onInput(self, keys)
+end
+
+function SquadAssignmentOverlay:onRenderFrame(dc, frame_rect)
+    SquadAssignmentOverlay.super.onRenderFrame(self, dc, frame_rect)
+    if self.dirty then
+        self:refresh_list()
+        self.dirty = false
+    end
+end
+
+OVERLAY_WIDGETS = {
+    squad_assignment=SquadAssignmentOverlay,
+}
+
+--[[
 local utils = require('utils')
 local units = require('plugins.sort.units')
 local items = require('plugins.sort.items')
@@ -51,5 +552,6 @@ function parse_ordering_spec(type,...)
 end
 
 make_sort_order = utils.make_sort_order
+]]
 
 return _ENV

plugins/overlay.cpp

@@ -1,3 +1,4 @@
+#include "df/enabler.h"
 #include "df/viewscreen_adopt_regionst.h"
 #include "df/viewscreen_choose_game_typest.h"
 #include "df/viewscreen_choose_start_sitest.h"
@@ -29,6 +30,7 @@ DFHACK_PLUGIN("overlay");
 DFHACK_PLUGIN_IS_ENABLED(is_enabled);
 
 REQUIRE_GLOBAL(world);
+REQUIRE_GLOBAL(enabler);
 
 namespace DFHack {
     DBG_DECLARE(overlay, control, DebugCategory::LINFO);
@@ -82,6 +84,8 @@ struct viewscreen_overlay : T {
                     });
         if (!input_is_handled)
             INTERPOSE_NEXT(feed)(input);
+        else
+            enabler->last_text_input[0] = '\0';
     }
     DEFINE_VMETHOD_INTERPOSE(void, render, ()) {
         INTERPOSE_NEXT(render)();

plugins/sort.cpp

@@ -1,43 +1,13 @@
-#include "Core.h"
-#include "Console.h"
-#include "Export.h"
 #include "PluginManager.h"
 
-#include "modules/Gui.h"
-#include "modules/Translation.h"
-#include "modules/Units.h"
-#include "modules/Job.h"
-
-#include "LuaTools.h"
-
-#include "DataDefs.h"
-#include "df/plotinfost.h"
-#include "df/world.h"
-#include "df/viewscreen_joblistst.h"
-#include "df/viewscreen_unitlistst.h"
-#include "df/viewscreen_layer_militaryst.h"
-#include "df/viewscreen_layer_noblelistst.h"
-#include "df/viewscreen_layer_overall_healthst.h"
-#include "df/viewscreen_layer_assigntradest.h"
-#include "df/viewscreen_tradegoodsst.h"
-#include "df/viewscreen_dwarfmodest.h"
-#include "df/viewscreen_petst.h"
-#include "df/viewscreen_storesst.h"
-#include "df/viewscreen_workshop_profilest.h"
-#include "df/layer_object_listst.h"
-#include "df/assign_trade_status.h"
-
-#include "MiscUtils.h"
-
-#include <stdlib.h>
-
 using std::vector;
 using std::string;
-using std::endl;
+
 using namespace DFHack;
-using namespace df::enums;
 
 DFHACK_PLUGIN("sort");
+
+/*
 REQUIRE_GLOBAL(plotinfo);
 REQUIRE_GLOBAL(world);
 REQUIRE_GLOBAL(ui_building_in_assign);
@@ -52,16 +22,18 @@ static bool item_list_hotkey(df::viewscreen *top);
 
 static command_result sort_units(color_ostream &out, vector <string> & parameters);
 static command_result sort_items(color_ostream &out, vector <string> & parameters);
+*/
 
-DFhackCExport command_result plugin_init (color_ostream &out, std::vector <PluginCommand> &commands)
+DFhackCExport command_result plugin_init (color_ostream &out, vector <PluginCommand> &commands)
 {
-    commands.push_back(PluginCommand(
-        "sort-units", "Sort the visible unit list.", sort_units, unit_list_hotkey));
-    commands.push_back(PluginCommand(
-        "sort-items", "Sort the visible item list.", sort_items, item_list_hotkey));
+    // commands.push_back(PluginCommand(
+    //     "sort-units", "Sort the visible unit list.", sort_units, unit_list_hotkey));
+    // commands.push_back(PluginCommand(
+    //     "sort-items", "Sort the visible item list.", sort_items, item_list_hotkey));
     return CR_OK;
 }
 
+/*
 DFhackCExport command_result plugin_shutdown ( color_ostream &out )
 {
     return CR_OK;
@@ -232,10 +204,7 @@ typedef void (*SortHandler)(color_ostream *pout, lua_State *L, int top,
 
 static std::map<std::string, SortHandler> unit_sorters;
 
-/*
- * Sort units in the 'u'nit list screen.
- */
-
+// Sort units in the 'u'nit list screen.
 DEFINE_SORT_HANDLER(unit_sorters, unitlist, "", units)
 {
     PARSE_SPEC("units", parameters);
@@ -250,10 +219,7 @@ DEFINE_SORT_HANDLER(unit_sorters, unitlist, "", units)
     }
 }
 
-/*
- * Sort units in the 'j'ob list screen.
- */
-
+//Sort units in the 'j'ob list screen.
 DEFINE_SORT_HANDLER(unit_sorters, joblist, "", jobs)
 {
     PARSE_SPEC("units", parameters);
@@ -275,10 +241,7 @@ DEFINE_SORT_HANDLER(unit_sorters, joblist, "", jobs)
     }
 }
 
-/*
- * Sort candidate units in the 'p'osition page of the 'm'ilitary screen.
- */
-
+// Sort candidate units in the 'p'osition page of the 'm'ilitary screen.
 DEFINE_SORT_HANDLER(unit_sorters, layer_military, "/Positions/Candidates", military)
 {
     auto &candidates = military->positions.candidates;
@@ -293,7 +256,6 @@ DEFINE_SORT_HANDLER(unit_sorters, layer_military, "/Positions/Candidates", milit
     }
 }
 
-
 DEFINE_SORT_HANDLER(unit_sorters, layer_noblelist, "/Appoint", nobles)
 {
     auto list2 = getLayerList(nobles, 1);
@@ -312,10 +274,7 @@ DEFINE_SORT_HANDLER(unit_sorters, layer_noblelist, "/Appoint", nobles)
     }
 }
 
-/*
- * Sort animal units in the Animal page of the 'z' status screen.
- */
-
+//Sort animal units in the Animal page of the 'z' status screen.
 DEFINE_SORT_HANDLER(unit_sorters, pet, "/List", animals)
 {
     PARSE_SPEC("units", parameters);
@@ -334,10 +293,7 @@ DEFINE_SORT_HANDLER(unit_sorters, pet, "/List", animals)
     }
 }
 
-/*
- * Sort candidate trainers in the Animal page of the 'z' status screen.
- */
-
+// Sort candidate trainers in the Animal page of the 'z' status screen.
 DEFINE_SORT_HANDLER(unit_sorters, pet, "/SelectTrainer", animals)
 {
     sort_null_first(parameters);
@@ -351,10 +307,7 @@ DEFINE_SORT_HANDLER(unit_sorters, pet, "/SelectTrainer", animals)
     }
 }
 
-/*
- * Sort units in the Health page of the 'z' status screen.
- */
-
+// Sort units in the Health page of the 'z' status screen.
 DEFINE_SORT_HANDLER(unit_sorters, layer_overall_health, "/Units", health)
 {
     auto list1 = getLayerList(health, 0);
@@ -371,10 +324,7 @@ DEFINE_SORT_HANDLER(unit_sorters, layer_overall_health, "/Units", health)
     }
 }
 
-/*
- * Sort burrow member candidate units in the 'w' sidebar mode.
- */
-
+// Sort burrow member candidate units in the 'w' sidebar mode.
 DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/Burrows/AddUnits", screen)
 {
     PARSE_SPEC("units", parameters);
@@ -387,10 +337,7 @@ DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/Burrows/AddUnits", screen)
     }
 }
 
-/*
- * Sort building owner candidate units in the 'q' sidebar mode, or cage assignment.
- */
-
+// Sort building owner candidate units in the 'q' sidebar mode, or cage assignment.
 DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/QueryBuilding/Some/Assign", screen)
 {
     sort_null_first(parameters);
@@ -410,10 +357,7 @@ DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/QueryBuilding/Some/Assign", scree
     }
 }
 
-/*
- * Sort units in the workshop 'q'uery 'P'rofile modification screen.
- */
-
+// Sort units in the workshop 'q'uery 'P'rofile modification screen.
 DEFINE_SORT_HANDLER(unit_sorters, workshop_profile, "/Unit", profile)
 {
     PARSE_SPEC("units", parameters);
@@ -425,10 +369,7 @@ DEFINE_SORT_HANDLER(unit_sorters, workshop_profile, "/Unit", profile)
     }
 }
 
-/*
- * Sort pen assignment candidate units in 'z'->'N'.
- */
-
+// Sort pen assignment candidate units in 'z'->'N'.
 DEFINE_SORT_HANDLER(unit_sorters, dwarfmode, "/ZonesPenInfo/Assign", screen)
 {
     PARSE_SPEC("units", parameters);
@@ -562,3 +503,4 @@ static command_result sort_items(color_ostream &out, vector <string> &parameters
 
     return CR_OK;
 }
+*/