Fix prefix skipping for dead-end code (#584)

pardeike · Mar 2, 2024 · 1dd31a8 · 1dd31a8
1 parent 3a3a91d
commit 1dd31a8
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/Harmony/Internal/MethodPatcher.cs b/Harmony/Internal/MethodPatcher.cs
@@ -164,7 +164,7 @@ internal MethodInfo CreateReplacement(out Dictionary<int, CodeInstruction> final
 
 			_ = AddPostfixes(privateVars, runOriginalVariable, false);
 
-			if (resultVariable is not null && hasReturnCode)
+			if (resultVariable is not null && (hasReturnCode || (methodEndsInDeadCode && skipOriginalLabel is not null)))
 				emitter.Emit(OpCodes.Ldloc, resultVariable);
 
 			var needsToStorePassthroughResult = AddPostfixes(privateVars, runOriginalVariable, true);
@@ -219,7 +219,7 @@ internal MethodInfo CreateReplacement(out Dictionary<int, CodeInstruction> final
 					emitter.Emit(OpCodes.Ldloc, resultVariable);
 			}
 
-			if (methodEndsInDeadCode == false || hasFinalizers || postfixes.Count > 0)
+			if (methodEndsInDeadCode == false || skipOriginalLabel is not null || hasFinalizers || postfixes.Count > 0)
 				emitter.Emit(OpCodes.Ret);
 
 			finalInstructions = emitter.GetInstructions();

diff --git a/HarmonyTests/Patching/Assets/Specials.cs b/HarmonyTests/Patching/Assets/Specials.cs
@@ -104,6 +104,8 @@ public class DeadEndCode_Patch1
 		public static void Prefix() => prefixCalled = true;
 
 		public static void Postfix() => postfixCalled = true;
+
+		public static bool PrefixWithControl() => false;
 	}
 
 	[HarmonyPatch(typeof(DeadEndCode), nameof(DeadEndCode.Method))]

diff --git a/HarmonyTests/Patching/Specials.cs b/HarmonyTests/Patching/Specials.cs
@@ -179,6 +179,9 @@ public void Test_PatchException()
 			Assert.NotNull(prefix);
 			var postfix = AccessTools.Method(typeof(DeadEndCode_Patch1), nameof(DeadEndCode_Patch1.Postfix));
 			Assert.NotNull(postfix);
+			var prefixWithControl =
+				AccessTools.Method(typeof(DeadEndCode_Patch1), nameof(DeadEndCode_Patch1.PrefixWithControl));
+			Assert.NotNull(postfix);
 
 			// run original
 			try
@@ -208,13 +211,16 @@ public void Test_PatchException()
 			}
 			Assert.True(DeadEndCode_Patch1.prefixCalled);
 
+
+			var isMonoRuntime = AccessTools.IsMonoRuntime;
+			var harmonyMethodPostfix = new HarmonyMethod(postfix);
 			// patch: +postfix
-			if (AccessTools.IsMonoRuntime)
+			if (isMonoRuntime)
 			{
 				// mono should fail
 				try
 				{
-					_ = instance.Patch(original, postfix: new HarmonyMethod(postfix));
+					_ = instance.Patch(original, postfix: harmonyMethodPostfix);
 					Assert.Fail("expecting patch exception");
 				}
 				catch (Exception ex)
@@ -226,7 +232,7 @@ public void Test_PatchException()
 			{
 				// non-mono can add postfix to methods ending in dead code
 
-				_ = instance.Patch(original, postfix: new HarmonyMethod(prefix));
+				_ = instance.Patch(original, postfix: harmonyMethodPostfix);
 				DeadEndCode_Patch1.prefixCalled = false;
 				DeadEndCode_Patch1.postfixCalled = false;
 				// run original
@@ -242,6 +248,16 @@ public void Test_PatchException()
 					Assert.False(DeadEndCode_Patch1.postfixCalled);
 				}
 			}
+
+			_ = instance.Patch(original,
+				prefix: new HarmonyMethod(prefixWithControl),
+				postfix: isMonoRuntime ? harmonyMethodPostfix : null
+			);
+			DeadEndCode_Patch1.prefixCalled = false;
+			DeadEndCode_Patch1.postfixCalled = false;
+			test.Method();
+			Assert.True(DeadEndCode_Patch1.prefixCalled);
+			Assert.True(DeadEndCode_Patch1.postfixCalled);
 		}
 
 		[Test]