Update code in AwaitUnsafeOnCompleted to use interface matching

Given that the jit can now avoid boxing on some interface calls to value types,
generalize the patterns introduced in AwaitUnsafeOnCompleted in #14718 by
using interfaces instead of checking for specific types.

Also move the catch-all processing back in line as the jit can now fold the
interface tests early and not pull in the EH portion of the method unless it is
needed.

src/mscorlib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs

@@ -35,7 +35,7 @@ internal ConfiguredValueTaskAwaitable(ValueTask<TResult> value, bool continueOnC
 
         /// <summary>Provides an awaiter for a <see cref="ConfiguredValueTaskAwaitable{TResult}"/>.</summary>
         [StructLayout(LayoutKind.Auto)]
-        public struct ConfiguredValueTaskAwaiter : ICriticalNotifyCompletion
+        public struct ConfiguredValueTaskAwaiter : ICriticalNotifyCompletion, IConfiguredValueTaskAwaiter
         {
             /// <summary>The value being awaited.</summary>
             private ValueTask<TResult> _value; // Methods are called on this; avoid making it readonly so as to avoid unnecessary copies
@@ -71,6 +71,9 @@ internal ConfiguredValueTaskAwaiter(ValueTask<TResult> value, bool continueOnCap
 
             /// <summary>Gets the task underlying <see cref="_value"/>.</summary>
             internal Task<TResult> AsTask() => _value.AsTask();
+
+            /// <summary>Gets the task underlying <see cref="_value"/>.</summary>
+            (Task, bool) IConfiguredValueTaskAwaiter.GetTask() => (_value.AsTask(), _continueOnCapturedContext);
         }
     }
 }

src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs

@@ -9,7 +9,7 @@
 namespace System.Runtime.CompilerServices
 {
     /// <summary>Provides an awaiter for a <see cref="ValueTask{TResult}"/>.</summary>
-    public struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion
+    public struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion, IValueTaskAwaiter
     {
         /// <summary>The value being awaited.</summary>
         private ValueTask<TResult> _value; // Methods are called on this; avoid making it readonly so as to avoid unnecessary copies
@@ -38,5 +38,8 @@ public struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion
 
         /// <summary>Gets the task underlying <see cref="_value"/>.</summary>
         internal Task<TResult> AsTask() => _value.AsTask();
+
+        /// <summary>Gets the task underlying <see cref="_value"/>.</summary>
+        Task IValueTaskAwaiter.GetTask() => _value.AsTask();
     }
 }

src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs

@@ -393,7 +393,7 @@ public void SetStateMachine(IAsyncStateMachine stateMachine)
         {
             IAsyncStateMachineBox box = GetStateMachineBox(ref stateMachine);
 
-            // TThe null tests here ensure that the jit can optimize away the interface
+            // The null tests here ensure that the jit can optimize away the interface
             // tests when TAwaiter is is a ref type.
             if ((null != (object)default(TAwaiter)) && (awaiter is ITaskAwaiter))
             {
@@ -405,62 +405,27 @@ public void SetStateMachine(IAsyncStateMachine stateMachine)
                 ref ConfiguredTaskAwaitable.ConfiguredTaskAwaiter ta = ref Unsafe.As<TAwaiter, ConfiguredTaskAwaitable.ConfiguredTaskAwaiter>(ref awaiter);
                 TaskAwaiter.UnsafeOnCompletedInternal(ta.m_task, box, ta.m_continueOnCapturedContext);
             }
-
-            // Handle common {Configured}ValueTaskAwaiter<T> types.  Unfortunately these need to be special-cased
-            // individually, as we don't have good way to extract the task from a ValueTaskAwaiter<T> when we don't
-            // know what the T is; we could make ValueTaskAwaiter<T> implement an IValueTaskAwaiter interface, but
-            // calling a GetTask method on that would end up boxing the awaiter.  This hard-coded list here is
-            // somewhat arbitrary and is based on types currently in use with ValueTask<T> in coreclr/corefx.
-            else if (typeof(TAwaiter) == typeof(ValueTaskAwaiter<int>))
-            {
-                var vta = (ValueTaskAwaiter<int>)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, continueOnCapturedContext: true);
-            }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<int>.ConfiguredValueTaskAwaiter))
-            {
-                var vta = (ConfiguredValueTaskAwaitable<int>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
-            }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<System.IO.Stream>.ConfiguredValueTaskAwaiter))
-            {
-                var vta = (ConfiguredValueTaskAwaitable<System.IO.Stream>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
-            }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<ArraySegment<byte>>.ConfiguredValueTaskAwaiter))
+            else if ((null != (object)default(TAwaiter)) && (awaiter is IValueTaskAwaiter))
             {
-                var vta = (ConfiguredValueTaskAwaitable<ArraySegment<byte>>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
+                Task t = ((IValueTaskAwaiter)awaiter).GetTask();
+                TaskAwaiter.UnsafeOnCompletedInternal(t, box, continueOnCapturedContext: true);
             }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<object>.ConfiguredValueTaskAwaiter))
+            else if ((null != (object)default(TAwaiter)) && (awaiter is IConfiguredValueTaskAwaiter))
             {
-                var vta = (ConfiguredValueTaskAwaitable<object>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
+                (Task task, bool continueOnCapturedContext) t = ((IConfiguredValueTaskAwaiter)awaiter).GetTask();
+                TaskAwaiter.UnsafeOnCompletedInternal(t.task, box, t.continueOnCapturedContext);
             }
-
             // The awaiter isn't specially known. Fall back to doing a normal await.
             else
             {
-                // TODO https://github.com/dotnet/coreclr/issues/14177:
-                // Move the code back into this method once the JIT is able to
-                // elide it successfully when one of the previous branches is hit.
-                AwaitArbitraryAwaiterUnsafeOnCompleted(ref awaiter, box);
-            }
-        }
-
-        /// <summary>Schedules the specified state machine to be pushed forward when the specified awaiter completes.</summary>
-        /// <typeparam name="TAwaiter">Specifies the type of the awaiter.</typeparam>
-        /// <param name="awaiter">The awaiter.</param>
-        /// <param name="box">The state machine box.</param>
-        private static void AwaitArbitraryAwaiterUnsafeOnCompleted<TAwaiter>(ref TAwaiter awaiter, IAsyncStateMachineBox box)
-            where TAwaiter : ICriticalNotifyCompletion
-        {
-            try
-            {
-                awaiter.UnsafeOnCompleted(box.MoveNextAction);
-            }
-            catch (Exception e)
-            {
-                AsyncMethodBuilderCore.ThrowAsync(e, targetContext: null);
+                try
+                {
+                    awaiter.UnsafeOnCompleted(box.MoveNextAction);
+                }
+                catch (Exception e)
+                {
+                    AsyncMethodBuilderCore.ThrowAsync(e, targetContext: null);
+                }
             }
         }
 

src/mscorlib/src/System/Runtime/CompilerServices/TaskAwaiter.cs

@@ -391,6 +391,22 @@ internal interface ITaskAwaiter { }
     /// </summary>
     internal interface IConfiguredTaskAwaiter { }
 
+    /// <summary>
+    /// Internal interface used to enable extract the Task from arbitrary ValueTask awaiters.
+    /// </summary>>
+    internal interface IValueTaskAwaiter
+    {
+        Task GetTask();
+    }
+
+    /// <summary>
+    /// Internal interface used to enable extract the Task from arbitrary configured ValueTask awaiters.
+    /// </summary>
+    internal interface IConfiguredValueTaskAwaiter
+    {
+        (Task task, bool continueOnCapturedContext) GetTask();
+    }
+
     /// <summary>Provides an awaitable object that allows for configured awaits on <see cref="System.Threading.Tasks.Task"/>.</summary>
     /// <remarks>This type is intended for compiler use only.</remarks>
     public struct ConfiguredTaskAwaitable