From 24c7f0ffecd28bcac1c0b0d6820ef40fddc7ecab Mon Sep 17 00:00:00 2001
From: Mira <42539014+MiraHell@users.noreply.github.com>
Date: Fri, 31 Jan 2025 18:54:01 +0300
Subject: [PATCH] Fix mitocholide runtime when applied to dismembered limb
 (#28010)

* fix runtime

* handle all cases

* move trait
---
 code/__HELPERS/trait_helpers.dm               |  2 ++
 code/_globalvars/traits.dm                    |  3 ++-
 code/datums/diseases/zombie_virus.dm          | 10 ++++++++++
 code/modules/surgery/organs/organ_external.dm |  4 ++--
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/code/__HELPERS/trait_helpers.dm b/code/__HELPERS/trait_helpers.dm
index 824927828226..e41cb2713498 100644
--- a/code/__HELPERS/trait_helpers.dm
+++ b/code/__HELPERS/trait_helpers.dm
@@ -314,6 +314,8 @@ Remember to update _globalvars/traits.dm if you're adding/removing/renaming trai
 /// Prevents seeing this item on examine when on a mob, or seeing it in the strip menu. It's like ABSTRACT, without making the item fail to interact in several ways. The item can still be stripped however, combine with no_strip unless you have a reason not to.
 #define TRAIT_SKIP_EXAMINE "skip_examine"
 
+/// A general trait for tracking whether a zombie owned the organ or limb
+#define TRAIT_I_WANT_BRAINS_ORGAN "zombie_organ"
 //****** OBJ TRAITS *****//
 
 ///An /obj that should not increase the "depth" of the search for adjacency,
diff --git a/code/_globalvars/traits.dm b/code/_globalvars/traits.dm
index c9a69c212ba9..1211797e7674 100644
--- a/code/_globalvars/traits.dm
+++ b/code/_globalvars/traits.dm
@@ -140,7 +140,8 @@ GLOBAL_LIST_INIT(traits_by_type, list(
 		"TRAIT_HYPOSPRAY_IMMUNE" = TRAIT_HYPOSPRAY_IMMUNE,
 		"TRAIT_ITEM_ACTIVE"	= TRAIT_ITEM_ACTIVE,
 		"TRAIT_NO_STRIP" = TRAIT_NO_STRIP,
-		"TRAIT_SKIP_EXAMINE" = TRAIT_SKIP_EXAMINE
+		"TRAIT_SKIP_EXAMINE" = TRAIT_SKIP_EXAMINE,
+		"TRAIT_I_WANT_BRAINS_ORGAN" = TRAIT_I_WANT_BRAINS_ORGAN
 	),
 
 	/turf = list(
diff --git a/code/datums/diseases/zombie_virus.dm b/code/datums/diseases/zombie_virus.dm
index 37b96d68420f..4f41eb8483d6 100644
--- a/code/datums/diseases/zombie_virus.dm
+++ b/code/datums/diseases/zombie_virus.dm
@@ -29,6 +29,8 @@
 		for(var/obj/item/organ/limb as anything in H.bodyparts)
 			if(!(limb.status & ORGAN_DEAD) && !limb.is_robotic())
 				limb.necrotize(TRUE, TRUE)
+			if(!HAS_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN))
+				ADD_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN, ZOMBIE_TRAIT)
 
 	if(!..())
 		return FALSE
@@ -105,11 +107,15 @@
 	for(var/obj/item/organ/limb as anything in H.bodyparts)
 		if(!(limb.status & ORGAN_DEAD) && !limb.vital && !limb.is_robotic())
 			limb.necrotize()
+		if(!HAS_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN))
+			ADD_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN, ZOMBIE_TRAIT)
 			return FALSE
 
 	for(var/obj/item/organ/limb as anything in H.bodyparts)
 		if(!(limb.status & ORGAN_DEAD) && !limb.is_robotic())
 			limb.necrotize(FALSE, TRUE)
+		if(!HAS_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN))
+			ADD_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN, ZOMBIE_TRAIT)
 			return FALSE
 
 	if(!HAS_TRAIT(affected_mob, TRAIT_I_WANT_BRAINS))
@@ -148,6 +154,10 @@
 /datum/disease/zombie/cure()
 	affected_mob.mind?.remove_antag_datum(/datum/antagonist/zombie)
 	REMOVE_TRAIT(affected_mob, TRAIT_I_WANT_BRAINS, ZOMBIE_TRAIT)
+	var/mob/living/carbon/human/H = affected_mob
+	for(var/obj/item/organ/limb as anything in H.bodyparts)
+		if(HAS_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN))
+			REMOVE_TRAIT(limb, TRAIT_I_WANT_BRAINS_ORGAN, ZOMBIE_TRAIT)
 	affected_mob.DeleteComponent(/datum/component/zombie_regen)
 	affected_mob.med_hud_set_health()
 	affected_mob.med_hud_set_status()
diff --git a/code/modules/surgery/organs/organ_external.dm b/code/modules/surgery/organs/organ_external.dm
index 36c4b86aaf93..3e17ce220f3c 100644
--- a/code/modules/surgery/organs/organ_external.dm
+++ b/code/modules/surgery/organs/organ_external.dm
@@ -394,7 +394,7 @@ This function completely restores a damaged organ to perfect condition.
 	surgeryize()
 	if(is_robotic())	//Robotic organs stay robotic.
 		status = ORGAN_ROBOT
-	else if(HAS_TRAIT(owner, TRAIT_I_WANT_BRAINS))
+	else if(HAS_TRAIT(src, TRAIT_I_WANT_BRAINS_ORGAN))
 		status = ORGAN_DEAD
 	else
 		status = 0
@@ -402,7 +402,7 @@ This function completely restores a damaged organ to perfect condition.
 	perma_injury = 0
 	brute_dam = 0
 	burn_dam = 0
-	if(!HAS_TRAIT(owner, TRAIT_I_WANT_BRAINS)) // zombies's wounds don't close. Because thats cool.
+	if(!HAS_TRAIT(src, TRAIT_I_WANT_BRAINS_ORGAN)) // zombies's wounds don't close. Because thats cool.
 		open = ORGAN_CLOSED //Closing all wounds.
 
 	// handle internal organs