From 7d41f78e788cb1d7bb6056656f5baede0cb6da30 Mon Sep 17 00:00:00 2001
From: Sergio Pedri <sergio0694@live.com>
Date: Sun, 2 Feb 2020 13:13:52 +0100
Subject: [PATCH] Added UI thread dispatch skip to Task/Task<T> overloads

---
 .../Helpers/DispatcherHelper.cs               | 75 +++++++++++++------
 1 file changed, 54 insertions(+), 21 deletions(-)

diff --git a/Microsoft.Toolkit.Uwp/Helpers/DispatcherHelper.cs b/Microsoft.Toolkit.Uwp/Helpers/DispatcherHelper.cs
index ca7d6dc365c..5500fcbbe1e 100644
--- a/Microsoft.Toolkit.Uwp/Helpers/DispatcherHelper.cs
+++ b/Microsoft.Toolkit.Uwp/Helpers/DispatcherHelper.cs
@@ -69,7 +69,7 @@ public static Task<T> ExecuteOnUIThreadAsync<T>(Func<Task<T>> function, CoreDisp
         /// <returns>Awaitable Task/></returns>
         public static Task ExecuteOnUIThreadAsync(this CoreApplicationView viewToExecuteOn, Action function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (viewToExecuteOn == null)
+            if (viewToExecuteOn is null)
             {
                 throw new ArgumentNullException(nameof(viewToExecuteOn));
             }
@@ -87,7 +87,7 @@ public static Task ExecuteOnUIThreadAsync(this CoreApplicationView viewToExecute
         /// <returns>Awaitable Task with type <typeparamref name="T"/></returns>
         public static Task<T> ExecuteOnUIThreadAsync<T>(this CoreApplicationView viewToExecuteOn, Func<T> function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (viewToExecuteOn == null)
+            if (viewToExecuteOn is null)
             {
                 throw new ArgumentNullException(nameof(viewToExecuteOn));
             }
@@ -104,7 +104,7 @@ public static Task<T> ExecuteOnUIThreadAsync<T>(this CoreApplicationView viewToE
         /// <returns>Awaitable Task</returns>
         public static Task ExecuteOnUIThreadAsync(this CoreApplicationView viewToExecuteOn, Func<Task> function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (viewToExecuteOn == null)
+            if (viewToExecuteOn is null)
             {
                 throw new ArgumentNullException(nameof(viewToExecuteOn));
             }
@@ -122,7 +122,7 @@ public static Task ExecuteOnUIThreadAsync(this CoreApplicationView viewToExecute
         /// <returns>Awaitable Task with type <typeparamref name="T"/></returns>
         public static Task<T> ExecuteOnUIThreadAsync<T>(this CoreApplicationView viewToExecuteOn, Func<Task<T>> function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (viewToExecuteOn == null)
+            if (viewToExecuteOn is null)
             {
                 throw new ArgumentNullException(nameof(viewToExecuteOn));
             }
@@ -139,16 +139,16 @@ public static Task<T> ExecuteOnUIThreadAsync<T>(this CoreApplicationView viewToE
         /// <returns>Awaitable Task</returns>
         public static Task AwaitableRunAsync(this CoreDispatcher dispatcher, Action function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (function == null)
+            if (function is null)
             {
                 throw new ArgumentNullException(nameof(function));
             }
 
+            /* Run the function directly when we have thread access.
+             * Also reuse Task.CompletedTask in case of success,
+             * to skip an unnecessary heap allocation for every invocation. */
             if (dispatcher.HasThreadAccess)
             {
-                /* Run the function directly when we have thread access.
-                 * Also reuse Task.CompletedTask in case of success,
-                 * to skip an unnecessary heap allocation for every invocation. */
                 try
                 {
                     function();
@@ -190,19 +190,17 @@ public static Task AwaitableRunAsync(this CoreDispatcher dispatcher, Action func
         /// <returns>Awaitable Task</returns>
         public static Task<T> AwaitableRunAsync<T>(this CoreDispatcher dispatcher, Func<T> function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (function == null)
+            if (function is null)
             {
                 throw new ArgumentNullException(nameof(function));
             }
 
+            // Skip the dispatch, if posssible
             if (dispatcher.HasThreadAccess)
             {
-                // Skip the dispatch, if posssible
                 try
                 {
-                    var result = function();
-
-                    return Task.FromResult(result);
+                    return Task.FromResult(function());
                 }
                 catch (Exception e)
                 {
@@ -236,20 +234,39 @@ public static Task<T> AwaitableRunAsync<T>(this CoreDispatcher dispatcher, Func<
         /// <returns>Awaitable Task</returns>
         public static Task AwaitableRunAsync(this CoreDispatcher dispatcher, Func<Task> function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (function == null)
+            if (function is null)
             {
                 throw new ArgumentNullException(nameof(function));
             }
 
+            /* If we have thread access, we can retrieve the task directly.
+             * We don't use ConfigureAwait(false) in this case, in order
+             * to let the caller continue its execution on the same thread
+             * after awaiting the task returned by this function. */
+            if (dispatcher.HasThreadAccess)
+            {
+                try
+                {
+                    if (function() is Task awaitableResult)
+                    {
+                        return awaitableResult;
+                    }
+
+                    return Task.FromException(new InvalidOperationException("The Task returned by function cannot be null."));
+                }
+                catch (Exception e)
+                {
+                    return Task.FromException(e);
+                }
+            }
+
             var taskCompletionSource = new TaskCompletionSource<object>();
 
             _ = dispatcher.RunAsync(priority, async () =>
             {
                 try
                 {
-                    var awaitableResult = function();
-
-                    if (awaitableResult != null)
+                    if (function() is Task awaitableResult)
                     {
                         await awaitableResult.ConfigureAwait(false);
 
@@ -279,20 +296,36 @@ public static Task AwaitableRunAsync(this CoreDispatcher dispatcher, Func<Task>
         /// <returns>Awaitable Task with type <typeparamref name="T"/></returns>
         public static Task<T> AwaitableRunAsync<T>(this CoreDispatcher dispatcher, Func<Task<T>> function, CoreDispatcherPriority priority = CoreDispatcherPriority.Normal)
         {
-            if (function == null)
+            if (function is null)
             {
                 throw new ArgumentNullException(nameof(function));
             }
 
+            // Skip the dispatch, if posssible
+            if (dispatcher.HasThreadAccess)
+            {
+                try
+                {
+                    if (function() is Task<T> awaitableResult)
+                    {
+                        return awaitableResult;
+                    }
+
+                    return Task.FromException<T>(new InvalidOperationException("The Task returned by function cannot be null."));
+                }
+                catch (Exception e)
+                {
+                    return Task.FromException<T>(e);
+                }
+            }
+
             var taskCompletionSource = new TaskCompletionSource<T>();
 
             _ = dispatcher.RunAsync(priority, async () =>
             {
                 try
                 {
-                    var awaitableResult = function();
-
-                    if (awaitableResult != null)
+                    if (function() is Task<T> awaitableResult)
                     {
                         var result = await awaitableResult.ConfigureAwait(false);