search and sort for elevate to barony page

docs/changelog.txt

@@ -54,6 +54,7 @@ Template for new versions:
 ## New Tools
 
 ## New Features
+- `sort`: search and sort for the "choose unit to elevate to the barony" screen. units are sorted by the number of item preferences they have and the units are annotated with the items that they have preferences for
 
 ## Fixes
 

library/modules/Gui.cpp

@@ -524,6 +524,8 @@ DEFINE_GET_FOCUS_STRING_HANDLER(dwarfmode)
         newFocusString += "/Diplomacy";
         if (game->main_interface.diplomacy.taking_requests)
             newFocusString += "/Requests";
+        else if (game->main_interface.diplomacy.selecting_land_holder_position)
+            newFocusString += "/ElevateLandHolder";
         else
             newFocusString += "/Default";
         focusStrings.push_back(newFocusString);

plugins/lua/sort.lua

@@ -1299,6 +1299,8 @@ OVERLAY_WIDGETS = {
     slab=require('plugins.sort.slab').SlabOverlay,
     world=require('plugins.sort.world').WorldOverlay,
     places=require('plugins.sort.places').PlacesOverlay,
+    elevate_barony=require('plugins.sort.diplomacy').DiplomacyOverlay,
+    elevate_barony_preferences=require('plugins.sort.diplomacy').PreferenceOverlay,
 }
 
 dfhack.onStateChange[GLOBAL_KEY] = function(sc)

plugins/lua/sort/diplomacy.lua

@@ -0,0 +1,165 @@
+local _ENV = mkmodule('plugins.sort.diplomacy')
+
+local overlay = require('plugins.overlay')
+local sortoverlay = require('plugins.sort.sortoverlay')
+local widgets = require('gui.widgets')
+
+local diplomacy = df.global.game.main_interface.diplomacy
+
+-- ----------------------
+-- DiplomacyOverlay
+--
+
+DiplomacyOverlay = defclass(DiplomacyOverlay, sortoverlay.SortOverlay)
+DiplomacyOverlay.ATTRS{
+    default_pos={x=25, y=7},
+    viewscreens='dwarfmode/Diplomacy',
+    frame={w=57, h=1},
+}
+
+function DiplomacyOverlay:init()
+    self:addviews{
+        widgets.BannerPanel{
+            view_id='panel',
+            frame={l=0, t=0, r=0, h=1},
+            visible=self:callback('get_key'),
+            subviews={
+                widgets.EditField{
+                    view_id='search',
+                    frame={l=1, t=0, r=1},
+                    label_text="Search: ",
+                    key='CUSTOM_ALT_S',
+                    text='',
+                    on_change=function(text) self:do_search(text) end,
+                },
+            },
+        },
+    }
+
+    self:register_handler('ELEVATE', diplomacy.land_holder_avail_hfid,
+        curry(sortoverlay.single_vector_search,
+            {
+                get_search_key_fn=self:callback('get_search_key'),
+                get_sort_fn=self:callback('get_sort'),
+            }))
+end
+
+function DiplomacyOverlay:get_key()
+    if diplomacy.selecting_land_holder_position then
+        return 'ELEVATE'
+    end
+end
+
+local function to_item_type_str(item_type)
+    return string.lower(df.item_type[item_type]):gsub('_', ' ')
+end
+
+local function make_item_description(item_type, subtype)
+    local itemdef = dfhack.items.getSubtypeDef(item_type, subtype)
+    return itemdef and string.lower(itemdef.name) or to_item_type_str(item_type)
+end
+
+local function get_preferences(unit)
+    if not unit then return {} end
+    local preferences = {}
+    for _, pref in ipairs(unit.status.current_soul.preferences) do
+        if pref.type == df.unit_preference.T_type.LikeItem and pref.active then
+            table.insert(preferences, make_item_description(pref.item_type, pref.item_subtype))
+        end
+    end
+    return preferences
+end
+
+local function get_unit(hfid)
+    local hf = df.historical_figure.find(hfid)
+    if not hf then return end
+    return df.unit.find(hf.unit_id)
+end
+
+function DiplomacyOverlay:get_search_key(hfid)
+    local unit = get_unit(hfid)
+    if not unit then return end
+    return ('%s %s'):format(sortoverlay.get_unit_search_key(unit),
+        table.concat(get_preferences(unit), ' '))
+end
+
+
+function DiplomacyOverlay:get_sort()
+    local _, sh = dfhack.screen.getWindowSize()
+    local list_height = sh - 17
+    local num_elems = list_height // 3
+    diplomacy.scroll_position_land_holder_hf = math.max(0,
+        math.min(diplomacy.scroll_position_land_holder_hf,
+            #diplomacy.land_holder_avail_hfid - num_elems))
+    return function(a, b)
+        local unita = get_unit(a)
+        local unitb = get_unit(b)
+        return #get_preferences(unita) < #get_preferences(unitb)
+    end
+end
+
+function DiplomacyOverlay:onRenderFrame(dc, rect)
+    local sw = dfhack.screen.getWindowSize()
+    local margin = (sw - 114) // 3
+    self.frame.w = 57 + margin
+    self.subviews.panel.frame.l = margin
+
+    DiplomacyOverlay.super.onRenderFrame(self, dc, rect)
+end
+
+-- ----------------------
+-- PreferenceOverlay
+--
+
+PreferenceOverlay = defclass(PreferenceOverlay, overlay.OverlayWidget)
+PreferenceOverlay.ATTRS{
+    default_pos={x=-34, y=9},
+    viewscreens='dwarfmode/Diplomacy/ElevateLandHolder',
+    default_enabled=true,
+    frame={w=29, h=1},
+}
+
+function PreferenceOverlay:init()
+    self:addviews{
+        widgets.Label{
+            view_id='annotations',
+            frame={t=0, l=0},
+            text='',
+            text_pen=COLOR_GRAY,
+        }
+    }
+end
+
+function PreferenceOverlay:onRenderFrame(dc, rect)
+    local sw, sh = dfhack.screen.getWindowSize()
+    local margin = (sw - 114) // 3
+    local list_height = sh - 17
+    local num_elems = list_height // 3
+    local max_elem = math.min(#diplomacy.land_holder_avail_hfid-1,
+        diplomacy.scroll_position_land_holder_hf+num_elems-1)
+    self.frame.w = sw - 85
+    self.frame.h = list_height
+
+    local annotations = {}
+    for idx=diplomacy.scroll_position_land_holder_hf,max_elem do
+        table.insert(annotations, NEWLINE)
+        table.insert(annotations, NEWLINE)
+        local prefs = get_preferences(get_unit(diplomacy.land_holder_avail_hfid[idx]))
+        table.insert(annotations, {text='[', pen=COLOR_RED})
+        if #prefs == 0 then
+            table.insert(annotations, 'no item preferences')
+        else
+            local pref_str = ('%d preference%s: %s'):format(#prefs,
+                #prefs > 1 and 's' or '', table.concat(prefs, ', '))
+            table.insert(annotations, pref_str:sub(1, self.frame.w-(margin+2)))
+        end
+        table.insert(annotations, {text=']', pen=COLOR_RED})
+        table.insert(annotations, NEWLINE)
+    end
+    self.subviews.annotations.frame.l = margin
+    self.subviews.annotations:setText(annotations)
+
+    PreferenceOverlay.super.onRenderFrame(self, dc, rect)
+end
+
+return _ENV

plugins/lua/sort/sortoverlay.lua

@@ -84,7 +84,6 @@ end
 
 -- handles saving/restoring search strings when the player moves between different contexts
 function SortOverlay:onRenderBody(dc)
-    if next(self.state) then
     local key, group = self:get_key()
     if self.state.cur_group ~= group then
         self.state.cur_group = group
@@ -96,7 +95,6 @@ function SortOverlay:onRenderBody(dc)
         self.subviews.search:setText(prev_text)
         self:do_search(self.subviews.search.text, true)
     end
-    end
     self.overlay_onupdate_max_freq_seconds = 0
     SortOverlay.super.onRenderBody(self, dc)
 end
@@ -132,7 +130,7 @@ function SortOverlay:do_search(text, force_full_search)
     if not key then return end
     local prev_text = ensure_key(self.state, key).prev_text
     -- some screens reset their contents between context switches; regardless,
-    -- a switch back to the context should results in an incremental search
+    -- a switch back to the context should result in an incremental search
     local incremental = not force_full_search and prev_text and text:startswith(prev_text)
     local handler = self.handlers[key]
     handler.search_fn(handler.vec, self.state[key], text, incremental)