From 52d3cb3ebf3e09d0fa70be0e5b9d8b8a8d7c9233 Mon Sep 17 00:00:00 2001
From: verysecrethero <hnorthh@gmail.com>
Date: Tue, 8 Mar 2016 16:33:08 +0100
Subject: [PATCH] GetPlayerAmmo is now a SetPlayerAmmo instead to fix some
 potential problems if the player has no weapon at all

---
 include/Game/Systems/AmmoPickupSystem.h |  2 +-
 src/Game/Systems/AmmoPickupSystem.cpp   | 19 +++++++++----------
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/include/Game/Systems/AmmoPickupSystem.h b/include/Game/Systems/AmmoPickupSystem.h
index d34427e3..1e61dd32 100644
--- a/include/Game/Systems/AmmoPickupSystem.h
+++ b/include/Game/Systems/AmmoPickupSystem.h
@@ -50,7 +50,7 @@ private:
     //helper methods
     bool DoesPlayerHaveMaxAmmo(EntityWrapper &player);
     PlayerClass DetermineClass(EntityWrapper &player);
-    int& GetPlayerAmmo(EntityWrapper &player);
+    void SetPlayerAmmo(EntityWrapper &player, int ammoGain);
     int GetPlayerMaxAmmo(EntityWrapper &player);
 };
 #endif
diff --git a/src/Game/Systems/AmmoPickupSystem.cpp b/src/Game/Systems/AmmoPickupSystem.cpp
index 21101b69..ea4561c6 100644
--- a/src/Game/Systems/AmmoPickupSystem.cpp
+++ b/src/Game/Systems/AmmoPickupSystem.cpp
@@ -67,16 +67,18 @@ bool AmmoPickupSystem::DoesPlayerHaveMaxAmmo(EntityWrapper &player) {
         return false;
     }
 }
-int& AmmoPickupSystem::GetPlayerAmmo(EntityWrapper &player) {
+void AmmoPickupSystem::SetPlayerAmmo(EntityWrapper &player, int ammoGain) {
+    int maxWeaponAmmo = GetPlayerMaxAmmo(player);
+
     PlayerClass playerClass = DetermineClass(player);
     if (playerClass == PlayerClass::Defender) {
-        return (int)player["DefenderWeapon"]["Ammo"];
+        (int&)player["DefenderWeapon"]["Ammo"] = std::min((int)player["DefenderWeapon"]["Ammo"] + ammoGain, maxWeaponAmmo);
     } else if (playerClass == PlayerClass::Sniper) {
-        return (int)player["SniperWeapon"]["Ammo"];
+        (int&)player["SniperWeapon"]["Ammo"] = std::min((int)player["SniperWeapon"]["Ammo"] + ammoGain, maxWeaponAmmo);
     } else if (playerClass == PlayerClass::Assault) {
-        return (int)player["AssaultWeapon"]["Ammo"];
+        (int&)player["AssaultWeapon"]["Ammo"] = std::min((int)player["AssaultWeapon"]["Ammo"] + ammoGain, maxWeaponAmmo);
     } else {
-        //TODO: should really return something here
+        //unknown class - ignore
     }
 }
 int AmmoPickupSystem::GetPlayerMaxAmmo(EntityWrapper &player) {
@@ -141,12 +143,9 @@ bool AmmoPickupSystem::OnAmmoPickup(Events::AmmoPickup & e)
     if (DoesPlayerHaveMaxAmmo(e.Player)) {
         return false;
     }
-    int maxWeaponAmmo = GetPlayerMaxAmmo(e.Player);
-    int& currentAmmo = GetPlayerAmmo(e.Player);
+    SetPlayerAmmo(e.Player, e.AmmoGain);
 
-    currentAmmo = std::min(currentAmmo + e.AmmoGain, maxWeaponAmmo);
-
-    return false;
+    return true;
 }
 
 bool AmmoPickupSystem::OnTriggerLeave(Events::TriggerLeave& e) {