From 56172044808a6a67235ebf2b2c8e8fe71d915a91 Mon Sep 17 00:00:00 2001
From: Myk Taylor <myk002@yahoo.com>
Date: Mon, 20 Nov 2023 12:53:58 -0800
Subject: [PATCH] handle exiting and reentering creature tabs

memory gets reallocated, and we need to keep up
---
 plugins/lua/sort/info.lua        | 26 ++++++++++++++------------
 plugins/lua/sort/sortoverlay.lua | 20 +++++++++++++++-----
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/plugins/lua/sort/info.lua b/plugins/lua/sort/info.lua
index 4e75b2381..8f4a2f3c5 100644
--- a/plugins/lua/sort/info.lua
+++ b/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)
-    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)
+local function free_allocated_data(data)
+    if data.saved_visible and data.saved_original and #data.saved_visible ~= #data.saved_original then
+        for _,elem in ipairs(data.saved_original) do
+            if not utils.linear_index(data.saved_visible, elem) then
+                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()
diff --git a/plugins/lua/sort/sortoverlay.lua b/plugins/lua/sort/sortoverlay.lua
index c234bd83d..ce50ac98a 100644
--- a/plugins/lua/sort/sortoverlay.lua
+++ b/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')
-            if cleanup_fn then
-                cleanup_fn(data)
-            end
+            do_cleanup(self.handlers, key, data)
         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 ''