Make logic closer to C# version. Address code comments from peterix.

plugins/autolabor/autolabor.cpp

@@ -18,6 +18,8 @@
 #include <df/unit_labor.h>
 #include <df/unit_skill.h>
 #include <df/job.h>
+#include <df/building.h>
+#include <df/workshop_type.h>
 #include <df/unit_misc_trait.h>
 
 using namespace DFHack;
@@ -28,6 +30,7 @@ using df::global::world;
 // our own, empty header.
 #include "autolabor.h"
 
+#define ARRAY_COUNT(array) (sizeof(array)/sizeof((array)[0]))
 
 static int enable_autolabor;
 
@@ -40,56 +43,11 @@ command_result autolabor (color_ostream &out, std::vector <std::string> & parame
 // The name string provided must correspond to the filename - autolabor.plug.so or autolabor.plug.dll in this case
 DFHACK_PLUGIN("autolabor");
 
-DFhackCExport command_result plugin_init ( color_ostream &out, std::vector <PluginCommand> &commands)
-{
-    // Fill the command list with your commands.
-    commands.clear();
-    commands.push_back(PluginCommand(
-        "autolabor", "Do nothing, look pretty.",
-        autolabor, false, /* true means that the command can't be used from non-interactive user interface */
-        // Extended help string. Used by CR_WRONG_USAGE and the help command:
-        "  This command does nothing at all.\n"
-        "Example:\n"
-        "  autolabor\n"
-        "    Does nothing.\n"
-    ));
-    return CR_OK;
-}
-
-// This is called right before the plugin library is removed from memory.
-DFhackCExport command_result plugin_shutdown ( color_ostream &out )
-{
-    // You *MUST* kill all threads you created before this returns.
-    // If everything fails, just return CR_FAILURE. Your plugin will be
-    // in a zombie state, but things won't crash.
-    return CR_OK;
-}
-
-// Called to notify the plugin about important state changes.
-// Invoked with DF suspended, and always before the matching plugin_onupdate.
-// More event codes may be added in the future.
-/*
-DFhackCExport command_result plugin_onstatechange(color_ostream &out, state_change_event event)
-{
-    switch (event) {
-    case SC_GAME_LOADED:
-        // initialize from the world just loaded
-        break;
-    case SC_GAME_UNLOADED:
-        // cleanup
-        break;
-    default:
-        break;
-    }
-    return CR_OK;
-}
-*/
-
 enum labor_mode {
+	FIXED,
+	AUTOMATIC,
 	EVERYONE,
 	HAULERS,
-	AUTOMATIC,
-	FIXED,
 };
 
 enum dwarf_state {
@@ -345,7 +303,7 @@ struct labor_info
 	int minimum_dwarfs;
 };
 
-static const struct labor_info labor_info[] = {
+static const struct labor_info labor_infos[] = {
     /* MINE */				{AUTOMATIC, true, 2},
     /* HAUL_STONE */		{HAULERS, false, 1},
     /* HAUL_WOOD */			{HAULERS, false, 1},
@@ -438,6 +396,34 @@ struct dwarf_info
 	bool has_exclusive_labor;
 };
 
+DFhackCExport command_result plugin_init ( color_ostream &out, std::vector <PluginCommand> &commands)
+{
+	assert(ARRAY_COUNT(labor_infos) > ENUM_LAST_ITEM(unit_labor));
+
+	// Fill the command list with your commands.
+    commands.clear();
+    commands.push_back(PluginCommand(
+        "autolabor", "Automatically manage dwarf labors.",
+        autolabor, false, /* true means that the command can't be used from non-interactive user interface */
+        // Extended help string. Used by CR_WRONG_USAGE and the help command:
+        "  autolabor enable\n"
+		"  autolabor disable\n"
+		"    Enables or disables the plugin.\n"
+		"Function:\n"
+		"  When enabled, autolabor periodically checks your dwarves and enables or\n"
+		"  disables labors. It tries to keep as many dwarves as possible busy but\n"
+		"  also tries to have dwarves specialize in specific skills.\n"
+		"  Warning: autolabor will override any manual changes you make to labors\n"
+		"  while it is enabled.\n"
+    ));
+    return CR_OK;
+}
+
+DFhackCExport command_result plugin_shutdown ( color_ostream &out )
+{
+	return CR_OK;
+}
+
 DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 {
 	static int step_count = 0;
@@ -454,10 +440,27 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 
 	std::vector<df::unit *> dwarfs;
 
+	bool has_butchers = false;
+	bool has_fishery = false;
+
+	for (int i = 0; i < world->buildings.all.size(); ++i)
+	{
+		df::building *build = world->buildings.all[i];
+		auto type = build->getType();
+		if (df::enums::building_type::Workshop == type)
+		{
+			auto subType = build->getSubtype();
+			if (df::enums::workshop_type::Butchers == subType)
+				has_butchers = true;
+			if (df::enums::workshop_type::Fishery == subType)
+				has_fishery = true;
+		}
+	}
+
     for (int i = 0; i < world->units.all.size(); ++i)
     {
         df::unit* cre = world->units.all[i];
-		if (cre->race == race && cre->civ_id == civ && !cre->flags1.bits.marauder && !cre->flags1.bits.diplomat && !cre->flags1.bits.merchant) {
+		if (cre->race == race && cre->civ_id == civ && !cre->flags1.bits.marauder && !cre->flags1.bits.diplomat && !cre->flags1.bits.merchant && !cre->flags1.bits.dead) {
 			dwarfs.push_back(cre);
         }
     }
@@ -479,7 +482,7 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 	{
 		assert(dwarfs[dwarf]->status.souls.size() > 0);
 
-		for (auto s = dwarfs[dwarf]->status.souls[0]->skills.begin(); s < dwarfs[dwarf]->status.souls[0]->skills.end(); s++)
+		for (auto s = dwarfs[dwarf]->status.souls[0]->skills.begin(); s != dwarfs[dwarf]->status.souls[0]->skills.end(); s++)
 		{
 			df::job_skill skill = (*s)->id;
 
@@ -548,9 +551,9 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 				continue;
 
 			assert(labor >= 0);
-			assert(labor < _countof(labor_info));
+			assert(labor < _countof(labor_infos));
 
-			if (labor_info[labor].is_exclusive && dwarfs[dwarf]->status.labors[labor])
+			if (labor_infos[labor].is_exclusive && dwarfs[dwarf]->status.labors[labor])
 				dwarf_info[dwarf].mastery_penalty -= 100;
 		}
 	}
@@ -602,68 +605,6 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 		state_count[dwarf_info[dwarf].state]++;
 	}
 
-	// Set about 1/3 of the dwarfs as haulers. The haulers have all HAULER labors enabled. Having a lot of haulers helps
-	// make sure that hauling jobs are handled quickly rather than building up.
-
-	int num_haulers = state_count[IDLE] + state_count[BUSY] / 3;
-	if (num_haulers < 1)
-		num_haulers = 1;
-
-	std::vector<int> hauler_ids;
-	for (int dwarf = 0; dwarf < n_dwarfs; dwarf++)
-	{
-		if (dwarf_info[dwarf].state == IDLE || dwarf_info[dwarf].state == BUSY)
-			hauler_ids.push_back(dwarf);
-	}
-
-	// Idle dwarves come first, then we sort from least-skilled to most-skilled.
-	
-	std::sort(hauler_ids.begin(), hauler_ids.end(), [&dwarf_info] (int i, int j) -> bool
-	{ 
-		if (dwarf_info[i].state == IDLE && dwarf_info[j].state != IDLE)
-			return true;
-		if (dwarf_info[i].state != IDLE && dwarf_info[j].state == IDLE)
-			return false;
-		return dwarf_info[i].mastery_penalty > dwarf_info[j].mastery_penalty; 
-	});
-
-	FOR_ENUM_ITEMS(unit_labor, labor)
-	{
-		if (labor == df::enums::unit_labor::NONE)
-			continue;
-
-		assert(labor >= 0);
-		assert(labor < _countof(labor_info));
-
-		if (labor_info[labor].mode != HAULERS)
-			continue;
-
-		for (int i = 0; i < num_haulers; i++)
-		{
-			assert(i < hauler_ids.size());
-
-			int dwarf = hauler_ids[i];
-
-			assert(dwarf >= 0);
-			assert(dwarf < n_dwarfs);
-
-			dwarfs[dwarf]->status.labors[labor] = true;
-			dwarf_info[dwarf].assigned_jobs++;
-		}
-
-		for (int i = num_haulers; i < hauler_ids.size(); i++)
-		{
-			assert(i < hauler_ids.size());
-
-			int dwarf = hauler_ids[i];
-			
-			assert(dwarf >= 0);
-			assert(dwarf < n_dwarfs);
-
-			dwarfs[dwarf]->status.labors[labor] = false;
-		}
-	}
-
 	// Generate labor -> skill mapping
 
 	df::job_skill labor_to_skill[ENUM_LAST_ITEM(unit_labor) + 1];
@@ -682,7 +623,7 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 		}
 	}
 
-	// Handle all skills except those marked HAULERS
+	std::vector<df::unit_labor> labors;
 
 	FOR_ENUM_ITEMS(unit_labor, labor)
 	{
@@ -690,11 +631,25 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 			continue;
 
 		assert(labor >= 0);
-		assert(labor < _countof(labor_info));
+		assert(labor < _countof(labor_infos));
+
+		labors.push_back(labor);
+	}
+
+	std::sort(labors.begin(), labors.end(), [] (int i, int j) { return labor_infos[i].mode < labor_infos[j].mode; });
+
+	// Handle all skills except those marked HAULERS
+
+	for (auto lp = labors.begin(); lp != labors.end(); ++lp)
+	{
+		auto labor = *lp;
+
+		assert(labor >= 0);
+		assert(labor < _countof(labor_infos));
 
 		df::job_skill skill = labor_to_skill[labor];
 
-		if (labor_info[labor].mode == HAULERS)
+		if (labor_infos[labor].mode == HAULERS)
 			continue;
 
 		int best_dwarf = 0;
@@ -702,20 +657,21 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 
 		std::vector<int> values(n_dwarfs);
 		std::vector<int> candidates;
+		std::vector<int> backup_candidates;
+		std::map<int, int> dwarf_skill;
+
+		auto mode = labor_infos[labor].mode;
+		if (AUTOMATIC == mode && state_count[IDLE] == 0)
+			mode = FIXED;
 
 		for (int dwarf = 0; dwarf < n_dwarfs; dwarf++)
 		{
-			if (labor_info[labor].mode == AUTOMATIC && state_count[IDLE] > 0 && dwarf_info[dwarf].state != IDLE)
+			if (dwarf_info[dwarf].state != IDLE && dwarf_info[dwarf].state != BUSY && mode != EVERYONE)
 				continue;
 
-			if (dwarf_info[dwarf].state != IDLE && dwarf_info[dwarf].state != BUSY)
+			if (labor_infos[labor].is_exclusive && dwarf_info[dwarf].has_exclusive_labor)
 				continue;
 
-			if (labor_info[labor].is_exclusive && dwarf_info[dwarf].has_exclusive_labor)
-				continue;
-
-			candidates.push_back(dwarf);
-
 			int value = dwarf_info[dwarf].mastery_penalty - dwarf_info[dwarf].assigned_jobs;
 
 			if (skill != df::enums::job_skill::NONE)
@@ -733,6 +689,8 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 					}
 				}
 
+				dwarf_skill[dwarf] = skill_level;
+
 				value += skill_level * 100;
 				value += skill_experience / 20;
 				if (skill_level > 0 || skill_experience > 0)
@@ -740,26 +698,63 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 				if (skill_level >= 15)
 					value += 1000 * (skill_level - 14);
 			}
+			else
+			{
+				dwarf_skill[dwarf] = 0;
+			}
 
 			if (dwarfs[dwarf]->status.labors[labor])
 			{
 				value += 5;
-				if (labor_info[labor].is_exclusive)
+				if (labor_infos[labor].is_exclusive)
 					value += 350;
 			}
 
 			values[dwarf] = value;
+
+			if (mode == AUTOMATIC && dwarf_info[dwarf].state != IDLE)
+				backup_candidates.push_back(dwarf);
+			else
+				candidates.push_back(dwarf);
+
+		}
+
+		if (candidates.size() == 0)
+		{
+			candidates = backup_candidates;
+			mode = FIXED;
 		}
 
-		if (labor_info[labor].mode != EVERYONE)
+		if (labor_infos[labor].mode != EVERYONE)
 			std::sort(candidates.begin(), candidates.end(), [&values] (int i, int j) { return values[i] > values[j]; });
 
 		for (int dwarf = 0; dwarf < n_dwarfs; dwarf++)
 		{
-			if (state_count[IDLE] > 0 && dwarf_info[dwarf].state == BUSY && labor_info[labor].mode == AUTOMATIC)
+			bool allow_labor = false;
+			
+			if (dwarf_info[dwarf].state == BUSY && 
+				  mode == AUTOMATIC && 
+				  (labor_infos[labor].is_exclusive || dwarf_skill[dwarf] > 0))
 			{
-				if (dwarfs[dwarf]->status.labors[labor])
+				allow_labor = true;
-					dwarf_info[dwarf].assigned_jobs++;
+			}
+
+			if (dwarf_info[dwarf].state == OTHER &&
+				  mode == AUTOMATIC && 
+				  dwarf_skill[dwarf] > 0 &&
+				  !dwarf_info[dwarf].is_best_noble)
+			{
+				allow_labor = true;
+			}
+
+			if (dwarfs[dwarf]->status.labors[labor] &&
+				allow_labor && 
+				!(labor_infos[labor].is_exclusive && dwarf_info[dwarf].has_exclusive_labor))
+			{
+				if (labor_infos[labor].is_exclusive)
+					dwarf_info[dwarf].has_exclusive_labor = true;
+
+				dwarf_info[dwarf].assigned_jobs++;
 			}
 			else
 			{
@@ -767,11 +762,17 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 			}
 		}
 
-		int minimum_dwarfs = labor_info[labor].minimum_dwarfs;
+		int minimum_dwarfs = labor_infos[labor].minimum_dwarfs;
 
-		if (labor_info[labor].mode == EVERYONE)
+		if (labor_infos[labor].mode == EVERYONE)
 			minimum_dwarfs = n_dwarfs;
 
+		// Special - don't assign hunt without a butchers, or fish without a fishery
+		if (df::enums::unit_labor::HUNT == labor && !has_butchers)
+			minimum_dwarfs = 0;
+		if (df::enums::unit_labor::FISH == labor && !has_fishery)
+			minimum_dwarfs = 0;
+
 		for (int i = 0; i < candidates.size() && i < minimum_dwarfs; i++)
 		{
 			int dwarf = candidates[i];
@@ -784,11 +785,73 @@ DFhackCExport command_result plugin_onupdate ( color_ostream &out )
 
 			dwarfs[dwarf]->status.labors[labor] = true;
 
-			if (labor_info[labor].is_exclusive)
+			if (labor_infos[labor].is_exclusive)
 				dwarf_info[dwarf].has_exclusive_labor = true;
 		}
 	}
 
+	// Set about 1/3 of the dwarfs as haulers. The haulers have all HAULER labors enabled. Having a lot of haulers helps
+	// make sure that hauling jobs are handled quickly rather than building up.
+
+	int num_haulers = state_count[IDLE] + state_count[BUSY] / 3;
+	if (num_haulers < 1)
+		num_haulers = 1;
+
+	std::vector<int> hauler_ids;
+	for (int dwarf = 0; dwarf < n_dwarfs; dwarf++)
+	{
+		if (dwarf_info[dwarf].state == IDLE || dwarf_info[dwarf].state == BUSY)
+			hauler_ids.push_back(dwarf);
+	}
+
+	// Idle dwarves come first, then we sort from least-skilled to most-skilled.
+	
+	std::sort(hauler_ids.begin(), hauler_ids.end(), [&dwarf_info] (int i, int j) -> bool
+	{ 
+		if (dwarf_info[i].state == IDLE && dwarf_info[j].state != IDLE)
+			return true;
+		if (dwarf_info[i].state != IDLE && dwarf_info[j].state == IDLE)
+			return false;
+		return dwarf_info[i].mastery_penalty > dwarf_info[j].mastery_penalty; 
+	});
+
+	FOR_ENUM_ITEMS(unit_labor, labor)
+	{
+		if (labor == df::enums::unit_labor::NONE)
+			continue;
+
+		assert(labor >= 0);
+		assert(labor < _countof(labor_infos));
+
+		if (labor_infos[labor].mode != HAULERS)
+			continue;
+
+		for (int i = 0; i < num_haulers; i++)
+		{
+			assert(i < hauler_ids.size());
+
+			int dwarf = hauler_ids[i];
+
+			assert(dwarf >= 0);
+			assert(dwarf < n_dwarfs);
+
+			dwarfs[dwarf]->status.labors[labor] = true;
+			dwarf_info[dwarf].assigned_jobs++;
+		}
+
+		for (int i = num_haulers; i < hauler_ids.size(); i++)
+		{
+			assert(i < hauler_ids.size());
+
+			int dwarf = hauler_ids[i];
+			
+			assert(dwarf >= 0);
+			assert(dwarf < n_dwarfs);
+
+			dwarfs[dwarf]->status.labors[labor] = false;
+		}
+	}
+
     return CR_OK;
 }