add support for location selector

that is, guildhalls and temples
also start refactoring code for better reuse

plugins/lua/sort.lua

@@ -2,6 +2,7 @@ local _ENV = mkmodule('plugins.sort')
 
 local info = require('plugins.sort.info')
 local gui = require('gui')
+local locationselector = require('plugins.sort.locationselector')
 local overlay = require('plugins.overlay')
 local setbelief = reqscript('modtools/set-belief')
 local textures = require('gui.textures')
@@ -1288,6 +1289,7 @@ OVERLAY_WIDGETS = {
     squad_annotation=SquadAnnotationOverlay,
     info=info.InfoOverlay,
     interrogation=info.InterrogationOverlay,
+    location_selector=locationselector.LocationSelectorOverlay,
 }
 
 dfhack.onStateChange[GLOBAL_KEY] = function(sc)

plugins/lua/sort/locationselector.lua

@@ -0,0 +1,67 @@
+local _ENV = mkmodule('plugins.sort.locationselector')
+
+local sortoverlay = require('plugins.sort.sortoverlay')
+local widgets = require('gui.widgets')
+
+local location_selector = df.global.game.main_interface.location_selector
+
+-- ----------------------
+-- LocationSelectorOverlay
+--
+
+LocationSelectorOverlay = defclass(LocationSelectorOverlay, sortoverlay.SortOverlay)
+LocationSelectorOverlay.ATTRS{
+    default_pos={x=48, y=7},
+    viewscreens='dwarfmode/LocationSelector',
+    frame={w=26, h=1},
+}
+
+local function get_religion_string(religion_id, religion_type)
+    if religion_id == -1 then return end
+    local entity
+    if religion_type == 0 then
+        entity = df.historical_figure.find(religion_id)
+    elseif religion_type == 1 then
+        entity = df.historical_entity.find(religion_id)
+    end
+    if not entity then return end
+    return dfhack.TranslateName(entity.name, true)
+end
+
+local function get_profession_string(profession)
+    return df.profession[profession]:gsub('_', ' ')
+end
+
+function LocationSelectorOverlay:init()
+    self:addviews{
+        widgets.BannerPanel{
+            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',
+                    on_change=function(text) self:do_search(text) end,
+                },
+            },
+        },
+    }
+
+    self:register_handler('TEMPLE', location_selector.valid_religious_practice_id,
+        curry(sortoverlay.flags_vector_search, {get_search_key_fn=get_religion_string},
+        location_selector.valid_religious_practice))
+    self:register_handler('GUILDHALL', location_selector.valid_craft_guild_type,
+        curry(sortoverlay.single_vector_search, {get_search_key_fn=get_profession_string}))
+end
+
+function LocationSelectorOverlay:get_key()
+    if location_selector.choosing_temple_religious_practice then
+        return 'TEMPLE'
+    elseif location_selector.choosing_craft_guild then
+        return 'GUILDHALL'
+    end
+end
+
+return _ENV

plugins/lua/sort/sortoverlay.lua

@@ -0,0 +1,169 @@
+local _ENV = mkmodule('plugins.sort.sortoverlay')
+
+local overlay = require('plugins.overlay')
+local utils = require('utils')
+
+local function copy_to_lua_table(vec)
+    local tab = {}
+    for k,v in ipairs(vec) do
+        tab[k+1] = v
+    end
+    return tab
+end
+
+-- ----------------------
+-- SortOverlay
+--
+
+SortOverlay = defclass(SortOverlay, overlay.OverlayWidget)
+SortOverlay.ATTRS{
+    default_enabled=true,
+    hotspot=true,
+    overlay_onupdate_max_freq_seconds=0,
+    -- subclasses expected to provide default_pos, viewscreens (single string), and frame
+}
+
+function SortOverlay:init()
+    self.state = {}
+    self.handlers = {}
+    -- subclasses expected to provide an EditField widget with view_id='search'
+end
+
+function SortOverlay:register_handler(key, vec, search_fn, restore_filtered_on_cleanup)
+    self.handlers[key] = {
+        vec=vec,
+        search_fn=search_fn,
+        restore_filtered_on_cleanup=restore_filtered_on_cleanup
+    }
+end
+
+local function restore_filtered(vec, data)
+    if not data.saved_visible or not data.saved_original then return end
+    for _,elem in ipairs(data.saved_original) do
+        if not utils.linear_index(data.saved_visible, elem) then
+            vec:insert('#', elem)
+        end
+    end
+end
+
+-- handles reset and clean up when the player exits the handled scope
+function SortOverlay:overlay_onupdate()
+    if self.overlay_onupdate_max_freq_seconds == 0 and
+        not dfhack.gui.matchFocusString(self.viewscreens, dfhack.gui.getDFViewscreen(true))
+    then
+        for key,data in pairs(self.state) do
+            if safe_index(self.handlers, key, 'restore_filtered_on_cleanup') then
+                restore_filtered(self.handlers[key].vec, data)
+            end
+        end
+        self.state = {}
+        self.subviews.search:setText('')
+        self.subviews.search:setFocus(false)
+        self.overlay_onupdate_max_freq_seconds = 300
+    end
+end
+
+-- returns the current context key for dereferencing the handler
+-- subclasses must override
+function SortOverlay:get_key()
+    return nil
+end
+
+-- handles saving/restoring search strings when the player moves between different contexts
+function SortOverlay:onRenderBody(dc)
+    if next(self.state) then
+        local key = self:get_key()
+        if self.state.cur_key ~= key then
+            self.state.cur_key = key
+            local prev_text = key and ensure_key(self.state, key).prev_text or ''
+            self.subviews.search:setText(prev_text)
+        end
+    end
+    self.overlay_onupdate_max_freq_seconds = 0
+    SortOverlay.super.onRenderBody(self, dc)
+end
+
+function SortOverlay:onInput(keys)
+    if keys._MOUSE_R and self.subviews.search.focus and self:get_key() then
+        self.subviews.search:setFocus(false)
+        return true
+    end
+    return SortOverlay.super.onInput(self, keys)
+end
+
+function SortOverlay:do_search(text, force_full_search)
+    if not force_full_search and not next(self.state) and text == '' then return end
+    -- the EditField state is guaranteed to be consistent with the current
+    -- context since when clicking to switch tabs, onRenderBody is always called
+    -- before this text_input callback, even if a key is pressed before the next
+    -- graphical frame would otherwise be printed. if this ever becomes untrue,
+    -- then we can add an on_char handler to the EditField that also checks for
+    -- context transitions.
+    local key = self:get_key()
+    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
+    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)
+    self.state[key].prev_text = text
+end
+
+local function filter_vec(fns, flags_vec, vec, text, erase_fn)
+    if fns.matches_filters_fn or text ~= '' then
+        local search_tokens = text:split()
+        for idx = #vec-1,0,-1 do
+            local flag = flags_vec and flags_vec[idx] or nil
+            local search_key = fns.get_search_key_fn(vec[idx], flag)
+            if (search_key and not utils.search_text(search_key, search_tokens)) or
+                (fns.matches_filters_fn and not fns.matches_filters_fn(vec[idx], flag))
+            then
+                erase_fn(idx)
+            end
+        end
+    end
+end
+
+function single_vector_search(fns, vec, data, text, incremental)
+    if not data.saved_original then
+        data.saved_original = copy_to_lua_table(vec)
+    elseif not incremental then
+        vec:assign(data.saved_original)
+    end
+    filter_vec(fns, nil, vec, text, function(idx) vec:erase(idx) end)
+    data.saved_visible = copy_to_lua_table(vec)
+    if fns.get_sort_fn then
+        table.sort(data.saved_visible, fns.get_sort_fn())
+        vec:assign(data.saved_visible)
+    end
+end
+
+-- doesn't support cleanup since nothing that uses this needs it yet
+function flags_vector_search(fns, flags_vec, vec, data, text, incremental)
+    local get_elem_id_fn = fns.get_elem_id_fn and fns.get_elem_id_fn(elem) or function(elem) return elem end
+    if not data.saved_original then
+        data.saved_original = copy_to_lua_table(vec)
+        data.saved_flags = copy_to_lua_table(flags_vec)
+        data.saved_idx_map = {}
+        for idx,elem in ipairs(data.saved_original) do
+            data.saved_idx_map[get_elem_id_fn(elem)] = idx  -- 1-based idx
+        end
+    else  -- sync flag changes to saved vector
+        for idx,elem in ipairs(vec) do  -- 0-based idx
+            data.saved_flags[data.saved_idx_map[get_elem_id_fn(elem)]] = flags_vec[idx]
+        end
+    end
+
+    if not incremental then
+        vec:assign(data.saved_original)
+        flags_vec:assign(data.saved_flags)
+    end
+
+    filter_vec(fns, flags_vec, vec, text, function(idx)
+        vec:erase(idx)
+        flags_vec:erase(idx)
+    end)
+end
+
+return _ENV