support search on the world artifacts screen

plugins/lua/sort.lua

@@ -1289,6 +1289,7 @@ OVERLAY_WIDGETS = {
     interrogation=require('plugins.sort.info').InterrogationOverlay,
     location_selector=require('plugins.sort.locationselector').LocationSelectorOverlay,
     unit_selector=require('plugins.sort.unitselector').UnitSelectorOverlay,
+    world=require('plugins.sort.world').WorldOverlay,
 }
 
 dfhack.onStateChange[GLOBAL_KEY] = function(sc)

plugins/lua/sort/info.lua

@@ -135,6 +135,15 @@ local function work_details_search(vec, data, text, incremental)
         vec, data, text, incremental)
 end
 
+local function cleanup_cri_unit(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
+
 -- ----------------------
 -- InfoOverlay
 --
@@ -177,12 +186,12 @@ function InfoOverlay:init()
                     get_search_key_fn=get_cri_unit_search_key,
                     get_sort_fn=get_sort
                 }),
-            true)
+            curry(cleanup_cri_unit, vec))
     end
 
     self:register_handler('JOBS', tasks.cri_job,
         curry(sortoverlay.single_vector_search, {get_search_key_fn=get_cri_unit_search_key}),
-        true)
+        curry(cleanup_cri_unit, vec))
     self:register_handler('PET_OT', creatures.atk_index,
         curry(sortoverlay.single_vector_search, {get_search_key_fn=get_race_name}))
     self:register_handler('PET_AT', creatures.trainer,

plugins/lua/sort/sortoverlay.lua

@@ -31,31 +31,23 @@ function SortOverlay:init()
     -- subclasses expected to provide an EditField widget with view_id='search'
 end
 
-function SortOverlay:register_handler(key, vec, search_fn, restore_filtered_on_cleanup)
+function SortOverlay:register_handler(key, vec, search_fn, cleanup_fn)
     self.handlers[key] = {
         vec=vec,
         search_fn=search_fn,
-        restore_filtered_on_cleanup=restore_filtered_on_cleanup
+        cleanup_fn=cleanup_fn
     }
 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)
+            local cleanup_fn = safe_index(self.handlers, key, 'cleanup_fn')
+            if cleanup_fn then
+                cleanup_fn(data)
             end
         end
         self:reset()
@@ -133,25 +125,33 @@ local function filter_vec(fns, flags_vec, vec, text, erase_fn)
 end
 
 function single_vector_search(fns, vec, data, text, incremental)
+    vec = utils.getval(vec)
     if not data.saved_original then
         data.saved_original = copy_to_lua_table(vec)
+        data.saved_original_size = #vec
     elseif not incremental then
         vec:assign(data.saved_original)
+        vec:resize(data.saved_original_size)
     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)
+        vec:resize(data.saved_visible_size)
     end
 end
 
--- doesn't support cleanup since nothing that uses this needs it yet
+-- doesn't support sorting 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 or function(elem) return elem end
+    flags_vec, vec = utils.getval(flags_vec), utils.getval(vec)
     if not data.saved_original then
+        -- we save the sizes since trailing nils get lost in the lua -> vec assignment
         data.saved_original = copy_to_lua_table(vec)
+        data.saved_original_size = #vec
         data.saved_flags = copy_to_lua_table(flags_vec)
+        data.saved_flags_size = #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
@@ -164,7 +164,9 @@ function flags_vector_search(fns, flags_vec, vec, data, text, incremental)
 
     if not incremental then
         vec:assign(data.saved_original)
+        vec:resize(data.saved_original_size)
         flags_vec:assign(data.saved_flags)
+        flags_vec:resize(data.saved_flags_size)
     end
 
     filter_vec(fns, flags_vec, vec, text, function(idx)

plugins/lua/sort/world.lua

@@ -0,0 +1,87 @@
+local _ENV = mkmodule('plugins.sort.world')
+
+local sortoverlay = require('plugins.sort.sortoverlay')
+local widgets = require('gui.widgets')
+
+-- ----------------------
+-- WorldOverlay
+--
+
+WorldOverlay = defclass(WorldOverlay, sortoverlay.SortOverlay)
+WorldOverlay.ATTRS{
+    default_pos={x=-18, y=2},
+    viewscreens='world/ARTIFACTS',
+    frame={w=40, h=1},
+}
+
+local function get_world_artifact_search_key(artifact, rumor)
+    local search_key = ('%s %s'):format(dfhack.TranslateName(artifact.name, true),
+        dfhack.items.getDescription(artifact.item, 0))
+    if rumor then
+        local hf = df.historical_figure.find(rumor.hfid)
+        if hf then
+            search_key = ('%s %s %s'):format(search_key,
+                dfhack.TranslateName(hf.name),
+                dfhack.TranslateName(hf.name, true))
+        end
+        local ws = df.world_site.find(rumor.stid)
+        if ws then
+            search_key = ('%s %s'):format(search_key,
+                dfhack.TranslateName(ws.name, true))
+        end
+    else
+        local hf = df.historical_figure.find(artifact.holder_hf)
+        if hf then
+            local unit = df.unit.find(hf.unit_id)
+            if unit then
+                search_key = ('%s %s'):format(search_key,
+                    dfhack.units.getReadableName(unit))
+            end
+        end
+    end
+    return search_key
+end
+
+local function cleanup_artifact_vectors(data)
+    print('cleanng up')
+    local vs_world = dfhack.gui.getDFViewscreen(true)
+    vs_world.artifact:assign(data.saved_original)
+    vs_world.artifact_arl:assign(data.saved_flags)
+end
+
+function WorldOverlay: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('ARTIFACTS',
+        function() return dfhack.gui.getDFViewscreen(true).artifact end,
+        curry(sortoverlay.flags_vector_search,
+            {
+                get_search_key_fn=get_world_artifact_search_key,
+                get_elem_id_fn=function(artifact_record) return artifact_record.id end,
+            },
+            function() return dfhack.gui.getDFViewscreen(true).artifact_arl end),
+        cleanup_artifact_vectors)
+end
+
+function WorldOverlay:get_key()
+    local scr = dfhack.gui.getDFViewscreen(true)
+    if scr.view_mode == df.world_view_mode_type.ARTIFACTS then
+        return 'ARTIFACTS'
+    end
+end
+
+return _ENV