add support for work details and interrogation

plugins/lua/sort/info.lua

@@ -7,8 +7,9 @@ local utils = require('utils')
 
 local info = df.global.game.main_interface.info
 local creatures = info.creatures
-local objects = info.artifacts
 local justice = info.justice
+local objects = info.artifacts
+local work_details = info.labor.work_details
 
 local state = {}
 
@@ -104,17 +105,19 @@ local function copy_to_lua_table(vec)
     return tab
 end
 
-local function general_search(vec, get_search_key_fn, get_sort_fn, data, filter, incremental)
+local function general_search(vec, get_search_key_fn, get_sort_fn, matches_filters_fn, data, filter, 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
-    if filter ~= '' then
+    if matches_filters_fn ~= DEFAULT_NIL or filter ~= '' then
         local search_tokens = filter:split()
         for idx = #vec-1,0,-1 do
             local search_key = get_search_key_fn(vec[idx])
-            if search_key and not utils.search_text(search_key, search_tokens) then
+            if (search_key and not utils.search_text(search_key, search_tokens)) or
+                (matches_filters_fn ~= DEFAULT_NIL and not matches_filters_fn(vec[idx]))
+            then
                 vec:erase(idx)
             end
         end
@@ -148,30 +151,76 @@ local function make_cri_unitst_handlers(vec)
     }
 end
 
-local function overall_training_search(data, filter, incremental)
+local function overall_training_search(matches_filters_fn, data, filter, incremental)
     general_search(creatures.atk_index, function(elem)
             local raw = df.creature_raw.find(elem)
             if not raw then return end
             return raw.name[1]
-        end, nil, data, filter, incremental)
+        end, nil, matches_filters_fn, data, filter, incremental)
 end
 
-local function assign_trainer_search(data, filter, incremental)
+local function assign_trainer_search(matches_filters_fn, data, filter, incremental)
     general_search(creatures.trainer, function(elem)
             if not elem then return end
             return ('%s %s'):format(dfhack.TranslateName(elem.name), dfhack.units.getProfessionName(elem))
-        end, nil, data, filter, incremental)
+        end, nil, matches_filters_fn, data, filter, incremental)
 end
 
-local function interrogating_search(data, filter, incremental)
+local function get_unit_search_key(unit)
+    return ('%s %s %s'):format(
+        dfhack.units.getReadableName(unit),  -- last name is in english
+        dfhack.units.getProfessionName(unit),
+        dfhack.TranslateName(unit.name, false, true))  -- get untranslated last name
 end
 
-local function convicting_search(data, filter, incremental)
+local function work_details_search(matches_filters_fn, data, filter, incremental)
-    general_search(justice.conviction_list, function(elem)
+    if work_details.selected_work_detail_index ~= data.selected then
-        return ('%s %s'):format(
+        data.saved_original = nil
-            dfhack.units.getReadableName(elem),
+        data.selected = work_details.selected_work_detail_index
-            dfhack.units.getProfessionName(elem))
+    end
-    end, nil, data, filter, incremental)
+    general_search(work_details.assignable_unit, get_unit_search_key,
+        nil, matches_filters_fn, data, filter, incremental)
+end
+
+-- independent implementation of search algorithm since we need to
+-- keep two vectors in sync
+local function interrogating_search(matches_filters_fn, data, filter, incremental)
+    local vec, flags_vec = justice.interrogation_list, justice.interrogation_list_flag
+    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, unit in ipairs(data.saved_original) do
+            data.saved_idx_map[unit.id] = idx  -- 1-based idx
+        end
+    else  -- sync flag changes to saved vector
+        for idx, unit in ipairs(vec) do  -- 0-based idx
+            data.saved_flags[data.saved_idx_map[unit.id]] = flags_vec[idx]
+        end
+    end
+
+    if not incremental then
+        vec:assign(data.saved_original)
+        flags_vec:assign(data.saved_flags)
+    end
+
+    if matches_filters_fn or filter ~= '' then
+        local search_tokens = filter:split()
+        for idx = #vec-1,0,-1 do
+            local search_key = get_unit_search_key(vec[idx])
+            if (search_key and not utils.search_text(search_key, search_tokens)) or
+                (matches_filters_fn and not matches_filters_fn(vec[idx], idx))
+            then
+                vec:erase(idx)
+                flags_vec:erase(idx)
+            end
+        end
+    end
+end
+
+local function convicting_search(matches_filters_fn, data, filter, incremental)
+    general_search(justice.conviction_list, get_unit_search_key,
+        nil, matches_filters_fn, data, filter, incremental)
 end
 
 local HANDLERS = {
@@ -181,6 +230,7 @@ local HANDLERS = {
     DECEASED=make_cri_unitst_handlers(creatures.cri_unit.DECEASED),
     PET_OT={search_fn=overall_training_search},
     PET_AT={search_fn=assign_trainer_search},
+    WORK_DETAILS={search_fn=work_details_search},
     INTERROGATING={search_fn=interrogating_search},
     CONVICTING={search_fn=convicting_search},
 }
@@ -213,6 +263,10 @@ local function get_key()
         return df.unit_list_mode_type[creatures.current_mode]
     elseif info.current_mode == df.info_interface_mode_type.ARTIFACTS then
         return df.artifacts_mode_type[objects.mode]
+    elseif info.current_mode == df.info_interface_mode_type.LABOR then
+        if info.labor.mode == df.labor_mode_type.WORK_DETAILS then
+            return 'WORK_DETAILS'
+        end
     end
 end
 
@@ -225,8 +279,8 @@ local function check_context(self)
     end
 end
 
-local function do_search(self, text)
+local function do_search(matches_filters_fn, text, force_full_search)
-    if not next(state) and text == '' then return end
+    if not force_full_search and not next(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
@@ -238,11 +292,34 @@ local function do_search(self, text)
     local prev_text = ensure_key(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 = prev_text and text:startswith(prev_text)
+    local incremental = not force_full_search and prev_text and text:startswith(prev_text)
-    HANDLERS[key].search_fn(state[key], text, incremental)
+    HANDLERS[key].search_fn(matches_filters_fn, state[key], text, incremental)
     state[key].prev_text = text
 end
 
+local function on_update(self)
+    if self.overlay_onupdate_max_freq_seconds == 0 and
+        not dfhack.gui.matchFocusString('dwarfmode/Info', dfhack.gui.getDFViewscreen(true))
+    then
+        for k,v in pairs(state) do
+            local cleanup_fn = safe_index(HANDLERS, k, 'cleanup_fn')
+            if cleanup_fn then cleanup_fn(v) end
+        end
+        state = {}
+        self.subviews.search:setText('')
+        self.subviews.search:setFocus(false)
+        self.overlay_onupdate_max_freq_seconds = 60
+    end
+end
+
+local function on_input(self, clazz, keys)
+    if keys._MOUSE_R and self.subviews.search.focus then
+        self.subviews.search:setFocus(false)
+        return true
+    end
+    return clazz.super.onInput(self, keys)
+end
+
 -- ----------------------
 -- InfoOverlay
 --
@@ -262,6 +339,7 @@ InfoOverlay.ATTRS{
         'dwarfmode/Info/ARTIFACTS/SYMBOLS',
         'dwarfmode/Info/ARTIFACTS/NAMED_OBJECTS',
         'dwarfmode/Info/ARTIFACTS/WRITTEN_CONTENT',
+        'dwarfmode/Info/LABOR/WORK_DETAILS',
     },
     hotspot=true,
     overlay_onupdate_max_freq_seconds=0,
@@ -279,30 +357,15 @@ function InfoOverlay:init()
                     frame={l=1, t=0, r=1},
                     label_text="Search: ",
                     key='CUSTOM_ALT_S',
-                    on_change=curry(do_search, self),
+                    on_change=curry(do_search, DEFAULT_NIL),
                 },
             },
         },
     }
 end
 
-local function cleanup()
-    for k,v in pairs(state) do
-        local cleanup_fn = safe_index(HANDLERS, k, 'cleanup_fn')
-        if cleanup_fn then cleanup_fn(v) end
-    end
-end
-
 function InfoOverlay:overlay_onupdate()
-    if next(state) and
+    on_update(self)
-        not dfhack.gui.matchFocusString('dwarfmode/Info', dfhack.gui.getDFViewscreen(true))
-    then
-        cleanup()
-        state = {}
-        self.subviews.search:setText('')
-        self.subviews.search:setFocus(false)
-        self.overlay_onupdate_max_freq_seconds = 60
-    end
 end
 
 local function resize_overlay(self)
@@ -320,11 +383,12 @@ end
 
 local function get_panel_offsets()
     local tabs_in_two_rows = is_tabs_in_two_rows()
-    local is_objects = info.current_mode == df.info_interface_mode_type.ARTIFACTS
+    local shift_right = info.current_mode == df.info_interface_mode_type.ARTIFACTS or
-    local l_offset = (not tabs_in_two_rows and is_objects) and 4 or 0
+        info.current_mode == df.info_interface_mode_type.LABOR
+    local l_offset = (not tabs_in_two_rows and shift_right) and 4 or 0
     local t_offset = 1
     if tabs_in_two_rows then
-        t_offset = is_objects and 0 or 3
+        t_offset = shift_right and 0 or 3
     end
     return l_offset, t_offset
 end
@@ -345,16 +409,19 @@ function InfoOverlay:onRenderBody(dc)
     if self:updateFrames() then
         self:updateLayout()
     end
+    if self.refresh_search then
+        self.refresh_search = nil
+        do_search(DEFAULT_NIL, self.subviews.search.text)
+    end
     self.overlay_onupdate_max_freq_seconds = 0
     InfoOverlay.super.onRenderBody(self, dc)
 end
 
 function InfoOverlay:onInput(keys)
-    if keys._MOUSE_R and self.subviews.search.focus then
+    if keys._MOUSE_L and get_key() == 'WORK_DETAILS' then
-        self.subviews.search:setFocus(false)
+        self.refresh_search = true
-        return true
     end
-    return InfoOverlay.super.onInput(self, keys)
+    return on_input(self, InfoOverlay, keys)
 end
 
 -- ----------------------
@@ -367,8 +434,14 @@ InterrogationOverlay.ATTRS{
     default_enabled=true,
     viewscreens='dwarfmode/Info/JUSTICE',
     frame={w=27, h=9},
+    hotspot=true,
+    overlay_onupdate_max_freq_seconds=0,
 }
 
+function InterrogationOverlay:overlay_onupdate()
+    on_update(self)
+end
+
 local function is_interrogate_or_convict()
     local key = get_key()
     return key == 'INTERROGATING' or key == 'CONVICTING'
@@ -388,39 +461,75 @@ function InterrogationOverlay:init()
                     frame={l=0, t=0, r=0},
                     label_text="Search: ",
                     key='CUSTOM_ALT_S',
-                    on_change=curry(do_search, self),
+                    on_change=curry(do_search, self:callback('matches_filters')),
+                },
+                widgets.ToggleHotkeyLabel{
+                    view_id='include_interviewed',
+                    frame={l=0, t=1, w=23},
+                    key='CUSTOM_SHIFT_I',
+                    label='Interviewed:',
+                    options={
+                        {label='Include', value=true, pen=COLOR_GREEN},
+                        {label='Exclude', value=false, pen=COLOR_RED},
+                    },
+                    visible=function() return justice.interrogating end,
+                    on_change=function()
+                        do_search(self:callback('matches_filters'), self.subviews.search.text, true)
+                    end,
                 },
                 widgets.CycleHotkeyLabel{
                     view_id='subset',
-                    frame={l=0, t=1, w=24},
+                    frame={l=0, t=2, w=28},
                     key='CUSTOM_SHIFT_F',
                     label='Show:',
                     options={
                         {label='All', value='all', pen=COLOR_GREEN},
-                        {label='Undead visitors', value='undead', pen=COLOR_RED},
+                        {label='Risky visitors', value='risky', pen=COLOR_RED},
                         {label='Other visitors', value='visitors', pen=COLOR_LIGHTRED},
                         {label='Residents', value='residents', pen=COLOR_YELLOW},
                         {label='Citizens', value='citizens', pen=COLOR_CYAN},
-                        {label='Animals', value='animals', pen=COLOR_MAGENTA},
+                        {label='Animals', value='animals', pen=COLOR_BLUE},
-                        {label='Deceased', value='deceased', pen=COLOR_BLUE},
+                        {label='Deceased or missing', value='deceased', pen=COLOR_MAGENTA},
-                    },
+                        {label='Others', value='others', pen=COLOR_GRAY},
-                },
-                widgets.ToggleHotkeyLabel{
-                    view_id='include_interviewed',
-                    frame={l=0, t=2, w=23},
-                    key='CUSTOM_SHIFT_I',
-                    label='Interviewed:',
-                    options={
-                        {label='Include', value=true, pen=COLOR_GREEN},
-                        {label='Exclude', value=false, pen=COLOR_RED},
                     },
-                    initial_option=true,
+                    on_change=function()
+                        do_search(self:callback('matches_filters'), self.subviews.search.text, true)
+                    end,
                 },
             },
         },
     }
 end
 
+local function is_risky(unit)
+    return false
+end
+
+function InterrogationOverlay:matches_filters(unit, idx)
+    if justice.interrogating then
+        local include_interviewed = self.subviews.include_interviewed:getOptionValue()
+        if not include_interviewed and justice.interrogation_list_flag[idx] == 2 then
+            return false
+        end
+    end
+    local subset = self.subviews.subset:getOptionValue()
+    if subset == 'all' then
+        return true
+    elseif dfhack.units.isDead(unit) or not dfhack.units.isActive(unit) then
+        return subset == 'deceased'
+    elseif dfhack.units.isVisiting(unit) then
+        local risky = is_risky(unit)
+        return (subset == 'risky' and risky) or (subset == 'visitors' and not risky)
+    elseif dfhack.units.isAnimal(unit) then
+        return subset == 'animals'
+    elseif dfhack.units.isCitizen(unit) then
+        return subset == 'citizens'
+    elseif dfhack.units.isOwnGroup(unit) then
+        return subset == 'residents'
+    end
+    return subset == 'others'
+end
+
 function InterrogationOverlay:render(dc)
     local sw = dfhack.screen.getWindowSize()
     local info_panel_border = 31 -- from edges of panel to screen edges
@@ -447,16 +556,16 @@ end
 function InterrogationOverlay:onRenderBody(dc)
     if next(state) then
         check_context(self)
+    else
+        self.subviews.include_interviewed:setOption(true, false)
+        self.subviews.subset:setOption('all')
     end
+    self.overlay_onupdate_max_freq_seconds = 0
     InterrogationOverlay.super.onRenderBody(self, dc)
 end
 
 function InterrogationOverlay:onInput(keys)
-    if keys._MOUSE_R and self.subviews.search.focus then
+    return on_input(self, InterrogationOverlay, keys)
-        self.subviews.search:setFocus(false)
-        return true
-    end
-    return InterrogationOverlay.super.onInput(self, keys)
 end
 
 return _ENV