Merge pull request #5649 from ab9rf/clean-sort-filter-function

`sort`: make removal of our filter functions more reliable

Author: ab9rf

docs/changelog.txt

@@ -59,6 +59,7 @@ Template for new versions:
 ## New Features
 
 ## Fixes
+- `sort`: Using the squad unit selector will no longer cause Dwarf Fortress to crash on exit
 
 ## Misc Improvements
 

plugins/sort.cpp

@@ -1,4 +1,5 @@
 #include <variant>
+#include <functional>
 
 #include "Debug.h"
 #include "LuaTools.h"
@@ -13,6 +14,7 @@
 #include "df/widget_unit_list.h"
 #include "df/world.h"
 
+
 using std::vector;
 using std::string;
 
@@ -28,7 +30,8 @@ namespace DFHack {
 }
 
 using item_or_unit = std::variant<df::unit*, df::item*>;
-using filter_vec_type = std::vector<std::function<bool(item_or_unit)>>;
+using filter_func = bool(item_or_unit);
+using filter_vec_type = std::vector<std::function<filter_func>>;
 
 // recreated here since our autogenerated df::sort_entry lacks template params
 struct sort_entry {
@@ -42,19 +45,10 @@ static const string DFHACK_SORT_IDENT = "dfhack_sort";
 // filter logic
 //
 
-static bool probing = false;
-static bool probe_result = false;
-
 static bool do_filter(const char *module_name, const char *fn_name, const item_or_unit &elem) {
     if (std::holds_alternative<df::item*>(elem)) return true;
     auto unit = std::get<df::unit*>(elem);
 
-    if (probing) {
-        TRACE(log).print("probe successful\n");
-        probe_result = true;
-        return false;
-    }
-
     bool ret = true;
     color_ostream &out = Core::getInstance().getConsole();
     Lua::CallLuaModuleFunction(out, module_name, fn_name, std::make_tuple(unit),
@@ -78,35 +72,28 @@ static bool do_work_animal_assignment_filter(item_or_unit elem) {
     return do_filter("plugins.sort.info", "do_work_animal_assignment_filter", elem);
 }
 
-static int32_t our_filter_idx(df::widget_unit_list *unitlist) {
-    if (world->units.active.empty())
-        return -1;
-
-    df::unit *u = world->units.active[0]; // any unit will do; we just need a sentinel
-    if (!u)
-        return -1;
-
-    probing = true;
-    probe_result = false;
+static int32_t our_filter_idx(filter_func* filter, df::widget_unit_list* unitlist)
+{
     int32_t idx = 0;
 
-    filter_vec_type *filter_vec = reinterpret_cast<filter_vec_type *>(&unitlist->filter_func);
+    filter_vec_type* filter_vec = reinterpret_cast<filter_vec_type*>(&unitlist->filter_func);
 
     TRACE(log).print("probing for our filter function\n");
-    for (auto& fn : *filter_vec) {
-        fn(item_or_unit(u));
-        if (probe_result) {
+
+    for (auto& fn : *filter_vec)
+    {
+        auto t = fn.target<decltype(filter)>();
+        if (t && *t == filter)
+        {
             TRACE(log).print("found our filter function at idx %d\n", idx);
-            break;
+            return idx;
         }
         ++idx;
     }
-
-    probing = false;
-    return probe_result ? idx : -1;
+    return -1;
 }
 
-static df::widget_unit_list * get_squad_unit_list() {
+static df::widget_unit_list* get_squad_unit_list() {
     return virtual_cast<df::widget_unit_list>(
         Gui::getWidget(&game->main_interface.unit_selector, "Unit selector"));
 }
@@ -182,8 +169,8 @@ DFhackCExport command_result plugin_init(color_ostream &out, vector<PluginComman
     return CR_OK;
 }
 
-static void remove_filter_function(color_ostream &out, const char *which, df::widget_unit_list *unitlist) {
-    int32_t idx = our_filter_idx(unitlist);
+static void remove_filter_function(color_ostream &out, filter_func* filter, const char *which, df::widget_unit_list *unitlist) {
+    int32_t idx = our_filter_idx(filter, unitlist);
     if (idx >= 0) {
         DEBUG(log,out).print("removing %s filter function\n", which);
         filter_vec_type *filter_vec = reinterpret_cast<filter_vec_type *>(&unitlist->filter_func);
@@ -202,22 +189,22 @@ static void remove_sort_function(color_ostream &out, const char *which, df::widg
 
 DFhackCExport command_result plugin_shutdown(color_ostream &out) {
     if (auto unitlist = get_squad_unit_list()) {
-        remove_filter_function(out, "squad", unitlist);
+        remove_filter_function(out, do_squad_filter, "squad", unitlist);
         remove_sort_function(out, "squad", unitlist);
     }
 
     if (auto unitlist = get_interrogate_unit_list("Open cases"))
-        remove_filter_function(out, "open cases interrogate", unitlist);
+        remove_filter_function(out, do_justice_filter, "open cases interrogate", unitlist);
     if (auto unitlist = get_interrogate_unit_list("Cold cases"))
-        remove_filter_function(out, "cold cases interrogate", unitlist);
+        remove_filter_function(out, do_justice_filter, "cold cases interrogate", unitlist);
 
     if (auto unitlist = get_convict_unit_list("Open cases"))
-        remove_filter_function(out, "open cases convict", unitlist);
+        remove_filter_function(out, do_justice_filter, "open cases convict", unitlist);
     if (auto unitlist = get_convict_unit_list("Cold cases"))
-        remove_filter_function(out, "cold cases convict", unitlist);
+        remove_filter_function(out, do_justice_filter, "cold cases convict", unitlist);
 
     if (auto unitlist = get_work_animal_assignment_unit_list())
-        remove_filter_function(out, "work animal assignment", unitlist);
+        remove_filter_function(out, do_work_animal_assignment_filter, "work animal assignment", unitlist);
 
     return CR_OK;
 }
@@ -228,7 +215,7 @@ DFhackCExport command_result plugin_shutdown(color_ostream &out) {
 
 static void sort_set_squad_filter_fn(color_ostream &out) {
     auto unitlist = get_squad_unit_list();
-    if (unitlist && our_filter_idx(unitlist) == -1) {
+    if (unitlist && our_filter_idx(do_squad_filter, unitlist) == -1) {
         DEBUG(log).print("adding squad filter function\n");
         auto filter_vec = reinterpret_cast<filter_vec_type *>(&unitlist->filter_func);
         filter_vec->emplace_back(do_squad_filter);
@@ -240,7 +227,7 @@ static void sort_set_squad_filter_fn(color_ostream &out) {
 }
 
 static void sort_set_justice_filter_fn(color_ostream &out, df::widget_unit_list *unitlist) {
-    if (unitlist && our_filter_idx(unitlist) == -1) {
+    if (unitlist && our_filter_idx(do_justice_filter, unitlist) == -1) {
         DEBUG(log).print("adding justice filter function\n");
         auto filter_vec = reinterpret_cast<filter_vec_type *>(&unitlist->filter_func);
         filter_vec->emplace_back(do_justice_filter);
@@ -249,7 +236,7 @@ static void sort_set_justice_filter_fn(color_ostream &out, df::widget_unit_list
 }
 
 static void sort_set_work_animal_assignment_filter_fn(color_ostream &out, df::widget_unit_list *unitlist) {
-    if (unitlist && our_filter_idx(unitlist) == -1) {
+    if (unitlist && our_filter_idx(do_work_animal_assignment_filter, unitlist) == -1) {
         DEBUG(log).print("adding work animal assignment filter function\n");
         auto filter_vec = reinterpret_cast<filter_vec_type *>(&unitlist->filter_func);
         filter_vec->emplace_back(do_work_animal_assignment_filter);