diff options
Diffstat (limited to 'src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs')
-rw-r--r-- | src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs | 86 |
1 files changed, 54 insertions, 32 deletions
diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs index 221a1a437..02b5910b7 100644 --- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs +++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs @@ -6,12 +6,16 @@ using System.Diagnostics; using System.Threading.Tasks; using System.Threading.Tasks.Sources; +#if !netstandard +using Internal.Runtime.CompilerServices; +#endif + namespace System.Runtime.CompilerServices { /// <summary>Provides an awaiter for a <see cref="ValueTask"/>.</summary> public readonly struct ValueTaskAwaiter : ICriticalNotifyCompletion #if CORECLR - , IValueTaskAwaiter + , IStateMachineBoxAwareAwaiter #endif { /// <summary>Shim used to invoke an <see cref="Action"/> passed as the state argument to a <see cref="Action{Object}"/>.</summary> @@ -48,13 +52,16 @@ namespace System.Runtime.CompilerServices /// <summary>Schedules the continuation action for this ValueTask.</summary> public void OnCompleted(Action continuation) { - if (_value.ObjectIsTask) + object obj = _value._obj; + Debug.Assert(obj == null || obj is Task || obj is IValueTaskSource); + + if (obj is Task t) { - _value.UnsafeGetTask().GetAwaiter().OnCompleted(continuation); + t.GetAwaiter().OnCompleted(continuation); } - else if (_value._obj != null) + else if (obj != null) { - _value.UnsafeGetValueTaskSource().OnCompleted(s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext | ValueTaskSourceOnCompletedFlags.FlowExecutionContext); + Unsafe.As<IValueTaskSource>(obj).OnCompleted(s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext | ValueTaskSourceOnCompletedFlags.FlowExecutionContext); } else { @@ -65,13 +72,16 @@ namespace System.Runtime.CompilerServices /// <summary>Schedules the continuation action for this ValueTask.</summary> public void UnsafeOnCompleted(Action continuation) { - if (_value.ObjectIsTask) + object obj = _value._obj; + Debug.Assert(obj == null || obj is Task || obj is IValueTaskSource); + + if (obj is Task t) { - _value.UnsafeGetTask().GetAwaiter().UnsafeOnCompleted(continuation); + t.GetAwaiter().UnsafeOnCompleted(continuation); } - else if (_value._obj != null) + else if (obj != null) { - _value.UnsafeGetValueTaskSource().OnCompleted(s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); + Unsafe.As<IValueTaskSource>(obj).OnCompleted(s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); } else { @@ -80,15 +90,18 @@ namespace System.Runtime.CompilerServices } #if CORECLR - void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box) + void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box) { - if (_value.ObjectIsTask) + object obj = _value._obj; + Debug.Assert(obj == null || obj is Task || obj is IValueTaskSource); + + if (obj is Task t) { - TaskAwaiter.UnsafeOnCompletedInternal(_value.UnsafeGetTask(), box, continueOnCapturedContext: true); + TaskAwaiter.UnsafeOnCompletedInternal(t, box, continueOnCapturedContext: true); } - else if (_value._obj != null) + else if (obj != null) { - _value.UnsafeGetValueTaskSource().OnCompleted(s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); + Unsafe.As<IValueTaskSource>(obj).OnCompleted(s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); } else { @@ -105,7 +118,7 @@ namespace System.Runtime.CompilerServices return; } - box.Invoke(null); + box.MoveNext(); }; #endif } @@ -113,7 +126,7 @@ namespace System.Runtime.CompilerServices /// <summary>Provides an awaiter for a <see cref="ValueTask{TResult}"/>.</summary> public readonly struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion #if CORECLR - , IValueTaskAwaiter + , IStateMachineBoxAwareAwaiter #endif { /// <summary>The value being awaited.</summary> @@ -139,13 +152,16 @@ namespace System.Runtime.CompilerServices /// <summary>Schedules the continuation action for this ValueTask.</summary> public void OnCompleted(Action continuation) { - if (_value.ObjectIsTask) + object obj = _value._obj; + Debug.Assert(obj == null || obj is Task<TResult> || obj is IValueTaskSource<TResult>); + + if (obj is Task<TResult> t) { - _value.UnsafeGetTask().GetAwaiter().OnCompleted(continuation); + t.GetAwaiter().OnCompleted(continuation); } - else if (_value._obj != null) + else if (obj != null) { - _value.UnsafeGetValueTaskSource().OnCompleted(ValueTaskAwaiter.s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext | ValueTaskSourceOnCompletedFlags.FlowExecutionContext); + Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ValueTaskAwaiter.s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext | ValueTaskSourceOnCompletedFlags.FlowExecutionContext); } else { @@ -156,13 +172,16 @@ namespace System.Runtime.CompilerServices /// <summary>Schedules the continuation action for this ValueTask.</summary> public void UnsafeOnCompleted(Action continuation) { - if (_value.ObjectIsTask) + object obj = _value._obj; + Debug.Assert(obj == null || obj is Task<TResult> || obj is IValueTaskSource<TResult>); + + if (obj is Task<TResult> t) { - _value.UnsafeGetTask().GetAwaiter().UnsafeOnCompleted(continuation); + t.GetAwaiter().UnsafeOnCompleted(continuation); } - else if (_value._obj != null) + else if (obj != null) { - _value.UnsafeGetValueTaskSource().OnCompleted(ValueTaskAwaiter.s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); + Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ValueTaskAwaiter.s_invokeActionDelegate, continuation, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); } else { @@ -171,15 +190,18 @@ namespace System.Runtime.CompilerServices } #if CORECLR - void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box) + void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box) { - if (_value.ObjectIsTask) + object obj = _value._obj; + Debug.Assert(obj == null || obj is Task<TResult> || obj is IValueTaskSource<TResult>); + + if (obj is Task<TResult> t) { - TaskAwaiter.UnsafeOnCompletedInternal(_value.UnsafeGetTask(), box, continueOnCapturedContext: true); + TaskAwaiter.UnsafeOnCompletedInternal(t, box, continueOnCapturedContext: true); } - else if (_value._obj != null) + else if (obj != null) { - _value.UnsafeGetValueTaskSource().OnCompleted(ValueTaskAwaiter.s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); + Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ValueTaskAwaiter.s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext); } else { @@ -190,12 +212,12 @@ namespace System.Runtime.CompilerServices } #if CORECLR - /// <summary>Internal interface used to enable optimizations from <see cref="AsyncTaskMethodBuilder"/> on <see cref="ValueTask"/>.</summary>> - internal interface IValueTaskAwaiter + /// <summary>Internal interface used to enable optimizations from <see cref="AsyncTaskMethodBuilder"/>.</summary>> + internal interface IStateMachineBoxAwareAwaiter { /// <summary>Invoked to set <see cref="ITaskCompletionAction.Invoke"/> of the <paramref name="box"/> as the awaiter's continuation.</summary> /// <param name="box">The box object.</param> void AwaitUnsafeOnCompleted(IAsyncStateMachineBox box); } #endif -}
\ No newline at end of file +} |