Merge pull request #4048 from myk002/myk_sort_input

[sort] prevent keys from bleeding through to the vanilla handler and prevent double free

docs/changelog.txt

@@ -58,6 +58,8 @@ Template for new versions:
 ## Fixes
 - `buildingplan`: fix choosing the wrong mechanism (or something that isn't a mechanism) when linking a lever and manually choosing a mechanism, but then canceling the selection
 - RemoteServer: don't shut down the socket prematurely, allowing continuing connections from, for example, dfhack-run
+- `sort`: fix potential crash when exiting and re-entering a creatures subtab with a search active
+- `sort`: prevent keyboard keys from affecting the UI when search is active and multiple keys are hit at once
 
 ## Misc Improvements
 - `buildingplan`: save magma safe mechanisms for when magma safety is requested when linking levers and pressure plates to targets

plugins/lua/sort/info.lua

@@ -125,13 +125,15 @@ local function work_details_search(vec, data, text, incremental)
         vec, data, text, incremental)
 end
 
-local function restore_allocated_data(vec, data)
+local function free_allocated_data(data)
-    if not data.saved_visible or not data.saved_original then return end
+    if data.saved_visible and data.saved_original and #data.saved_visible ~= #data.saved_original then
-    for _,elem in ipairs(data.saved_original) do
+        for _,elem in ipairs(data.saved_original) do
-        if not utils.linear_index(data.saved_visible, elem) then
+            if not utils.linear_index(data.saved_visible, elem) then
-            vec:insert('#', elem)
+                elem:delete()
+            end
         end
     end
+    data.saved_original, data.saved_visible = nil, nil
 end
 
 local function serialize_skills(unit)
@@ -311,12 +313,12 @@ function InfoOverlay:init()
                     get_search_key_fn=get_cri_unit_search_key,
                     get_sort_fn=get_sort
                 }),
-            curry(restore_allocated_data, vec))
+            free_allocated_data)
     end
 
     self:register_handler('JOBS', tasks.cri_job,
         curry(sortoverlay.single_vector_search, {get_search_key_fn=get_cri_unit_search_key}),
-        curry(restore_allocated_data, tasks.cri_job))
+        free_allocated_data)
     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,
@@ -345,14 +347,14 @@ function InfoOverlay:get_key()
     if info.current_mode == df.info_interface_mode_type.CREATURES then
         if creatures.current_mode == df.unit_list_mode_type.PET then
             if creatures.showing_overall_training then
-                return 'PET_OT'
+                return 'PET_OT', 'cre'
             elseif creatures.adding_trainer then
-                return 'PET_AT'
+                return 'PET_AT', 'cre'
             elseif creatures.assign_work_animal then
-                return 'PET_WA'
+                return 'PET_WA', 'cre'
             end
         end
-        return df.unit_list_mode_type[creatures.current_mode]
+        return df.unit_list_mode_type[creatures.current_mode], 'cre'
     elseif info.current_mode == df.info_interface_mode_type.JOBS then
         return 'JOBS'
     elseif info.current_mode == df.info_interface_mode_type.ARTIFACTS then
@@ -459,7 +461,7 @@ function CandidatesOverlay:init()
 
     self:register_handler('CANDIDATE', administrators.candidate,
         curry(sortoverlay.single_vector_search, {get_search_key_fn=get_candidate_search_key}),
-        curry(restore_allocated_data, administrators.candidate))
+        free_allocated_data)
 end
 
 function CandidatesOverlay:get_key()

plugins/lua/sort/sortoverlay.lua

@@ -46,16 +46,22 @@ function SortOverlay:register_handler(key, vec, search_fn, cleanup_fn)
     }
 end
 
+local function do_cleanup(handlers, key, data)
+    if not key or not data then return end
+    local cleanup_fn = safe_index(handlers, key, 'cleanup_fn')
+    if cleanup_fn then
+        cleanup_fn(data)
+    end
+    data.saved_original = nil
+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
-            local cleanup_fn = safe_index(self.handlers, key, 'cleanup_fn')
+            do_cleanup(self.handlers, key, data)
-            if cleanup_fn then
-                cleanup_fn(data)
-            end
         end
         self:reset()
         self.overlay_onupdate_max_freq_seconds = 300
@@ -77,7 +83,11 @@ 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()
+        local key, group = self:get_key()
+        if self.state.cur_group ~= group then
+            self.state.cur_group = group
+            do_cleanup(self.handlers, self.state.cur_key, self.state[self.state.cur_key])
+        end
         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 ''
@@ -89,12 +99,23 @@ function SortOverlay:onRenderBody(dc)
     SortOverlay.super.onRenderBody(self, dc)
 end
 
+local function is_mouse_key(keys)
+    return keys._MOUSE_L
+        or keys._MOUSE_R
+        or keys._MOUSE_M
+        or keys.CONTEXT_SCROLL_UP
+        or keys.CONTEXT_SCROLL_DOWN
+        or keys.CONTEXT_SCROLL_PAGEUP
+        or keys.CONTEXT_SCROLL_PAGEDOWN
+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)
+    return SortOverlay.super.onInput(self, keys) or
+        (self.subviews.search.focus and not is_mouse_key(keys))
 end
 
 function SortOverlay:do_search(text, force_full_search)