From ab9b6320cb67843980d4c4e6f81323e8b5b04959 Mon Sep 17 00:00:00 2001
From: Boondorl <boondorl@live.com>
Date: Thu, 7 Nov 2024 18:54:23 -0500
Subject: [PATCH] Allow easier piece weapon replacing

Checks for replacements on weapons instead of using the given weapon class as is (also verifies said replacement is a weapon).
---
 .../zscript/actors/inventory/weaponpiece.zs   | 24 +++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/wadsrc/static/zscript/actors/inventory/weaponpiece.zs b/wadsrc/static/zscript/actors/inventory/weaponpiece.zs
index 2f2434d6c..2a2cd9fe6 100644
--- a/wadsrc/static/zscript/actors/inventory/weaponpiece.zs
+++ b/wadsrc/static/zscript/actors/inventory/weaponpiece.zs
@@ -59,6 +59,13 @@ class WeaponPiece : Inventory
 	
 	property number: PieceValue;
 	property weapon: WeaponClass;
+
+	// Account for weapon replacers, but make sure it's still a Weapon
+	clearscope class<Weapon> GetWeaponClass() const
+	{
+		class<Weapon> type = WeaponClass ? (class<Weapon>)(GetReplacement(WeaponClass)) : null;
+		return type ? type : WeaponClass;
+	}
 	
 	//==========================================================================
 	//
@@ -74,7 +81,11 @@ class WeaponPiece : Inventory
 			return false;
 		}
 
-		let Defaults = GetDefaultByType(WeaponClass);
+		class<Weapon> type = GetWeaponClass();
+		if (!type)
+			return false;
+
+		let Defaults = GetDefaultByType(type);
 
 		bool gaveSome = !!(toucher.GiveAmmo (Defaults.AmmoType1, Defaults.AmmoGive1) +
 						   toucher.GiveAmmo (Defaults.AmmoType2, Defaults.AmmoGive2));
@@ -94,11 +105,15 @@ class WeaponPiece : Inventory
 
 	override bool TryPickup (in out Actor toucher)
 	{
+		class<Weapon> type = GetWeaponClass();
+		if (!type)
+			return false;
+
 		Inventory item;
 		WeaponHolder hold = NULL;
 		bool shouldStay = ShouldStay ();
 		int gaveAmmo;
-		let Defaults = GetDefaultByType(WeaponClass);
+		let Defaults = GetDefaultByType(type);
 
 		FullWeapon = NULL;
 		for(item=toucher.Inv; item; item=item.Inv)
@@ -106,6 +121,7 @@ class WeaponPiece : Inventory
 			hold = WeaponHolder(item);
 			if (hold != null)
 			{
+				// Intentionally check against the unreplaced class
 				if (hold.PieceWeapon == WeaponClass) 
 				{
 					break;
@@ -153,9 +169,9 @@ class WeaponPiece : Inventory
 		// Check if  weapon assembled
 		if (hold.PieceMask == (1 << Defaults.health) - 1)
 		{
-			if (!toucher.FindInventory (WeaponClass))
+			if (!toucher.FindInventory (type))
 			{
-				FullWeapon= Weapon(Spawn(WeaponClass));
+				FullWeapon= Weapon(Spawn(type));
 				
 				// The weapon itself should not give more ammo to the player.
 				FullWeapon.AmmoGive1 = 0;