From 869ef3c79785e8a989cfaa6b12287a9a55855bf6 Mon Sep 17 00:00:00 2001
From: DarkPenguin <96709220+DarkPenguin24@users.noreply.github.com>
Date: Tue, 1 Mar 2022 18:53:50 +0100
Subject: [PATCH] Fix instant win if no roles were selected / present (#305)

---
 hooks/RoleManager.cpp | 4 +++-
 user/utility.cpp      | 5 +++--
 user/utility.h        | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/hooks/RoleManager.cpp b/hooks/RoleManager.cpp
index b7210641..47dabbfb 100644
--- a/hooks/RoleManager.cpp
+++ b/hooks/RoleManager.cpp
@@ -11,7 +11,7 @@
 void dRoleManager_SelectRoles(RoleManager* __this, MethodInfo* method) {
 	std::vector<uint8_t> assignedPlayers;
 	auto allPlayers = GetAllPlayerControl();
-	auto roleRates = RoleRates((*Game::pGameOptionsData)->fields.RoleOptions);
+	auto roleRates = RoleRates((*Game::pGameOptionsData)->fields);
 
 	AssignPreChosenRoles(roleRates, assignedPlayers);
 	AssignRoles(roleRates.ShapeshifterCount, roleRates.ShapeshifterChance, RoleTypes__Enum::Shapeshifter, allPlayers, assignedPlayers);
@@ -45,12 +45,14 @@ void AssignPreChosenRoles(RoleRates& roleRates, std::vector<uint8_t>& assignedPl
 			if (roleRates.ShapeshifterCount == 0)
 				continue;
 			roleRates.ShapeshifterCount--;
+			roleRates.ImposterCount--;
 		}
 		else if (trueRole == RoleTypes__Enum::Impostor)
 		{
 			if (roleRates.ImposterCount == 0)
 				continue;
 			roleRates.ImposterCount--;
+			roleRates.ShapeshifterCount--;
 		}
 		else if (trueRole == RoleTypes__Enum::Scientist)
 		{
diff --git a/user/utility.cpp b/user/utility.cpp
index 03e87273..184ffaa2 100644
--- a/user/utility.cpp
+++ b/user/utility.cpp
@@ -15,8 +15,9 @@ int randi(int lo, int hi) {
 	return lo + i;
 }
 
-RoleRates::RoleRates(RoleOptionsData* roleOptions) {
-	auto roleRates = roleOptions->fields.roleRates;
+RoleRates::RoleRates(GameOptionsData__Fields gameOptionsDataFields) {
+	this->ImposterCount = gameOptionsDataFields.NumImpostors;
+	auto roleRates = gameOptionsDataFields.RoleOptions->fields.roleRates;
 	if (roleRates->fields.count != 0) {
 		auto vectors = roleRates->fields.entries[0].vector;
 		for (auto iVector = 0; iVector < 32; iVector++)
diff --git a/user/utility.h b/user/utility.h
index ab3e1839..af39dcf5 100644
--- a/user/utility.h
+++ b/user/utility.h
@@ -38,7 +38,7 @@ class RoleRates {
 	int EngineerCount = 0;
 	int EngineerChance = 0;
 	int MaxCrewmates = 15;
-	RoleRates(RoleOptionsData* roleOptions);
+	RoleRates(GameOptionsData__Fields gameOptionsDataFields);
 };
 
 class PlayerSelection {