1 files changed, 24 insertions, 11 deletions
diff --git a/src/openvic-simulation/military/LeaderTrait.cpp b/src/openvic-simulation/military/LeaderTrait.cpp
index 978a3ed..08bb923 100644
--- a/src/openvic-simulation/military/LeaderTrait.cpp
+++ b/src/openvic-simulation/military/LeaderTrait.cpp
@@ -26,20 +26,33 @@ bool LeaderTraitManager::add_leader_trait(
 }
 
 bool LeaderTraitManager::load_leader_traits_file(ModifierManager const& modifier_manager, ast::NodeCPtr root) {
-	using enum LeaderTrait::trait_type_t;
-	const auto trait_callback = [this, &modifier_manager](LeaderTrait::trait_type_t type) -> key_value_callback_t {
-		return [this, &modifier_manager, type](std::string_view trait_identifier, ast::NodeCPtr value) -> bool {
-			ModifierValue modifiers;
-			bool ret =
-				modifier_manager.expect_whitelisted_modifier_value(move_variable_callback(modifiers), allowed_modifiers)(value);
-			ret &= add_leader_trait(trait_identifier, type, std::move(modifiers));
-			return ret;
-		};
+	const auto trait_callback = [this, &modifier_manager](LeaderTrait::trait_type_t type) -> NodeCallback auto {
+		return expect_dictionary_reserve_length(
+			leader_traits,
+			[this, &modifier_manager, type](std::string_view trait_identifier, ast::NodeCPtr value) -> bool {
+				static const string_set_t allowed_modifiers = {
+					"attack", "defence", "morale", "organisation", "reconnaissance",
+					"speed", "attrition", "experience", "reliability"
+				};
+
+				ModifierValue modifiers;
+				bool ret = modifier_manager.expect_whitelisted_modifier_value(
+					move_variable_callback(modifiers), allowed_modifiers
+				)(value);
+
+				ret &= add_leader_trait(trait_identifier, type, std::move(modifiers));
+				return ret;
+			}
+		);
 	};
+
+	using enum LeaderTrait::trait_type_t;
+
 	const bool ret = expect_dictionary_keys(
-		"personality", ONE_EXACTLY, expect_dictionary(trait_callback(PERSONALITY)),
-		"background", ONE_EXACTLY, expect_dictionary(trait_callback(BACKGROUND))
+		"personality", ONE_EXACTLY, trait_callback(PERSONALITY),
+		"background", ONE_EXACTLY, trait_callback(BACKGROUND)
 	)(root);
+
 	lock_leader_traits();
 
 	return ret;