From b0b7ffba5086da0d8681bffb3783ecf4fdcd4ecb Mon Sep 17 00:00:00 2001
From: lutepickle <28810-lutepickle@users.noreply.gitgud.io>
Date: Sat, 16 Jul 2022 14:50:01 -0700
Subject: [PATCH] Associate the post-birth vagina morph with the one that was
 pregnant

---
 .../RJW_Menstruation/RJW_Menstruation/Patch/RJW_Patch.cs   | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/1.3/source/RJW_Menstruation/RJW_Menstruation/Patch/RJW_Patch.cs b/1.3/source/RJW_Menstruation/RJW_Menstruation/Patch/RJW_Patch.cs
index a985598..6052d15 100644
--- a/1.3/source/RJW_Menstruation/RJW_Menstruation/Patch/RJW_Patch.cs
+++ b/1.3/source/RJW_Menstruation/RJW_Menstruation/Patch/RJW_Patch.cs
@@ -145,11 +145,14 @@ namespace RJW_Menstruation
     [HarmonyPatch(typeof(Hediff_BasePregnancy), nameof(Hediff_BasePregnancy.PostBirth))]
     public static class RJW_Patch_PostBirth
     {
-        public static void Postfix(Pawn mother, Pawn baby)
+        public static void Postfix(Hediff_BasePregnancy __instance, Pawn mother, Pawn baby)
         {
             if (Configurations.EnableBirthVaginaMorph)
             {
-                Hediff vagina = mother.health.hediffSet.hediffs.FirstOrFallback(x => x.def.defName.ToLower().Contains("vagina"));
+                // The comp still has the pregnancy attached at this point in the process
+                Hediff vagina = __instance.GetMenstruationComp()?.parent;
+                if (vagina == null) vagina = mother.health.hediffSet.hediffs.FirstOrFallback(x => x.def.defName.ToLower().Contains("vagina"));
+                if (vagina == null) return;
                 float morph = Mathf.Max(baby.BodySize - Mathf.Pow(vagina.Severity * mother.BodySize, 2), 0f);
                 vagina.Severity += morph * Configurations.VaginaMorphPower;
             }