Merge pull request #514 from Banane9/master

Adds support for targeting an async method's internal MoveNext method
pardeike · Mar 22, 2023 · 4e50067 · 4e50067
2 parents 2eef3d0 + b329cbc
commit 4e50067
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 6 deletions.
diff --git a/Harmony/Internal/PatchTools.cs b/Harmony/Internal/PatchTools.cs
@@ -84,8 +84,16 @@ internal static MethodBase GetOriginalMethod(this HarmonyMethod attr)
 					case MethodType.Enumerator:
 						if (attr.methodName is null)
 							return null;
-						var method = AccessTools.DeclaredMethod(attr.declaringType, attr.methodName, attr.argumentTypes);
-						return AccessTools.EnumeratorMoveNext(method);
+						var enumMethod = AccessTools.DeclaredMethod(attr.declaringType, attr.methodName, attr.argumentTypes);
+						return AccessTools.EnumeratorMoveNext(enumMethod);
+
+#if NET45_OR_GREATER
+					case MethodType.Async:
+						if (attr.methodName is null)
+							return null;
+						var asyncMethod = AccessTools.DeclaredMethod(attr.declaringType, attr.methodName, attr.argumentTypes);
+						return AccessTools.AsyncMoveNext(asyncMethod);
+#endif
 				}
 			}
 			catch (AmbiguousMatchException ex)

diff --git a/Harmony/Public/Attributes.cs b/Harmony/Public/Attributes.cs
@@ -17,8 +17,12 @@ public enum MethodType
 		Constructor,
 		/// <summary>This is a static constructor</summary>
 		StaticConstructor,
-		/// <summary>This targets the MoveNext method of the enumerator result</summary>
-		Enumerator
+		/// <summary>This targets the MoveNext method of the enumerator result, that actually contains the method's implementation</summary>
+		Enumerator,
+#if NET45_OR_GREATER
+		/// <summary>This targets the MoveNext method of the async state machine, that actually contains the method's implementation</summary>
+		Async
+#endif
 	}
 
 	/// <summary>Specifies the type of argument</summary>

diff --git a/Harmony/Tools/AccessTools.cs b/Harmony/Tools/AccessTools.cs
@@ -8,6 +8,7 @@
 using System.Reflection;
 using System.Reflection.Emit;
 using System.Runtime.Serialization;
+using System.Runtime.CompilerServices;
 using System.Threading;
 
 namespace HarmonyLib
@@ -466,9 +467,9 @@ public static MethodInfo Method(string typeColonName, Type[] parameters = null,
 			return Method(info.type, info.name, parameters, generics);
 		}
 
-		/// <summary>Gets the <see cref="IEnumerator.MoveNext" /> method of an enumerator method</summary>
+		/// <summary>Gets the <see cref="IEnumerator.MoveNext"/> method of an enumerator method</summary>
 		/// <param name="method">Enumerator method that creates the enumerator <see cref="IEnumerator" /></param>
-		/// <returns>The internal <see cref="IEnumerator.MoveNext" /> method of the enumerator or <b>null</b> if no valid enumerator is detected</returns>
+		/// <returns>The internal <see cref="IEnumerator.MoveNext"/> method of the enumerator or <b>null</b> if no valid enumerator is detected</returns>
 		public static MethodInfo EnumeratorMoveNext(MethodBase method)
 		{
 			if (method is null)
@@ -498,6 +499,37 @@ public static MethodInfo EnumeratorMoveNext(MethodBase method)
 			return Method(type, nameof(IEnumerator.MoveNext));
 		}
 
+#if NET45_OR_GREATER
+		/// <summary>Gets the <see cref="IAsyncStateMachine.MoveNext"/> method of an async method's state machine</summary>
+		/// <param name="method">Async method that creates the state machine internally</param>
+		/// <returns>The internal <see cref="IAsyncStateMachine.MoveNext"/> method of the async state machine or <b>null</b> if no valid async method is detected</returns>
+		public static MethodInfo AsyncMoveNext(MethodBase method)
+		{
+			if (method is null)
+			{
+				FileLog.Debug("AccessTools.AsyncMoveNext: method is null");
+				return null;
+			}
+
+			var asyncAttribute = method.GetCustomAttribute<AsyncStateMachineAttribute>();
+			if (asyncAttribute is null)
+			{
+				FileLog.Debug($"AccessTools.AsyncMoveNext: Could not find AsyncStateMachine for {method.FullDescription()}");
+				return null;
+			}
+
+			var asyncStateMachineType = asyncAttribute.StateMachineType;
+			var asyncMethodBody = DeclaredMethod(asyncStateMachineType, nameof(IAsyncStateMachine.MoveNext));
+			if (asyncMethodBody is null)
+			{
+				FileLog.Debug($"AccessTools.AsyncMoveNext: Could not find async method body for {method.FullDescription()}");
+				return null;
+			}
+
+			return asyncMethodBody;
+		}
+#endif
+
 		/// <summary>Gets the names of all method that are declared in a type</summary>
 		/// <param name="type">The declaring class/type</param>
 		/// <returns>A list of method names</returns>