7 files changed, 93 insertions, 113 deletions

diff --git a/src/openvic-simulation/map/Map.cpp b/src/openvic-simulation/map/Map.cpp
index fde3b3a..cdedebd 100644
--- a/src/openvic-simulation/map/Map.cpp
+++ b/src/openvic-simulation/map/Map.cpp
@@ -108,42 +108,30 @@ void Map::lock_water_provinces() {
 	Logger::info("Locked water provinces after registering ", water_provinces.size());
 }
 
-bool Map::add_region(std::string_view identifier, std::vector<std::string_view> const& province_identifiers) {
+bool Map::add_region(std::string_view identifier, Region::provinces_t const& provinces) {
 	if (identifier.empty()) {
 		Logger::error("Invalid region identifier - empty!");
 		return false;
 	}
-	Region::provinces_t provinces;
-	provinces.reserve(province_identifiers.size());
-	bool meta = false, ret = true;
-	for (const std::string_view province_identifier : province_identifiers) {
-		Province* province = get_province_by_identifier(province_identifier);
-		if (province != nullptr) {
-			if (std::find(provinces.begin(), provinces.end(), province) != provinces.end()) {
-				Logger::error("Duplicate province identifier ", province_identifier, " in region ", identifier);
-				ret = false;
-			} else {
-				if (province->get_has_region()) {
-					meta = true;
-				}
-				provinces.push_back(province);
-			}
-		} else {
-			Logger::error("Invalid province identifier ", province_identifier, " for region ", identifier);
-			ret = false;
-		}
-	}
 	if (provinces.empty()) {
 		Logger::warning("No valid provinces in list for ", identifier);
-		return ret;
+		return true;
 	}
 
-	if (!meta) {
-		for (Province* province : provinces) {
-			province->has_region = true;
+	const bool meta = std::any_of(provinces.begin(), provinces.end(), std::bind_front(&Province::get_has_region));
+
+	Region region { identifier, provinces.front()->get_colour(), meta };
+	bool ret = region.add_provinces(provinces);
+	region.lock();
+	if (regions.add_item(std::move(region))) {
+		if (!meta) {
+			for (Province const* province : provinces) {
+				remove_province_const(province)->has_region = true;
+			}
 		}
+	} else {
+		ret = false;
 	}
-	ret &= regions.add_item({ identifier, std::move(provinces), meta });
 	return ret;
 }
 
@@ -387,36 +375,22 @@ bool Map::load_province_positions(BuildingTypeManager const& building_type_manag
 bool Map::load_region_file(ast::NodeCPtr root) {
 	const bool ret = expect_dictionary_reserve_length(regions,
 		[this](std::string_view region_identifier, ast::NodeCPtr region_node) -> bool {
-			std::vector<std::string_view> province_identifiers;
+			Region::provinces_t provinces;
 			bool ret = expect_list_reserve_length(
-				province_identifiers, expect_identifier(vector_callback(province_identifiers))
+				provinces, expect_province_identifier(vector_callback_pointer(provinces))
 			)(region_node);
-			ret &= add_region(region_identifier, province_identifiers);
+			ret &= add_region(region_identifier, provinces);
 			return ret;
 		}
 	)(root);
 	lock_regions();
-	for (Region& region : regions.get_items()) {
+	for (Region const& region : regions.get_items()) {
 		if (!region.meta) {
-			for (Province* province : region.get_provinces()) {
-				if (!province->get_has_region()) {
-					Logger::error("Province in non-meta region without has_region set: ", province->get_identifier());
-					province->has_region = true;
-				}
-				province->region = &region;
+			for (Province const* province : region.get_provinces()) {
+				remove_province_const(province)->region = &region;
 			}
 		}
 	}
-	for (Province& province : provinces.get_items()) {
-		const bool region_null = province.get_region() == nullptr;
-		if (province.get_has_region() == region_null) {
-			Logger::error(
-				"Province has_region / region mismatch: has_region = ", province.get_has_region(),
-				", region = ", province.get_region()
-			);
-			province.has_region = !region_null;
-		}
-	}
 	return ret;
 }
 
diff --git a/src/openvic-simulation/map/Map.hpp b/src/openvic-simulation/map/Map.hpp
index e477b38..523a8e7 100644
--- a/src/openvic-simulation/map/Map.hpp
+++ b/src/openvic-simulation/map/Map.hpp
@@ -87,6 +87,15 @@ namespace OpenVic {
 		bool add_province(std::string_view identifier, colour_t colour);
 		IDENTIFIER_REGISTRY_NON_CONST_ACCESSORS_CUSTOM_INDEX_OFFSET(province, 1);
 
+		/* This provides a safe way to remove the const qualifier of a Province const*, via a non-const Map.
+		 * It uses a const_cast (the fastest/simplest solution), but this could also be done without it by looking up the
+		 * Province* using the Province const*'s index. Requiring a non-const Map ensures that this function can only be
+		 * used where the Province* could already be accessed by other means, such as the index method, preventing
+		 * misleading code, or in the worst case undefined behaviour. */
+		constexpr Province* remove_province_const(Province const* province) {
+			return const_cast<Province*>(province);
+		}
+
 		bool set_water_province(std::string_view identifier);
 		bool set_water_province_list(std::vector<std::string_view> const& list);
 		void lock_water_provinces();
@@ -97,8 +106,7 @@ namespace OpenVic {
 		Province* get_selected_province();
 		Province::index_t get_selected_province_index() const;
 
-		bool add_region(std::string_view identifier, std::vector<std::string_view> const& province_identifiers);
-		IDENTIFIER_REGISTRY_NON_CONST_ACCESSORS(region)
+		bool add_region(std::string_view identifier, Region::provinces_t const& provinces);
 
 		bool add_mapmode(std::string_view identifier, Mapmode::colour_func_t colour_func);
 
diff --git a/src/openvic-simulation/map/Province.hpp b/src/openvic-simulation/map/Province.hpp
index 8a9f9e9..58fe29e 100644
--- a/src/openvic-simulation/map/Province.hpp
+++ b/src/openvic-simulation/map/Province.hpp
@@ -87,7 +87,7 @@ namespace OpenVic {
 	private:
 		/* Immutable attributes (unchanged after initial game load) */
 		const index_t PROPERTY(index);
-		Region* PROPERTY(region);
+		Region const* PROPERTY(region);
 		bool PROPERTY(on_map);
 		bool PROPERTY(has_region);
 		bool PROPERTY_CUSTOM_PREFIX(water, is);
diff --git a/src/openvic-simulation/map/Region.cpp b/src/openvic-simulation/map/Region.cpp
index 18a47a9..e356c89 100644
--- a/src/openvic-simulation/map/Region.cpp
+++ b/src/openvic-simulation/map/Region.cpp
@@ -4,9 +4,7 @@
 
 using namespace OpenVic;
 
-ProvinceSet::ProvinceSet(provinces_t&& new_provinces) : provinces { std::move(new_provinces) } {}
-
-bool ProvinceSet::add_province(Province* province) {
+bool ProvinceSet::add_province(Province const* province) {
 	if (locked) {
 		Logger::error("Cannot add province to province set - locked!");
 		return false;
@@ -16,13 +14,39 @@ bool ProvinceSet::add_province(Province* province) {
 		return false;
 	}
 	if (contains_province(province)) {
-		Logger::error("Cannot add province ", province->get_identifier(), " to province set - already in the set!");
+		Logger::warning("Cannot add province ", province->get_identifier(), " to province set - already in the set!");
 		return false;
 	}
 	provinces.push_back(province);
 	return true;
 }
 
+bool ProvinceSet::add_provinces(provinces_t const& new_provinces) {
+	bool ret = true;
+	for (Province const* province : new_provinces) {
+		ret &= add_province(province);
+	}
+	return ret;
+}
+
+bool ProvinceSet::remove_province(Province const* province) {
+	if (locked) {
+		Logger::error("Cannot remove province from province set - locked!");
+		return false;
+	}
+	if (province == nullptr) {
+		Logger::error("Cannot remove province from province set - null province!");
+		return false;
+	}
+	const provinces_t::const_iterator it = std::find(provinces.begin(), provinces.end(), province);
+	if (it == provinces.end()) {
+		Logger::warning("Cannot remove province ", province->get_identifier(), " from province set - already not in the set!");
+		return false;
+	}
+	provinces.erase(it);
+	return true;
+}
+
 void ProvinceSet::lock(bool log) {
 	if (locked) {
 		Logger::error("Failed to lock province set - already locked!");
@@ -67,14 +91,8 @@ ProvinceSet::provinces_t const& ProvinceSet::get_provinces() const {
 	return provinces;
 }
 
-static constexpr colour_t ERROR_REGION_COLOUR { colour_t::max_value, 0, 0 };
-
-Region::Region(std::string_view new_identifier, provinces_t&& new_provinces, bool new_meta)
-	: HasIdentifierAndColour {
-		new_identifier, new_provinces.size() > 0 ? new_provinces.front()->get_colour() : ERROR_REGION_COLOUR, false
-	}, ProvinceSet { std::move(new_provinces) }, meta { new_meta } {
-	lock();
-}
+Region::Region(std::string_view new_identifier, colour_t new_colour, bool new_meta)
+	: HasIdentifierAndColour { new_identifier, new_colour, false }, meta { new_meta } {}
 
 bool Region::get_meta() const {
 	return meta;
diff --git a/src/openvic-simulation/map/Region.hpp b/src/openvic-simulation/map/Region.hpp
index 9119b93..e4565c9 100644
--- a/src/openvic-simulation/map/Region.hpp
+++ b/src/openvic-simulation/map/Region.hpp
@@ -5,16 +5,18 @@
 namespace OpenVic {
 
 	struct ProvinceSet {
-		using provinces_t = std::vector<Province*>;
+		using provinces_t = std::vector<Province const*>;
 
 	private:
 		provinces_t provinces;
 		bool locked = false;
 
 	public:
-		ProvinceSet(provinces_t&& new_provinces = {});
-
-		bool add_province(Province* province);
+		/* Returns true if the province is successfully added, false if not (including if it's already in the set). */
+		bool add_province(Province const* province);
+		bool add_provinces(provinces_t const& new_provinces);
+		/* Returns true if the province is successfully removed, false if not (including if it's not in the set). */
+		bool remove_province(Province const* province);
 		void lock(bool log = false);
 		bool is_locked() const;
 		void reset();
@@ -38,7 +40,7 @@ namespace OpenVic {
 		 */
 		const bool meta;
 
-		Region(std::string_view new_identifier, provinces_t&& new_provinces, bool new_meta);
+		Region(std::string_view new_identifier, colour_t new_colour, bool new_meta);
 
 	public:
 		Region(Region&&) = default;
diff --git a/src/openvic-simulation/map/State.cpp b/src/openvic-simulation/map/State.cpp
index 896008e..faf8d1b 100644
--- a/src/openvic-simulation/map/State.cpp
+++ b/src/openvic-simulation/map/State.cpp
@@ -8,76 +8,60 @@ State::State(
 	Country const* owner, Province const* capital, Region::provinces_t&& provinces, Province::colony_status_t colony_status
 ) : owner { owner }, capital { capital }, provinces { std::move(provinces) }, colony_status { colony_status } {}
 
-StateSet::StateSet(Region const& new_region) : region { new_region } {
+/* Whether two provinces in the same region should be grouped into the same state or not.
+ * (Assumes both provinces non-null.) */
+static bool provinces_belong_in_same_state(Province const* lhs, Province const* rhs) {
+	return lhs->get_owner() == rhs->get_owner() && lhs->get_colony_status() == rhs->get_colony_status();
+}
+
+StateSet::StateSet(Map& map, Region const& new_region) : region { new_region } {
 	if (region.get_meta()) {
 		Logger::error("Cannot use meta region as state template!");
 	}
 
 	std::vector<Region::provinces_t> temp_provinces;
-	bool in_state = false;
 
-	for (Province* province : region.get_provinces()) {
+	for (Province const* province : region.get_provinces()) {
 		// add to existing state if shared owner & status...
 		for (Region::provinces_t& provinces : temp_provinces) {
-			if (provinces[0] == province) {
+			if (provinces_belong_in_same_state(provinces[0], province)) {
 				provinces.push_back(province);
-				in_state = true;
-				break;
+				// jump to the end of the outer loop, skipping the new state code
+				goto loop_end;
 			}
 		}
-		if (in_state) {
-			in_state = false;
-		} else {
-			// ...otherwise start a new state
-			temp_provinces.push_back({ province });
-		}
+		// ...otherwise start a new state
+		temp_provinces.push_back({ province });
+	loop_end:;
+		/* Either the province was added to an existing state and the program jumped to here,
+		 * or it was used to create a new state and the program arrived here normally. */
 	}
 
 	for (Region::provinces_t& provinces : temp_provinces) {
-		states.push_back({
+		states.emplace_back(
 			/* TODO: capital province logic */
 			provinces[0]->get_owner(), provinces[0], std::move(provinces), provinces[0]->get_colony_status()
-		});
+		);
 	}
 
 	// Go back and assign each new state to its provinces.
 	for (State const& state : states) {
-		for (Province* province : state.get_provinces()) {
-			province->set_state(&state);
+		for (Province const* province : state.get_provinces()) {
+			map.remove_province_const(province)->set_state(&state);
 		}
 	}
 }
 
-bool StateSet::add_state(State&& state) {
-	const auto existing = std::find(states.begin(), states.end(), state);
-	if (existing != states.end()) {
-		Logger::error("Attempted to add existing state!");
-		return false;
-	}
-	states.push_back(std::move(state));
-	return true;
-}
-
-bool StateSet::remove_state(State const* state) {
-	const auto existing = std::find(states.begin(), states.end(), *state);
-	if (existing == states.end()) {
-		Logger::error("Attempted to remove non-existant state!");
-		return false;
-	}
-	states.erase(existing);
-	return true;
-}
-
 StateSet::states_t& StateSet::get_states() {
 	return states;
 }
 
-void StateManager::generate_states(Map const& map) {
+void StateManager::generate_states(Map& map) {
 	regions.clear();
 	regions.reserve(map.get_region_count());
 	for(Region const& region : map.get_regions()) {
 		if (!region.get_meta()) {
-			regions.push_back(StateSet(region));
+			regions.emplace_back(map, region);
 		}
 	}
 	Logger::info("Generated states.");
diff --git a/src/openvic-simulation/map/State.hpp b/src/openvic-simulation/map/State.hpp
index bb23b2d..bae83f7 100644
--- a/src/openvic-simulation/map/State.hpp
+++ b/src/openvic-simulation/map/State.hpp
@@ -21,22 +21,16 @@ namespace OpenVic {
 		);
 	};
 
-	inline bool operator==(const State& lhs, const State& rhs) {
-		return (lhs.get_owner() == rhs.get_owner() && lhs.get_colony_status() == rhs.get_colony_status());
-	}
-
 	struct StateSet {
 		using states_t = std::deque<State>;
 
 	private:
 		Region const& PROPERTY(region);
-		states_t states;
+		states_t PROPERTY(states);
 
 	public:
-		StateSet(Region const& new_region);
+		StateSet(Map& map, Region const& new_region);
 
-		bool add_state(State&& state);
-		bool remove_state(State const* state);
 		states_t& get_states();
 	};
 
@@ -49,6 +43,6 @@ namespace OpenVic {
 		/* Creates states from current province gamestate & regions, sets province state value.
 		 * After this function, the `regions` property is unmanaged and must be carefully updated and
 		 * validated by functions that modify it. */
-		void generate_states(Map const& map);
+		void generate_states(Map& map);
 	};
 } // namespace OpenVic