#12 Successfully add support for async methods

Miista · Jan 27, 2024 · 31e537d · 31e537d
1 parent 63cfc59
commit 31e537d
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 6 deletions.
diff --git a/src/Pose/PoseContext.cs b/src/Pose/PoseContext.cs
@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Reflection;
 using System.Reflection.Emit;
+using System.Threading.Tasks;
 using Pose.IL;
 
 namespace Pose
@@ -30,5 +31,27 @@ public static void Isolate(Action entryPoint, params Shim[] shims)
             Console.WriteLine("----------------------------- Invoking ----------------------------- ");
             methodInfo.CreateDelegate(delegateType).DynamicInvoke(entryPoint.Target);
         }
+
+        public static async Task Isolate(Func<Task> entryPoint, params Shim[] shims)
+        {
+            if (shims == null || shims.Length == 0)
+            {
+                await entryPoint.Invoke();
+                return;
+            }
+
+            Shims = shims;
+            StubCache = new Dictionary<MethodBase, DynamicMethod>();
+
+            var delegateType = typeof(Func<Task>);
+            var rewriter = MethodRewriter.CreateRewriter(entryPoint.Method, false);
+            Console.WriteLine("----------------------------- Rewriting ----------------------------- ");
+            var methodInfo = (MethodInfo)(rewriter.Rewrite());
+
+            Console.WriteLine("----------------------------- Invoking ----------------------------- ");
+
+            // ReSharper disable once PossibleNullReferenceException
+            await (methodInfo.CreateDelegate(delegateType).DynamicInvoke() as Task);
+        }
     }
 }
diff --git a/src/Pose/Shim.cs b/src/Pose/Shim.cs
@@ -63,7 +63,8 @@ public static Shim Replace<T>(Expression<Func<T>> expression, bool setter = fals
 
         private static Shim ReplaceImpl<T>(Expression<T> expression, bool setter)
         {
-            // TODO: Figure out if method is an async method. Do that by finding the attribute on the method which designates the state machine
+            // We could find out whether the method is an async method by checking whether it has the AsyncStateMachineAttribute.
+            // However, it seems that this is not necessary.
             var methodBase = ShimHelper.GetMethodFromExpression(expression.Body, setter, out var instance);
             return new Shim(methodBase, instance) { _setter = setter };
         }

diff --git a/src/Sandbox/Program.cs b/src/Sandbox/Program.cs
@@ -25,15 +25,19 @@ public static void Main(string[] args)
                 }, dateTimeShim);
 #elif NETCOREAPP2_0
             Console.WriteLine("2.0");
-            var dateTimeShim = Shim.Replace(() => DateTime.Now).With(() => new DateTime(2004, 1, 1));
-            var asyncShim = Shim.Replace(() => GetIntAsync()).With(() => Task.FromResult(10));
+            //var dateTimeShim = Shim.Replace(() => DateTime.Now).With(() => new DateTime(2004, 1, 1));
+            var asyncShim = Shim.Replace(() => GetIntAsync()).With(() =>
+            {
+                Console.WriteLine("This actually works!!!");
+                return Task.FromResult(15);
+            });
             PoseContext.Isolate(
-                () =>
+                async () =>
                 {
-                    var result = GetIntAsync().GetAwaiter().GetResult();
+                    var result = await GetIntAsync();
                     Console.WriteLine($"Result: {result}");
                     //Console.WriteLine(DateTime.Now);
-                }, dateTimeShim, asyncShim);
+                }, asyncShim);
 #elif NET6_0
             Console.WriteLine("6.0");
             var dateTimeShim = Shim.Replace(() => DateTime.Now).With(() => new DateTime(2004, 1, 1));